// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2025 Cloudflare */

/*
 * All of these benchmarks operate on tries with keys in the range
 * [0, args.nr_entries), i.e. there are no gaps or partially filled
 * branches of the trie for any key < args.nr_entries.
 *
 * This gives an idea of worst-case behaviour.
 */

#include <argp.h>
#include <linux/time64.h>
#include <linux/if_ether.h>
#include "lpm_trie_bench.skel.h"
#include "lpm_trie_map.skel.h"
#include "bench.h"
#include "testing_helpers.h"
#include "progs/lpm_trie.h"

static struct ctx {
        struct lpm_trie_bench *bench;
} ctx;

static struct {
        __u32 nr_entries;
        __u32 prefixlen;
        bool random;
} args = {
        .nr_entries = 0,
        .prefixlen = 32,
        .random = false,
};

enum {
        ARG_NR_ENTRIES = 9000,
        ARG_PREFIX_LEN,
        ARG_RANDOM,
};

static const struct argp_option opts[] = {
        { "nr_entries", ARG_NR_ENTRIES, "NR_ENTRIES", 0,
          "Number of unique entries in the LPM trie" },
        { "prefix_len", ARG_PREFIX_LEN, "PREFIX_LEN", 0,
          "Number of prefix bits to use in the LPM trie" },
        { "random", ARG_RANDOM, NULL, 0, "Access random keys during op" },
        {},
};

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

        switch (key) {
        case ARG_NR_ENTRIES:
                ret = strtol(arg, NULL, 10);
                if (ret < 1 || ret > UINT_MAX) {
                        fprintf(stderr, "Invalid nr_entries count.");
                        argp_usage(state);
                }
                args.nr_entries = ret;
                break;
        case ARG_PREFIX_LEN:
                ret = strtol(arg, NULL, 10);
                if (ret < 1 || ret > UINT_MAX) {
                        fprintf(stderr, "Invalid prefix_len value.");
                        argp_usage(state);
                }
                args.prefixlen = ret;
                break;
        case ARG_RANDOM:
                args.random = true;
                break;
        default:
                return ARGP_ERR_UNKNOWN;
        }
        return 0;
}

const struct argp bench_lpm_trie_map_argp = {
        .options = opts,
        .parser = lpm_parse_arg,
};

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

        if (args.nr_entries == 0) {
                fprintf(stderr, "Missing --nr_entries parameter\n");
                exit(1);
        }

        if ((1UL << args.prefixlen) < args.nr_entries) {
                fprintf(stderr, "prefix_len value too small for nr_entries\n");
                exit(1);
        }
}

static void lpm_insert_validate(void)
{
        validate_common();

        if (env.producer_cnt != 1) {
                fprintf(stderr, "lpm-trie-insert requires a single producer\n");
                exit(1);
        }

        if (args.random) {
                fprintf(stderr, "lpm-trie-insert does not support --random\n");
                exit(1);
        }
}

static void lpm_delete_validate(void)
{
        validate_common();

        if (env.producer_cnt != 1) {
                fprintf(stderr, "lpm-trie-delete requires a single producer\n");
                exit(1);
        }

        if (args.random) {
                fprintf(stderr, "lpm-trie-delete does not support --random\n");
                exit(1);
        }
}

static void lpm_free_validate(void)
{
        validate_common();

        if (env.producer_cnt != 1) {
                fprintf(stderr, "lpm-trie-free requires a single producer\n");
                exit(1);
        }

        if (args.random) {
                fprintf(stderr, "lpm-trie-free does not support --random\n");
                exit(1);
        }
}

static struct trie_key *keys;
static __u32 *vals;

static void fill_map(int map_fd)
{
        int err;

        DECLARE_LIBBPF_OPTS(bpf_map_batch_opts, opts,
                .elem_flags = 0,
                .flags = 0,
        );

        err = bpf_map_update_batch(map_fd, keys, vals, &args.nr_entries, &opts);
        if (err) {
                fprintf(stderr, "failed to batch update keys to map: %d\n",
                        -err);
                exit(1);
        }
}

