// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2023 Isovalent */

#include <sys/random.h>
#include <argp.h>
#include "bench.h"
#include "bpf_hashmap_lookup.skel.h"
#include "bpf_util.h"

/* BPF triggering benchmarks */
static struct ctx {
        struct bpf_hashmap_lookup *skel;
} ctx;

/* only available to kernel, so define it here */
#define BPF_MAX_LOOPS (1<<23)

#define MAX_KEY_SIZE 1024 /* the size of the key map */

static struct {
        __u32 key_size;
        __u32 map_flags;
        __u32 max_entries;
        __u32 nr_entries;
        __u32 nr_loops;
} args = {
        .key_size = 4,
        .map_flags = 0,
        .max_entries = 1000,
        .nr_entries = 500,
        .nr_loops = 1000000,
};

enum {
        ARG_KEY_SIZE = 8001,
        ARG_MAP_FLAGS,
        ARG_MAX_ENTRIES,
        ARG_NR_ENTRIES,
        ARG_NR_LOOPS,
};

static const struct argp_option opts[] = {
        { "key_size", ARG_KEY_SIZE, "KEY_SIZE", 0,
          "The hashmap key size (max 1024)"},
        { "map_flags", ARG_MAP_FLAGS, "MAP_FLAGS", 0,
          "The hashmap flags passed to BPF_MAP_CREATE"},
        { "max_entries", ARG_MAX_ENTRIES, "MAX_ENTRIES", 0,
          "The hashmap max entries"},
        { "nr_entries", ARG_NR_ENTRIES, "NR_ENTRIES", 0,
          "The number of entries to insert/lookup"},
        { "nr_loops", ARG_NR_LOOPS, "NR_LOOPS", 0,
          "The number of loops for the benchmark"},
        {},
};

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

        switch (key) {
        case ARG_KEY_SIZE:
                ret = strtol(arg, NULL, 10);
                if (ret < 1 || ret > MAX_KEY_SIZE) {
                        fprintf(stderr, "invalid key_size");
                        argp_usage(state);
                }
                args.key_size = ret;
                break;
        case ARG_MAP_FLAGS:
                ret = strtol(arg, NULL, 0);
                if (ret < 0 || ret > UINT_MAX) {
                        fprintf(stderr, "invalid map_flags");
                        argp_usage(state);
                }
                args.map_flags = ret;
                break;
        case ARG_MAX_ENTRIES:
                ret = strtol(arg, NULL, 10);
                if (ret < 1 || ret > UINT_MAX) {
                        fprintf(stderr, "invalid max_entries");
                        argp_usage(state);
                }
                args.max_entries = ret;
                break;
        case ARG_NR_ENTRIES:
                ret = strtol(arg, NULL, 10);
                if (ret < 1 || ret > UINT_MAX) {
                        fprintf(stderr, "invalid nr_entries");
                        argp_usage(state);
                }
                args.nr_entries = ret;
                break;
        case ARG_NR_LOOPS:
                ret = strtol(arg, NULL, 10);
                if (ret < 1 || ret > BPF_MAX_LOOPS) {
                        fprintf(stderr, "invalid nr_loops: %ld (min=1 max=%u)\n",
                                ret, BPF_MAX_LOOPS);
                        argp_usage(state);
                }
                args.nr_loops = ret;
                break;
        default:
                return ARGP_ERR_UNKNOWN;
        }

        return 0;
}

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

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

        if (args.nr_entries > args.max_entries) {
                fprintf(stderr, "args.nr_entries is too big! (max %u, got %u)\n",
                        args.max_entries, args.nr_entries);
                exit(1);
        }
}

static void *producer(void *input)
{
        while (true) {
                /* trigger the bpf program */
                syscall(__NR_getpgid);
        }
        return NULL;
}

static void measure(struct bench_res *res)
{
}

static inline void patch_key(u32 i, u32 *key)
{
#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
        *key = i + 1;
#else
        *key = __builtin_bswap32(i + 1);
#endif
        /* the rest of key is random */
}

