// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#include <test_progs.h>
#include <unistd.h>

#include <libarena/common.h>
#include <libarena/asan.h>
#include <libarena/buddy.h>
#include <libarena/userspace.h>

#include "libarena/libarena.skel.h"

static void run_libarena_test(struct libarena *skel, struct bpf_program *prog,
		const char *name)
{
	int ret;

	if (strstr(name, "test_buddy")) {
		/* Buddy tests initialize the allocator directly. */
		ret = libarena_run_prog(bpf_program__fd(skel->progs.arena_buddy_destroy));
		if (!ASSERT_OK(ret, "arena_buddy_destroy"))
			return;
	} else {
		ret = libarena_run_prog(bpf_program__fd(skel->progs.arena_buddy_reset));
		if (!ASSERT_OK(ret, "arena_buddy_reset"))
			return;
	}

	ret = libarena_run_prog(bpf_program__fd(prog));

	ASSERT_OK(ret, name);
}

static void *run_libarena_parallel_prog(void *arg)
{
	struct bpf_program *prog = arg;

	return (void *)(long)libarena_run_prog(bpf_program__fd(prog));
}

/* Max suffix is ceil((lg 2^32) / (lg 10)) + sizeof("__") = 10 + 2 = 12. */
#define MAX_PARTEST_SUFFIX (12)
#define MAX_PARTEST_NAME (1024)
#define MAX_PARTEST_PREFIX (MAX_PARTEST_NAME - MAX_PARTEST_SUFFIX)

static int run_libarena_parallel_fini(struct libarena *skel, const char *name,
				      size_t prefixlen)
{
	char tdname[MAX_PARTEST_NAME];
	struct bpf_program *fini_prog;
	int ret;

	ret = snprintf(tdname, sizeof(tdname), "%.*s__fini", (int)prefixlen, name);
	if (!ASSERT_LT(ret, sizeof(tdname), "partest fini name"))
		return -ENAMETOOLONG;

	fini_prog = bpf_object__find_program_by_name(skel->obj, tdname);
	if (!ASSERT_TRUE(fini_prog, "partest fini prog"))
		return -ENOENT;

	ret = libarena_run_prog(bpf_program__fd(fini_prog));
	ASSERT_OK(ret, tdname);

	return ret;
}

static int run_libarena_parallel_test_workers(struct libarena *skel,
		const char *name, size_t prefixlen)
{
	pthread_t *threads = NULL, *tmp_threads;
	char tdname[MAX_PARTEST_NAME];
	struct bpf_program *tdprog;
	uint32_t nthreads;
	void *thread_ret;
	int ret, err = 0;
	int worker_err;
	int i;

	for (nthreads = 0; nthreads < UINT_MAX; nthreads++) {
		ret = snprintf(tdname, sizeof(tdname), "%.*s__%u", (int)prefixlen,
			       name, nthreads);
		if (!ASSERT_LT(ret, sizeof(tdname), "test worker name")) {
			err = -ENAMETOOLONG;
			break;
		}

		/* 
		 * We enumerate the worker threads for a given test with __0, __1,
		 * and so on. The suffixes always start from 0 and are contiguous,
		 * so if we don't find a program with the requested name we have
		 * discovered all available worker programs.
		 */
		tdprog = bpf_object__find_program_by_name(skel->obj, tdname);
		if (!tdprog)
			break;

		/* Bump the alloc array to accommodate the new thread. */
		tmp_threads = realloc(threads, (nthreads + 1) * sizeof(*threads));
		if (!ASSERT_TRUE(tmp_threads, "realloc")) {
			err = -ENOMEM;
			break;
		}
		threads = tmp_threads;

		ret = pthread_create(&threads[nthreads], NULL,
				     run_libarena_parallel_prog,
				     tdprog);
		if (!ASSERT_OK(ret, "pthread_create")) {
			err = ret;
			break;
		}
	}


	for (i = 0; i < nthreads; i++) {
		ret = pthread_join(threads[i], &thread_ret);
		if (!ASSERT_OK(ret, "pthread_join")) {
			err = err ?: ret;
			continue;
		}

		worker_err = (long)thread_ret;

		/*
		 * A worker that bails out because another one already gave up
		 * reports -EINTR. It is collateral damage that carries no
		 * information, so skip it entirely: never let it become the
		 * reported error, and don't log it either.
		 */
		if (!worker_err || worker_err == -EINTR)
			continue;

		if (!err)
			err = worker_err;

		fprintf(stdout, "%.*s__%d returned %d\n", (int)prefixlen, name,
			i, worker_err);
	}

	free(threads);

	return err;
}