static void empty_map(int map_fd)
{
        int err;

        DECLARE_LIBBPF_OPTS(bpf_map_batch_opts, opts,
                .elem_flags = 0,
                .flags = 0,
        );

        err = bpf_map_delete_batch(map_fd, keys, &args.nr_entries, &opts);
        if (err) {
                fprintf(stderr, "failed to batch delete keys for map: %d\n",
                        -err);
                exit(1);
        }
}

static void attach_prog(void)
{
        int i;

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

        ctx.bench->bss->nr_entries = args.nr_entries;
        ctx.bench->bss->prefixlen = args.prefixlen;
        ctx.bench->bss->random = args.random;

        if (lpm_trie_bench__attach(ctx.bench)) {
                fprintf(stderr, "failed to attach skeleton\n");
                exit(1);
        }

        keys = calloc(args.nr_entries, sizeof(*keys));
        vals = calloc(args.nr_entries, sizeof(*vals));

        for (i = 0; i < args.nr_entries; i++) {
                struct trie_key *k = &keys[i];
                __u32 *v = &vals[i];

                k->prefixlen = args.prefixlen;
                k->data = i;
                *v = 1;
        }
}

static void attach_prog_and_fill_map(void)
{
        int fd;

        attach_prog();

        fd = bpf_map__fd(ctx.bench->maps.trie_map);
        fill_map(fd);
}

static void lpm_noop_setup(void)
{
        attach_prog();
        ctx.bench->bss->op = LPM_OP_NOOP;
}

static void lpm_baseline_setup(void)
{
        attach_prog();
        ctx.bench->bss->op = LPM_OP_BASELINE;
}

static void lpm_lookup_setup(void)
{
        attach_prog_and_fill_map();
        ctx.bench->bss->op = LPM_OP_LOOKUP;
}

static void lpm_insert_setup(void)
{
        attach_prog();
        ctx.bench->bss->op = LPM_OP_INSERT;
}

static void lpm_update_setup(void)
{
        attach_prog_and_fill_map();
        ctx.bench->bss->op = LPM_OP_UPDATE;
}

static void lpm_delete_setup(void)
{
        attach_prog_and_fill_map();
        ctx.bench->bss->op = LPM_OP_DELETE;
}

static void lpm_free_setup(void)
{
        attach_prog();
        ctx.bench->bss->op = LPM_OP_FREE;
}

static void lpm_measure(struct bench_res *res)
{
        res->hits = atomic_swap(&ctx.bench->bss->hits, 0);
        res->duration_ns = atomic_swap(&ctx.bench->bss->duration_ns, 0);
}

static void bench_reinit_map(void)
{
        int fd = bpf_map__fd(ctx.bench->maps.trie_map);

        switch (ctx.bench->bss->op) {
        case LPM_OP_INSERT:
                /* trie_map needs to be emptied */
                empty_map(fd);
                break;
        case LPM_OP_DELETE:
                /* trie_map needs to be refilled */
                fill_map(fd);
                break;
        default:
                fprintf(stderr, "Unexpected REINIT return code for op %d\n",
                                ctx.bench->bss->op);
                exit(1);
        }
}

/* For NOOP, BASELINE, LOOKUP, INSERT, UPDATE, and DELETE */
static void *lpm_producer(void *unused __always_unused)
{
        int err;
        char in[ETH_HLEN]; /* unused */

        LIBBPF_OPTS(bpf_test_run_opts, opts, .data_in = in,
                    .data_size_in = sizeof(in), .repeat = 1, );

        while (true) {
                int fd = bpf_program__fd(ctx.bench->progs.run_bench);
                err = bpf_prog_test_run_opts(fd, &opts);
                if (err) {
                        fprintf(stderr, "failed to run BPF prog: %d\n", err);
                        exit(1);
                }

                /* Check for kernel error code */
                if ((int)opts.retval < 0) {
                        fprintf(stderr, "BPF prog returned error: %d\n",
                                opts.retval);
                        exit(1);
                }

                switch (opts.retval) {
                case LPM_BENCH_SUCCESS:
                        break;
                case LPM_BENCH_REINIT_MAP:
                        bench_reinit_map();
                        break;
                default:
                        fprintf(stderr, "Unexpected BPF prog return code %d for op %d\n",
                                        opts.retval, ctx.bench->bss->op);
                        exit(1);
                }
        }

        return NULL;
}

