// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2022 Meta Platforms, Inc. and affiliates. */

#include <argp.h>
#include <linux/btf.h>

#include "local_storage_bench.skel.h"
#include "bench.h"

#include <test_btf.h>

static struct {
	__u32 nr_maps;
	__u32 hashmap_nr_keys_used;
} args = {
	.nr_maps = 1000,
	.hashmap_nr_keys_used = 1000,
};

enum {
	ARG_NR_MAPS = 6000,
	ARG_HASHMAP_NR_KEYS_USED = 6001,
};

static const struct argp_option opts[] = {
	{ "nr_maps", ARG_NR_MAPS, "NR_MAPS", 0,
		"Set number of local_storage maps"},
	{ "hashmap_nr_keys_used", ARG_HASHMAP_NR_KEYS_USED, "NR_KEYS",
		0, "When doing hashmap test, set number of hashmap keys test uses"},
	{},
};

static error_t parse_arg(int key, char *arg, struct argp_state *state)
{
	long ret;

	switch (key) {
	case ARG_NR_MAPS:
		ret = strtol(arg, NULL, 10);
		if (ret < 1 || ret > UINT_MAX) {
			fprintf(stderr, "invalid nr_maps");
			argp_usage(state);
		}
		args.nr_maps = ret;
		break;
	case ARG_HASHMAP_NR_KEYS_USED:
		ret = strtol(arg, NULL, 10);
		if (ret < 1 || ret > UINT_MAX) {
			fprintf(stderr, "invalid hashmap_nr_keys_used");
			argp_usage(state);
		}
		args.hashmap_nr_keys_used = ret;
		break;
	default:
		return ARGP_ERR_UNKNOWN;
	}

	return 0;
}

const struct argp bench_local_storage_argp = {
	.options = opts,
	.parser = parse_arg,
};

/* Keep in sync w/ array of maps in bpf */
#define MAX_NR_MAPS 1000
/* keep in sync w/ same define in bpf */
#define HASHMAP_SZ 4194304

static void validate(void)
{
	if (env.producer_cnt != 1) {
		fprintf(stderr, "benchmark doesn't support multi-producer!\n");
		exit(1);
	}
	if (env.consumer_cnt != 0) {
		fprintf(stderr, "benchmark doesn't support consumer!\n");
		exit(1);
	}

	if (args.nr_maps > MAX_NR_MAPS) {
		fprintf(stderr, "nr_maps must be <= 1000\n");
		exit(1);
	}

	if (args.hashmap_nr_keys_used > HASHMAP_SZ) {
		fprintf(stderr, "hashmap_nr_keys_used must be <= %u\n", HASHMAP_SZ);
		exit(1);
	}
}

static struct {
	struct local_storage_bench *skel;
	void *bpf_obj;
	struct bpf_map *array_of_maps;
} ctx;

static void prepopulate_hashmap(int fd)
{
	int i, key, val;

	/* local_storage gets will have BPF_LOCAL_STORAGE_GET_F_CREATE flag set, so
	 * populate the hashmap for a similar comparison
	 */
	for (i = 0; i < HASHMAP_SZ; i++) {
		key = val = i;
		if (bpf_map_update_elem(fd, &key, &val, 0)) {
			fprintf(stderr, "Error prepopulating hashmap (key %d)\n", key);
			exit(1);
		}
	}
}

