// SPDX-License-Identifier: GPL-2.0
/*
 * Stress every LRU lock-failure and orphan-recovery.
 * perf_event NMI BPF on every online CPU does
 * update+delete on a small LRU map; userspace threads on every CPU do
 * the same from syscall context.
 */
#define _GNU_SOURCE
#include <pthread.h>
#include <sched.h>
#include <sys/syscall.h>
#include <linux/perf_event.h>
#include <test_progs.h>
#include "testing_helpers.h"
#include "lru_lock_nmi.skel.h"

#define MAP_ENTRIES	64
#define KEY_RANGE	(MAP_ENTRIES * 2)
#define STRESS_NS	(500 * 1000 * 1000ULL)

struct hammer_arg {
	int map_fd;
	int cpu;
	__u64 deadline_ns;
};

struct refill_arg {
	int map_fd;
	int cpu;
	int per_cpu_quota;
	int update_errors;
};

/*
 * Pin the calling thread to @cpu. Uses dynamically-allocated CPU sets so
 * we stay correct on hosts with @cpu >= CPU_SETSIZE (default 1024).
 */
static int pin_to_cpu(int cpu)
{
	cpu_set_t *cs;
	size_t cs_size;
	int err;

	cs = CPU_ALLOC(cpu + 1);
	if (!cs)
		return -ENOMEM;
	cs_size = CPU_ALLOC_SIZE(cpu + 1);

	CPU_ZERO_S(cs_size, cs);
	CPU_SET_S(cpu, cs_size, cs);
	err = pthread_setaffinity_np(pthread_self(), cs_size, cs);
	CPU_FREE(cs);
	return err;
}

static void *hammer_thread(void *p)
{
	struct hammer_arg *a = p;
	int nr_possible_cpus = libbpf_num_possible_cpus();
	__u64 val[nr_possible_cpus];
	unsigned int seed;
	__u32 key;

	memset(val, 0, sizeof(val));
	pin_to_cpu(a->cpu);

	seed = (unsigned int)a->cpu ^ (unsigned int)(uintptr_t)pthread_self();

	while (get_time_ns() < a->deadline_ns) {
		bool do_update = rand_r(&seed) & 1;

		key = rand_r(&seed) % KEY_RANGE;
		if (do_update)
			bpf_map_update_elem(a->map_fd, &key, val, BPF_ANY);
		else
			bpf_map_delete_elem(a->map_fd, &key);
	}
	return NULL;
}

static void *refill_thread(void *p)
{
	struct refill_arg *a = p;
	int nr_possible_cpus = libbpf_num_possible_cpus();
	__u64 val[nr_possible_cpus];
	__u32 start, end, key;

	memset(val, 0, sizeof(val));
	pin_to_cpu(a->cpu);

	start = (__u32)a->cpu * (__u32)a->per_cpu_quota;
	end   = start + (__u32)a->per_cpu_quota;
	for (key = start; key < end; key++)
		if (bpf_map_update_elem(a->map_fd, &key, val, BPF_ANY))
			a->update_errors++;
	return NULL;
}

/*
 * Drain the map, then refill it with each CPU inserting only its own
 * quota of keys.
 * After refill, lookup every key we inserted - a stranded node on any
 * CPU's pool would have forced eviction.
 */
static int drain_then_verify_capacity(int map_fd, int nr_cpus)
{
	int per_cpu_quota = MAP_ENTRIES / nr_cpus;
	int total = per_cpu_quota * nr_cpus;
	int nr_possible_cpus = libbpf_num_possible_cpus();
	pthread_t threads[nr_cpus];
	struct refill_arg args[nr_cpus];
	__u64 val[nr_possible_cpus];
	int i, hits = 0, nthreads = 0;
	__u32 key;

	memset(val, 0, sizeof(val));

	for (key = 0; key < KEY_RANGE; key++)
		bpf_map_delete_elem(map_fd, &key);

	for (i = 0; i < nr_cpus; i++) {
		args[i] = (struct refill_arg){
			.map_fd = map_fd,
			.cpu = i,
			.per_cpu_quota = per_cpu_quota,
		};
		if (pthread_create(&threads[nthreads], NULL, refill_thread, &args[i]) == 0)
			nthreads++;
	}
	for (i = 0; i < nthreads; i++)
		pthread_join(threads[i], NULL);

	for (i = 0; i < nr_cpus; i++)
		if (args[i].update_errors)
			return -ENOMEM;

	for (key = 0; key < (__u32)total; key++)
		if (bpf_map_lookup_elem(map_fd, &key, val) == 0)
			hits++;

	return hits == total ? 0 : -EIO;
}