static void *lpm_free_producer(void *unused __always_unused)
{
        while (true) {
                struct lpm_trie_map *skel;

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

                fill_map(bpf_map__fd(skel->maps.trie_free_map));
                lpm_trie_map__destroy(skel);
        }

        return NULL;
}

/*
 * The standard bench op_report_*() functions assume measurements are
 * taken over a 1-second interval but operations that modify the map
 * (INSERT, DELETE, and FREE) cannot run indefinitely without
 * "resetting" the map to the initial state. Depending on the size of
 * the map, this likely needs to happen before the 1-second timer fires.
 *
 * Calculate the fraction of a second over which the op measurement was
 * taken (to ignore any time spent doing the reset) and report the
 * throughput results per second.
 */
static void frac_second_report_progress(int iter, struct bench_res *res,
                                        long delta_ns, double rate_divisor,
                                        char rate)
{
        double hits_per_sec, hits_per_prod;

        hits_per_sec = res->hits / rate_divisor /
                (res->duration_ns / (double)NSEC_PER_SEC);
        hits_per_prod = hits_per_sec / env.producer_cnt;

        printf("Iter %3d (%7.3lfus): ", iter,
               (delta_ns - NSEC_PER_SEC) / 1000.0);
        printf("hits %8.3lf%c/s (%7.3lf%c/prod)\n", hits_per_sec, rate,
               hits_per_prod, rate);
}

static void frac_second_report_final(struct bench_res res[], int res_cnt,
                                     double lat_divisor, double rate_divisor,
                                     char rate, const char *unit)
{
        double hits_mean = 0.0, hits_stddev = 0.0;
        double latency = 0.0;
        int i;

        for (i = 0; i < res_cnt; i++) {
                double val = res[i].hits / rate_divisor /
                             (res[i].duration_ns / (double)NSEC_PER_SEC);
                hits_mean += val / (0.0 + res_cnt);
                latency += res[i].duration_ns / res[i].hits / (0.0 + res_cnt);
        }

        if (res_cnt > 1) {
                for (i = 0; i < res_cnt; i++) {
                        double val =
                                res[i].hits / rate_divisor /
                                (res[i].duration_ns / (double)NSEC_PER_SEC);
                        hits_stddev += (hits_mean - val) * (hits_mean - val) /
                                       (res_cnt - 1.0);
                }

                hits_stddev = sqrt(hits_stddev);
        }
        printf("Summary: throughput %8.3lf \u00B1 %5.3lf %c ops/s (%7.3lf%c ops/prod), ",
               hits_mean, hits_stddev, rate, hits_mean / env.producer_cnt,
               rate);
        printf("latency %8.3lf %s/op\n",
               latency / lat_divisor / env.producer_cnt, unit);
}

static void insert_ops_report_progress(int iter, struct bench_res *res,
                                       long delta_ns)
{
        double rate_divisor = 1000000.0;
        char rate = 'M';

        frac_second_report_progress(iter, res, delta_ns, rate_divisor, rate);
}

static void delete_ops_report_progress(int iter, struct bench_res *res,
                                       long delta_ns)
{
        double rate_divisor = 1000000.0;
        char rate = 'M';

        frac_second_report_progress(iter, res, delta_ns, rate_divisor, rate);
}

static void free_ops_report_progress(int iter, struct bench_res *res,
                                     long delta_ns)
{
        double rate_divisor = 1000.0;
        char rate = 'K';

        frac_second_report_progress(iter, res, delta_ns, rate_divisor, rate);
}