static void __setup(struct bpf_program *prog, bool hashmap)
{
	struct bpf_map *inner_map;
	int i, fd, mim_fd, err;

	LIBBPF_OPTS(bpf_map_create_opts, create_opts);

	if (!hashmap)
		create_opts.map_flags = BPF_F_NO_PREALLOC;

	ctx.skel->rodata->num_maps = args.nr_maps;
	ctx.skel->rodata->hashmap_num_keys = args.hashmap_nr_keys_used;
	inner_map = bpf_map__inner_map(ctx.array_of_maps);
	create_opts.btf_key_type_id = bpf_map__btf_key_type_id(inner_map);
	create_opts.btf_value_type_id = bpf_map__btf_value_type_id(inner_map);

	err = local_storage_bench__load(ctx.skel);
	if (err) {
		fprintf(stderr, "Error loading skeleton\n");
		goto err_out;
	}

	create_opts.btf_fd = bpf_object__btf_fd(ctx.skel->obj);

	mim_fd = bpf_map__fd(ctx.array_of_maps);
	if (mim_fd < 0) {
		fprintf(stderr, "Error getting map_in_map fd\n");
		goto err_out;
	}

	for (i = 0; i < args.nr_maps; i++) {
		if (hashmap)
			fd = bpf_map_create(BPF_MAP_TYPE_HASH, NULL, sizeof(int),
					    sizeof(int), HASHMAP_SZ, &create_opts);
		else
			fd = bpf_map_create(BPF_MAP_TYPE_TASK_STORAGE, NULL, sizeof(int),
					    sizeof(int), 0, &create_opts);
		if (fd < 0) {
			fprintf(stderr, "Error creating map %d: %d\n", i, fd);
			goto err_out;
		}

		if (hashmap)
			prepopulate_hashmap(fd);

		err = bpf_map_update_elem(mim_fd, &i, &fd, 0);
		if (err) {
			fprintf(stderr, "Error updating array-of-maps w/ map %d\n", i);
			goto err_out;
		}
	}

	if (!bpf_program__attach(prog)) {
		fprintf(stderr, "Error attaching bpf program\n");
		goto err_out;
	}

	return;
err_out:
	exit(1);
}

static void hashmap_setup(void)
{
	struct local_storage_bench *skel;

	setup_libbpf();

	skel = local_storage_bench__open();
	ctx.skel = skel;
	ctx.array_of_maps = skel->maps.array_of_hash_maps;
	skel->rodata->use_hashmap = 1;
	skel->rodata->interleave = 0;

	__setup(skel->progs.get_local, true);
}

static void local_storage_cache_get_setup(void)
{
	struct local_storage_bench *skel;

	setup_libbpf();

	skel = local_storage_bench__open();
	ctx.skel = skel;
	ctx.array_of_maps = skel->maps.array_of_local_storage_maps;
	skel->rodata->use_hashmap = 0;
	skel->rodata->interleave = 0;

	__setup(skel->progs.get_local, false);
}

static void local_storage_cache_get_interleaved_setup(void)
{
	struct local_storage_bench *skel;

	setup_libbpf();

	skel = local_storage_bench__open();
	ctx.skel = skel;
	ctx.array_of_maps = skel->maps.array_of_local_storage_maps;
	skel->rodata->use_hashmap = 0;
	skel->rodata->interleave = 1;

	__setup(skel->progs.get_local, false);
}

static void measure(struct bench_res *res)
{
	res->hits = atomic_swap(&ctx.skel->bss->hits, 0);
	res->important_hits = atomic_swap(&ctx.skel->bss->important_hits, 0);
}

static inline void trigger_bpf_program(void)
{
	syscall(__NR_getpgid);
}

static void *producer(void *input)
{
	while (true)
		trigger_bpf_program();

	return NULL;
}

/* cache sequential and interleaved get benchs test local_storage get
 * performance, specifically they demonstrate performance cliff of
 * current list-plus-cache local_storage model.
 *
 * cache sequential get: call bpf_task_storage_get on n maps in order
 * cache interleaved get: like "sequential get", but interleave 4 calls to the
 *	'important' map (idx 0 in array_of_maps) for every 10 calls. Goal
 *	is to mimic environment where many progs are accessing their local_storage
 *	maps, with 'our' prog needing to access its map more often than others
 */
const struct bench bench_local_storage_cache_seq_get = {
	.name = "local-storage-cache-seq-get",
	.argp = &bench_local_storage_argp,
	.validate = validate,
	.setup = local_storage_cache_get_setup,
	.producer_thread = producer,
	.measure = measure,
	.report_progress = local_storage_report_progress,
	.report_final = local_storage_report_final,
};

const struct bench bench_local_storage_cache_interleaved_get = {
	.name = "local-storage-cache-int-get",
	.argp = &bench_local_storage_argp,
	.validate = validate,
	.setup = local_storage_cache_get_interleaved_setup,
	.producer_thread = producer,
	.measure = measure,
	.report_progress = local_storage_report_progress,
	.report_final = local_storage_report_final,
};

const struct bench bench_local_storage_cache_hashmap_control = {
	.name = "local-storage-cache-hashmap-control",
	.argp = &bench_local_storage_argp,
	.validate = validate,
	.setup = hashmap_setup,
	.producer_thread = producer,
	.measure = measure,
	.report_progress = local_storage_report_progress,
	.report_final = local_storage_report_final,
};