static void run_variant(enum bpf_map_type type, __u32 map_flags, const char *name)
{
	struct perf_event_attr attr = {
		.size = sizeof(attr),
		.type = PERF_TYPE_HARDWARE,
		.config = PERF_COUNT_HW_CPU_CYCLES,
		.freq = 1,
	};
	int nr_cpus, max_cpus = 64;
	struct bpf_link *links[max_cpus];
	pthread_t threads[max_cpus];
	struct hammer_arg args[max_cpus];
	struct lru_lock_nmi *skel = NULL;
	int map_fd, i, err, nr_threads = 0, pmu_fd = -1;
	__u64 deadline;

	nr_cpus = libbpf_num_possible_cpus();
	if (!ASSERT_GT(nr_cpus, 0, "num_cpus"))
		return;

	if (nr_cpus > max_cpus)
		nr_cpus = max_cpus;

	if (!test__start_subtest(name))
		return;

	memset(links, 0, sizeof(links));
	skel = lru_lock_nmi__open();
	if (!ASSERT_OK_PTR(skel, "skel_open"))
		goto cleanup;

	err = bpf_map__set_type(skel->maps.lru_map, type);
	if (!ASSERT_OK(err, "set_type"))
		goto cleanup;
	err = bpf_map__set_map_flags(skel->maps.lru_map, map_flags);
	if (!ASSERT_OK(err, "set_flags"))
		goto cleanup;
	err = bpf_map__set_max_entries(skel->maps.lru_map, MAP_ENTRIES);
	if (!ASSERT_OK(err, "set_max_entries"))
		goto cleanup;

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

	skel->bss->hits = 0;
	map_fd = bpf_map__fd(skel->maps.lru_map);
	attr.sample_freq = read_perf_max_sample_freq();

	for (i = 0; i < nr_cpus; i++) {
		pmu_fd = syscall(__NR_perf_event_open, &attr, -1, i, -1, 0);
		if (pmu_fd < 0) {
			if (i == 0 &&
			    (errno == ENOENT || errno == EOPNOTSUPP)) {
				test__skip();
				goto cleanup;
			}
			continue;
		}
		/* libbpf takes ownership of pfd on success */
		links[i] = bpf_program__attach_perf_event(skel->progs.oncpu, pmu_fd);
		if (!links[i])
			close(pmu_fd);
	}

	deadline = get_time_ns() + STRESS_NS;
	for (i = 0; i < nr_cpus; i++) {
		args[i].map_fd = map_fd;
		args[i].cpu = i;
		args[i].deadline_ns = deadline;
		if (pthread_create(&threads[nr_threads], NULL, hammer_thread, &args[i]) == 0)
			nr_threads++;
	}
	for (i = 0; i < nr_threads; i++)
		pthread_join(threads[i], NULL);

	for (i = 0; i < nr_cpus; i++) {
		if (links[i]) {
			bpf_link__destroy(links[i]);
			links[i] = NULL;
		}
	}

	ASSERT_GT(skel->bss->hits, 0, "nmi_bpf_ran");
	ASSERT_OK(drain_then_verify_capacity(map_fd, nr_cpus), "drain_then_verify_capacity");

cleanup:
	for (i = 0; i < nr_cpus; i++) {
		if (links[i])
			bpf_link__destroy(links[i]);
	}
	lru_lock_nmi__destroy(skel);
}

void serial_test_lru_lock_nmi(void)
{
	run_variant(BPF_MAP_TYPE_LRU_HASH, 0, "common_lru");
	run_variant(BPF_MAP_TYPE_LRU_HASH, BPF_F_NO_COMMON_LRU, "no_common_lru");
	run_variant(BPF_MAP_TYPE_LRU_PERCPU_HASH, 0, "percpu_lru");
}