// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2025 Meta Platforms, Inc. and affiliates. */
#include <test_progs.h>
#include <string.h>
#include <stdio.h>
#include "task_work_stress.skel.h"
#include <linux/bpf.h>
#include <linux/perf_event.h>
#include <sys/syscall.h>
#include <time.h>
#include <stdlib.h>
#include <stdatomic.h>

struct test_data {
	int prog_fd;
	atomic_int exit;
};

void *runner(void *test_data)
{
	struct test_data *td = test_data;
	int err = 0;
	LIBBPF_OPTS(bpf_test_run_opts, opts);

	while (!err && !atomic_load(&td->exit))
		err = bpf_prog_test_run_opts(td->prog_fd, &opts);

	return NULL;
}

static int get_env_int(const char *str, int def)
{
	const char *s = getenv(str);
	char *end;
	int retval;

	if (!s || !*s)
		return def;
	errno = 0;
	retval = strtol(s, &end, 10);
	if (errno || *end || retval < 0)
		return def;
	return retval;
}

static void task_work_run(bool enable_delete)
{
	struct task_work_stress *skel;
	struct bpf_program *scheduler, *deleter;
	int nthreads = 16;
	int test_time_s = get_env_int("BPF_TASK_WORK_TEST_TIME", 1);
	pthread_t tid[nthreads], tid_del;
	bool started[nthreads], started_del = false;
	struct test_data td_sched = { .exit = 0 }, td_del = { .exit = 1 };
	int i, err;

	skel = task_work_stress__open();
	if (!ASSERT_OK_PTR(skel, "task_work__open"))
		return;

	scheduler = bpf_object__find_program_by_name(skel->obj, "schedule_task_work");
	bpf_program__set_autoload(scheduler, true);

	deleter = bpf_object__find_program_by_name(skel->obj, "delete_task_work");
	bpf_program__set_autoload(deleter, true);

	err = task_work_stress__load(skel);
	if (!ASSERT_OK(err, "skel_load"))
		goto cleanup;

	for (i = 0; i < nthreads; ++i)
		started[i] = false;

	td_sched.prog_fd = bpf_program__fd(scheduler);
	for (i = 0; i < nthreads; ++i) {
		if (pthread_create(&tid[i], NULL, runner, &td_sched) != 0) {
			fprintf(stderr, "could not start thread");
			goto cancel;
		}
		started[i] = true;
	}

	if (enable_delete)
		atomic_store(&td_del.exit, 0);

	td_del.prog_fd = bpf_program__fd(deleter);
	if (pthread_create(&tid_del, NULL, runner, &td_del) != 0) {
		fprintf(stderr, "could not start thread");
		goto cancel;
	}
	started_del = true;

	/* Run stress test for some time */
	sleep(test_time_s);

cancel:
	atomic_store(&td_sched.exit, 1);
	atomic_store(&td_del.exit, 1);
	for (i = 0; i < nthreads; ++i) {
		if (started[i])
			pthread_join(tid[i], NULL);
	}

	if (started_del)
		pthread_join(tid_del, NULL);

	ASSERT_GT(skel->bss->callback_scheduled, 0, "work scheduled");
	/* Some scheduling attempts should have failed due to contention */
	ASSERT_GT(skel->bss->schedule_error, 0, "schedule error");

	if (enable_delete) {
		/* If delete thread is enabled, it has cancelled some callbacks */
		ASSERT_GT(skel->bss->delete_success, 0, "delete success");
		ASSERT_LT(skel->bss->callback_success, skel->bss->callback_scheduled, "callbacks");
	} else {
		/* Without delete thread number of scheduled callbacks is the same as fired */
		ASSERT_EQ(skel->bss->callback_success, skel->bss->callback_scheduled, "callbacks");
	}

cleanup:
	task_work_stress__destroy(skel);
}

void test_task_work_stress(void)
{
	if (test__start_subtest("no_delete"))
		task_work_run(false);
	if (test__start_subtest("with_delete"))
		task_work_run(true);
}