static bool libarena_parallel_test_enabled(struct libarena *skel,
					   const char *prefix,
					   size_t prefixlen)
{
	struct bpf_program *prog;
	char progname[MAX_PARTEST_NAME];
	int ret;

	ret = snprintf(progname, sizeof(progname), "%.*s__enabled", (int)prefixlen,
		       prefix);
	if (!ASSERT_LT(ret, sizeof(progname), "partest enabled name"))
		return false;

	prog = bpf_object__find_program_by_name(skel->obj, progname);
	if (!prog)
		return true;

	ret = libarena_run_prog(bpf_program__fd(prog));
	if (ret == -EOPNOTSUPP)
		return false;
	if (!ASSERT_OK(ret, progname))
		return false;
	return true;
}

static void run_libarena_parallel_test(struct libarena *skel, struct bpf_program *prog,
		const char *name)
{
	char testname[MAX_PARTEST_NAME];
	size_t prefixlen;
	const char *pos;
	int ret;

	/*
	 * We annotate the initialization prog with __init. If the current prog does
	 * not match, it is one of the parallel threads instead and is ignored.
	 *
	 * We assume the test writer knows what they are doing and do not add __init
	 * randomly in the middle of a test name.
	 */
	pos = strstr(name, "__init");
	if (!pos)
		return;

	prefixlen = pos - name;
	if (!ASSERT_LT(prefixlen, MAX_PARTEST_PREFIX, "partest prefix too long"))
		return;

	/* The name of the test without the __init suffix. Looks nicer in the test log. */
	ret = snprintf(testname, sizeof(testname), "%.*s", (int)prefixlen, name);
	if (!ASSERT_LT(ret, sizeof(testname), "partest test name"))
		return;

	if (!test__start_subtest(testname))
		return;

	if (!libarena_parallel_test_enabled(skel, testname, prefixlen)) {
		test__skip();
		return;
	}

	ret = libarena_run_prog(bpf_program__fd(skel->progs.arena_buddy_reset));
	if (!ASSERT_OK(ret, "arena_buddy_reset"))
		return;

	ret = libarena_run_prog(bpf_program__fd(prog));
	if (!ASSERT_OK(ret, testname))
		return;

	ret = run_libarena_parallel_test_workers(skel, name, prefixlen);

	ASSERT_OK(ret, testname);

	run_libarena_parallel_fini(skel, name, prefixlen);
}

void serial_test_libarena(void)
{
	struct arena_alloc_reserve_args args;
	struct libarena *skel;
	struct bpf_program *prog;
	int ret;

	skel = libarena__open_and_load();
	if (!ASSERT_OK_PTR(skel, "open_and_load"))
		return;

	ret = libarena__attach(skel);
	if (!ASSERT_OK(ret, "attach"))
		goto out;

	args.nr_pages = ARENA_RESERVE_PAGES_DFL;

	ret = libarena_run_prog_args(bpf_program__fd(skel->progs.arena_alloc_reserve),
			&args, sizeof(args));
	if (!ASSERT_OK(ret, "arena_alloc_reserve"))
		goto out;

	bpf_object__for_each_program(prog, skel->obj) {
		const char *name = bpf_program__name(prog);

		/*
		 * Handle parallel test progs separately. For those
		 * progs it's not a matter of test/skip, because each
		 * parallel test prog includes an initialization prog
		 * and a set of progs to be run in parallel. For the
		 * latter we do not record them as skipped or run,
		 * because we run them all at once when we come across
		 * the initialization prog. For more details on how we
		 * discover the progs see the comment on
		 * run_libarena_parallel_test.
		 */
		if (libarena_is_parallel_test_prog(name)) {
			run_libarena_parallel_test(skel, prog, name);
			continue;
		}

		if (!libarena_is_test_prog(name))
			continue;

		if (!test__start_subtest(name))
			continue;

		run_libarena_test(skel, prog, name);
	}

out:
	libarena__destroy(skel);
}