static void insert_ops_report_final(struct bench_res res[], int res_cnt)
{
        double lat_divisor = 1.0;
        double rate_divisor = 1000000.0;
        const char *unit = "ns";
        char rate = 'M';

        frac_second_report_final(res, res_cnt, lat_divisor, rate_divisor, rate,
                                 unit);
}

static void delete_ops_report_final(struct bench_res res[], int res_cnt)
{
        double lat_divisor = 1.0;
        double rate_divisor = 1000000.0;
        const char *unit = "ns";
        char rate = 'M';

        frac_second_report_final(res, res_cnt, lat_divisor, rate_divisor, rate,
                                 unit);
}

static void free_ops_report_final(struct bench_res res[], int res_cnt)
{
        double lat_divisor = 1000000.0;
        double rate_divisor = 1000.0;
        const char *unit = "ms";
        char rate = 'K';

        frac_second_report_final(res, res_cnt, lat_divisor, rate_divisor, rate,
                                 unit);
}

/* noop bench measures harness-overhead */
const struct bench bench_lpm_trie_noop = {
        .name = "lpm-trie-noop",
        .argp = &bench_lpm_trie_map_argp,
        .validate = validate_common,
        .setup = lpm_noop_setup,
        .producer_thread = lpm_producer,
        .measure = lpm_measure,
        .report_progress = ops_report_progress,
        .report_final = ops_report_final,
};

/* baseline overhead for lookup and update */
const struct bench bench_lpm_trie_baseline = {
        .name = "lpm-trie-baseline",
        .argp = &bench_lpm_trie_map_argp,
        .validate = validate_common,
        .setup = lpm_baseline_setup,
        .producer_thread = lpm_producer,
        .measure = lpm_measure,
        .report_progress = ops_report_progress,
        .report_final = ops_report_final,
};

/* measure cost of doing a lookup on existing entries in a full trie */
const struct bench bench_lpm_trie_lookup = {
        .name = "lpm-trie-lookup",
        .argp = &bench_lpm_trie_map_argp,
        .validate = validate_common,
        .setup = lpm_lookup_setup,
        .producer_thread = lpm_producer,
        .measure = lpm_measure,
        .report_progress = ops_report_progress,
        .report_final = ops_report_final,
};

/* measure cost of inserting new entries into an empty trie */
const struct bench bench_lpm_trie_insert = {
        .name = "lpm-trie-insert",
        .argp = &bench_lpm_trie_map_argp,
        .validate = lpm_insert_validate,
        .setup = lpm_insert_setup,
        .producer_thread = lpm_producer,
        .measure = lpm_measure,
        .report_progress = insert_ops_report_progress,
        .report_final = insert_ops_report_final,
};

/* measure cost of updating existing entries in a full trie */
const struct bench bench_lpm_trie_update = {
        .name = "lpm-trie-update",
        .argp = &bench_lpm_trie_map_argp,
        .validate = validate_common,
        .setup = lpm_update_setup,
        .producer_thread = lpm_producer,
        .measure = lpm_measure,
        .report_progress = ops_report_progress,
        .report_final = ops_report_final,
};

/* measure cost of deleting existing entries from a full trie */
const struct bench bench_lpm_trie_delete = {
        .name = "lpm-trie-delete",
        .argp = &bench_lpm_trie_map_argp,
        .validate = lpm_delete_validate,
        .setup = lpm_delete_setup,
        .producer_thread = lpm_producer,
        .measure = lpm_measure,
        .report_progress = delete_ops_report_progress,
        .report_final = delete_ops_report_final,
};

/* measure cost of freeing a full trie */
const struct bench bench_lpm_trie_free = {
        .name = "lpm-trie-free",
        .argp = &bench_lpm_trie_map_argp,
        .validate = lpm_free_validate,
        .setup = lpm_free_setup,
        .producer_thread = lpm_free_producer,
        .measure = lpm_measure,
        .report_progress = free_ops_report_progress,
        .report_final = free_ops_report_final,
};