static void setup(void)
{
        struct bpf_link *link;
        int map_fd;
        int ret;
        int i;

        setup_libbpf();

        ctx.skel = bpf_hashmap_lookup__open();
        if (!ctx.skel) {
                fprintf(stderr, "failed to open skeleton\n");
                exit(1);
        }

        bpf_map__set_max_entries(ctx.skel->maps.hash_map_bench, args.max_entries);
        bpf_map__set_key_size(ctx.skel->maps.hash_map_bench, args.key_size);
        bpf_map__set_value_size(ctx.skel->maps.hash_map_bench, 8);
        bpf_map__set_map_flags(ctx.skel->maps.hash_map_bench, args.map_flags);

        ctx.skel->bss->nr_entries = args.nr_entries;
        ctx.skel->bss->nr_loops = args.nr_loops / args.nr_entries;

        if (args.key_size > 4) {
                for (i = 1; i < args.key_size/4; i++)
                        ctx.skel->bss->key[i] = 2654435761 * i;
        }

        ret = bpf_hashmap_lookup__load(ctx.skel);
        if (ret) {
                bpf_hashmap_lookup__destroy(ctx.skel);
                fprintf(stderr, "failed to load map: %s", strerror(-ret));
                exit(1);
        }

        /* fill in the hash_map */
        map_fd = bpf_map__fd(ctx.skel->maps.hash_map_bench);
        for (u64 i = 0; i < args.nr_entries; i++) {
                patch_key(i, ctx.skel->bss->key);
                bpf_map_update_elem(map_fd, ctx.skel->bss->key, &i, BPF_ANY);
        }

        link = bpf_program__attach(ctx.skel->progs.benchmark);
        if (!link) {
                fprintf(stderr, "failed to attach program!\n");
                exit(1);
        }
}

static inline double events_from_time(u64 time)
{
        if (time)
                return args.nr_loops * 1000000000llu / time / 1000000.0L;

        return 0;
}

static int compute_events(u64 *times, double *events_mean, double *events_stddev, u64 *mean_time)
{
        int i, n = 0;

        *events_mean = 0;
        *events_stddev = 0;
        *mean_time = 0;

        for (i = 0; i < 32; i++) {
                if (!times[i])
                        break;
                *mean_time += times[i];
                *events_mean += events_from_time(times[i]);
                n += 1;
        }
        if (!n)
                return 0;

        *mean_time /= n;
        *events_mean /= n;

        if (n > 1) {
                for (i = 0; i < n; i++) {
                        double events_i = *events_mean - events_from_time(times[i]);
                        *events_stddev += events_i * events_i / (n - 1);
                }
                *events_stddev = sqrt(*events_stddev);
        }

        return n;
}

static void hashmap_report_final(struct bench_res res[], int res_cnt)
{
        unsigned int nr_cpus = bpf_num_possible_cpus();
        double events_mean, events_stddev;
        u64 mean_time;
        int i, n;

        for (i = 0; i < nr_cpus; i++) {
                n = compute_events(ctx.skel->bss->percpu_times[i], &events_mean,
                                   &events_stddev, &mean_time);
                if (n == 0)
                        continue;

                if (env.quiet) {
                        /* we expect only one cpu to be present */
                        if (env.affinity)
                                printf("%.3lf\n", events_mean);
                        else
                                printf("cpu%02d %.3lf\n", i, events_mean);
                } else {
                        printf("cpu%02d: lookup %.3lfM ± %.3lfM events/sec"
                               " (approximated from %d samples of ~%lums)\n",
                               i, events_mean, 2*events_stddev,
                               n, mean_time / 1000000);
                }
        }
}

const struct bench bench_bpf_hashmap_lookup = {
        .name = "bpf-hashmap-lookup",
        .argp = &bench_hashmap_lookup_argp,
        .validate = validate,
        .setup = setup,
        .producer_thread = producer,
        .measure = measure,
        .report_progress = NULL,
        .report_final = hashmap_report_final,
};