// SPDX-License-Identifier: GPL-2.0

/*
 * Test module for lockless object pool
 *
 * Copyright: wuqiang.matt@bytedance.com
 */

#include <linux/errno.h>
#include <linux/module.h>
#include <linux/moduleparam.h>
#include <linux/completion.h>
#include <linux/kthread.h>
#include <linux/slab.h>
#include <linux/vmalloc.h>
#include <linux/delay.h>
#include <linux/hrtimer.h>
#include <linux/objpool.h>

#define OT_NR_MAX_BULK (16)

/* memory usage */
struct ot_mem_stat {
        atomic_long_t alloc;
        atomic_long_t free;
};

/* object allocation results */
struct ot_obj_stat {
        unsigned long nhits;
        unsigned long nmiss;
};

/* control & results per testcase */
struct ot_data {
        struct rw_semaphore start;
        struct completion wait;
        struct completion rcu;
        atomic_t nthreads ____cacheline_aligned_in_smp;
        atomic_t stop ____cacheline_aligned_in_smp;
        struct ot_mem_stat kmalloc;
        struct ot_mem_stat vmalloc;
        struct ot_obj_stat objects;
        u64    duration;
};

/* testcase */
struct ot_test {
        int async; /* synchronous or asynchronous */
        int mode; /* only mode 0 supported */
        int objsz; /* object size */
        int duration; /* ms */
        int delay; /* ms */
        int bulk_normal;
        int bulk_irq;
        unsigned long hrtimer; /* ms */
        const char *name;
        struct ot_data data;
};

/* per-cpu worker */
struct ot_item {
        struct objpool_head *pool; /* pool head */
        struct ot_test *test; /* test parameters */

        void (*worker)(struct ot_item *item, int irq);

        /* hrtimer control */
        ktime_t hrtcycle;
        struct hrtimer hrtimer;

        int bulk[2]; /* for thread and irq */
        int delay;
        u32 niters;

        /* summary per thread */
        struct ot_obj_stat stat[2]; /* thread and irq */
        u64 duration;
};

/*
 * memory leakage checking
 */

static void *ot_kzalloc(struct ot_test *test, long size)
{
        void *ptr = kzalloc(size, GFP_KERNEL);

        if (ptr)
                atomic_long_add(size, &test->data.kmalloc.alloc);
        return ptr;
}

static void ot_kfree(struct ot_test *test, void *ptr, long size)
{
        if (!ptr)
                return;
        atomic_long_add(size, &test->data.kmalloc.free);
        kfree(ptr);
}

static void ot_mem_report(struct ot_test *test)
{
        long alloc, free;

        pr_info("memory allocation summary for %s\n", test->name);

        alloc = atomic_long_read(&test->data.kmalloc.alloc);
        free = atomic_long_read(&test->data.kmalloc.free);
        pr_info("  kmalloc: %lu - %lu = %lu\n", alloc, free, alloc - free);

        alloc = atomic_long_read(&test->data.vmalloc.alloc);
        free = atomic_long_read(&test->data.vmalloc.free);
        pr_info("  vmalloc: %lu - %lu = %lu\n", alloc, free, alloc - free);
}

/* user object instance */
struct ot_node {
        void *owner;
        unsigned long data;
        unsigned long refs;
        unsigned long payload[32];
};

/* user objpool manager */
struct ot_context {
        struct objpool_head pool; /* objpool head */
        struct ot_test *test; /* test parameters */
        void *ptr; /* user pool buffer */
        unsigned long size; /* buffer size */
        struct rcu_head rcu;
};

static DEFINE_PER_CPU(struct ot_item, ot_pcup_items);

static int ot_init_data(struct ot_data *data)
{
        memset(data, 0, sizeof(*data));
        init_rwsem(&data->start);
        init_completion(&data->wait);
        init_completion(&data->rcu);
        atomic_set(&data->nthreads, 1);

        return 0;
}

static int ot_init_node(void *nod, void *context)
{
        struct ot_context *sop = context;
        struct ot_node *on = nod;

        on->owner = &sop->pool;
        return 0;
}

static enum hrtimer_restart ot_hrtimer_handler(struct hrtimer *hrt)
{
        struct ot_item *item = container_of(hrt, struct ot_item, hrtimer);
        struct ot_test *test = item->test;

        if (atomic_read_acquire(&test->data.stop))
                return HRTIMER_NORESTART;

        /* do bulk-testings for objects pop/push */
        item->worker(item, 1);

        hrtimer_forward_now(hrt, item->hrtcycle);
        return HRTIMER_RESTART;
}

static void ot_start_hrtimer(struct ot_item *item)
{
        if (!item->test->hrtimer)
                return;
        hrtimer_start(&item->hrtimer, item->hrtcycle, HRTIMER_MODE_REL);
}

static void ot_stop_hrtimer(struct ot_item *item)
{
        if (!item->test->hrtimer)
                return;
        hrtimer_cancel(&item->hrtimer);
}

static int ot_init_hrtimer(struct ot_item *item, unsigned long hrtimer)
{
        struct hrtimer *hrt = &item->hrtimer;

        if (!hrtimer)
                return -ENOENT;

        item->hrtcycle = ktime_set(0, hrtimer * 1000000UL);
        hrtimer_setup(hrt, ot_hrtimer_handler, CLOCK_MONOTONIC, HRTIMER_MODE_REL);
        return 0;
}

static int ot_init_cpu_item(struct ot_item *item,
                        struct ot_test *test,
                        struct objpool_head *pool,
                        void (*worker)(struct ot_item *, int))
{
        memset(item, 0, sizeof(*item));
        item->pool = pool;
        item->test = test;
        item->worker = worker;

        item->bulk[0] = test->bulk_normal;
        item->bulk[1] = test->bulk_irq;
        item->delay = test->delay;

        /* initialize hrtimer */
        ot_init_hrtimer(item, item->test->hrtimer);
        return 0;
}

static int ot_thread_worker(void *arg)
{
        struct ot_item *item = arg;
        struct ot_test *test = item->test;
        ktime_t start;

        atomic_inc(&test->data.nthreads);
        down_read(&test->data.start);
        up_read(&test->data.start);
        start = ktime_get();
        ot_start_hrtimer(item);
        do {
                if (atomic_read_acquire(&test->data.stop))
                        break;
                /* do bulk-testings for objects pop/push */
                item->worker(item, 0);
        } while (!kthread_should_stop());
        ot_stop_hrtimer(item);
        item->duration = (u64) ktime_us_delta(ktime_get(), start);
        if (atomic_dec_and_test(&test->data.nthreads))
                complete(&test->data.wait);

        return 0;
}

static void ot_perf_report(struct ot_test *test, u64 duration)
{
        struct ot_obj_stat total, normal = {0}, irq = {0};
        int cpu, nthreads = 0;

        pr_info("\n");
        pr_info("Testing summary for %s\n", test->name);

        for_each_possible_cpu(cpu) {
                struct ot_item *item = per_cpu_ptr(&ot_pcup_items, cpu);
                if (!item->duration)
                        continue;
                normal.nhits += item->stat[0].nhits;
                normal.nmiss += item->stat[0].nmiss;
                irq.nhits += item->stat[1].nhits;
                irq.nmiss += item->stat[1].nmiss;
                pr_info("CPU: %d  duration: %lluus\n", cpu, item->duration);
                pr_info("\tthread:\t%16lu hits \t%16lu miss\n",
                        item->stat[0].nhits, item->stat[0].nmiss);
                pr_info("\tirq:   \t%16lu hits \t%16lu miss\n",
                        item->stat[1].nhits, item->stat[1].nmiss);
                pr_info("\ttotal: \t%16lu hits \t%16lu miss\n",
                        item->stat[0].nhits + item->stat[1].nhits,
                        item->stat[0].nmiss + item->stat[1].nmiss);
                nthreads++;
        }

        total.nhits = normal.nhits + irq.nhits;
        total.nmiss = normal.nmiss + irq.nmiss;

        pr_info("ALL: \tnthreads: %d  duration: %lluus\n", nthreads, duration);
        pr_info("SUM: \t%16lu hits \t%16lu miss\n",
                total.nhits, total.nmiss);

        test->data.objects = total;
        test->data.duration = duration;
}

/*
 * synchronous test cases for objpool manipulation
 */

/* objpool manipulation for synchronous mode (percpu objpool) */
static struct ot_context *ot_init_sync_m0(struct ot_test *test)
{
        struct ot_context *sop = NULL;
        int max = num_possible_cpus() << 3;
        gfp_t gfp = GFP_KERNEL;

        sop = (struct ot_context *)ot_kzalloc(test, sizeof(*sop));
        if (!sop)
                return NULL;
        sop->test = test;
        if (test->objsz < 512)
                gfp = GFP_ATOMIC;

        if (objpool_init(&sop->pool, max, test->objsz,
                         gfp, sop, ot_init_node, NULL)) {
                ot_kfree(test, sop, sizeof(*sop));
                return NULL;
        }
        WARN_ON(max != sop->pool.nr_objs);

        return sop;
}

static void ot_fini_sync(struct ot_context *sop)
{
        objpool_fini(&sop->pool);
        ot_kfree(sop->test, sop, sizeof(*sop));
}

static struct {
        struct ot_context * (*init)(struct ot_test *oc);
        void (*fini)(struct ot_context *sop);
} g_ot_sync_ops[] = {
        {.init = ot_init_sync_m0, .fini = ot_fini_sync},
};

/*
 * synchronous test cases: performance mode
 */

static void ot_bulk_sync(struct ot_item *item, int irq)
{
        struct ot_node *nods[OT_NR_MAX_BULK];
        int i;

        for (i = 0; i < item->bulk[irq]; i++)
                nods[i] = objpool_pop(item->pool);

        if (!irq && (item->delay || !(++(item->niters) & 0x7FFF)))
                msleep(item->delay);

        while (i-- > 0) {
                struct ot_node *on = nods[i];
                if (on) {
                        on->refs++;
                        objpool_push(on, item->pool);
                        item->stat[irq].nhits++;
                } else {
                        item->stat[irq].nmiss++;
                }
        }
}

static int ot_start_sync(struct ot_test *test)
{
        struct ot_context *sop;
        ktime_t start;
        u64 duration;
        unsigned long timeout;
        int cpu;

        /* initialize objpool for syncrhonous testcase */
        sop = g_ot_sync_ops[test->mode].init(test);
        if (!sop)
                return -ENOMEM;

        /* grab rwsem to block testing threads */
        down_write(&test->data.start);

        for_each_possible_cpu(cpu) {
                struct ot_item *item = per_cpu_ptr(&ot_pcup_items, cpu);
                struct task_struct *work;

                ot_init_cpu_item(item, test, &sop->pool, ot_bulk_sync);

                /* skip offline cpus */
                if (!cpu_online(cpu))
                        continue;

                work = kthread_run_on_cpu(ot_thread_worker, item,
                                          cpu, "ot_worker_%d");
                if (IS_ERR(work))
                        pr_err("failed to create thread for cpu %d\n", cpu);
        }

        /* wait a while to make sure all threads waiting at start line */
        msleep(20);

        /* in case no threads were created: memory insufficient ? */
        if (atomic_dec_and_test(&test->data.nthreads))
                complete(&test->data.wait);

        // sched_set_fifo_low(current);

        /* start objpool testing threads */
        start = ktime_get();
        up_write(&test->data.start);

        /* yeild cpu to worker threads for duration ms */
        timeout = msecs_to_jiffies(test->duration);
        schedule_timeout_interruptible(timeout);

        /* tell workers threads to quit */
        atomic_set_release(&test->data.stop, 1);

        /* wait all workers threads finish and quit */
        wait_for_completion(&test->data.wait);
        duration = (u64) ktime_us_delta(ktime_get(), start);

        /* cleanup objpool */
        g_ot_sync_ops[test->mode].fini(sop);

        /* report testing summary and performance results */
        ot_perf_report(test, duration);

        /* report memory allocation summary */
        ot_mem_report(test);

        return 0;
}

/*
 * asynchronous test cases: pool lifecycle controlled by refcount
 */

static void ot_fini_async_rcu(struct rcu_head *rcu)
{
        struct ot_context *sop = container_of(rcu, struct ot_context, rcu);
        struct ot_test *test = sop->test;

        /* here all cpus are aware of the stop event: test->data.stop = 1 */
        WARN_ON(!atomic_read_acquire(&test->data.stop));

        objpool_fini(&sop->pool);
        complete(&test->data.rcu);
}

static void ot_fini_async(struct ot_context *sop)
{
        /* make sure the stop event is acknowledged by all cores */
        call_rcu(&sop->rcu, ot_fini_async_rcu);
}

static int ot_objpool_release(struct objpool_head *head, void *context)
{
        struct ot_context *sop = context;

        WARN_ON(!head || !sop || head != &sop->pool);

        /* do context cleaning if needed */
        if (sop)
                ot_kfree(sop->test, sop, sizeof(*sop));

        return 0;
}

static struct ot_context *ot_init_async_m0(struct ot_test *test)
{
        struct ot_context *sop = NULL;
        int max = num_possible_cpus() << 3;
        gfp_t gfp = GFP_KERNEL;

        sop = (struct ot_context *)ot_kzalloc(test, sizeof(*sop));
        if (!sop)
                return NULL;
        sop->test = test;
        if (test->objsz < 512)
                gfp = GFP_ATOMIC;

        if (objpool_init(&sop->pool, max, test->objsz, gfp, sop,
                         ot_init_node, ot_objpool_release)) {
                ot_kfree(test, sop, sizeof(*sop));
                return NULL;
        }
        WARN_ON(max != sop->pool.nr_objs);

        return sop;
}

static struct {
        struct ot_context * (*init)(struct ot_test *oc);
        void (*fini)(struct ot_context *sop);
} g_ot_async_ops[] = {
        {.init = ot_init_async_m0, .fini = ot_fini_async},
};

static void ot_nod_recycle(struct ot_node *on, struct objpool_head *pool,
                        int release)
{
        struct ot_context *sop;

        on->refs++;

        if (!release) {
                /* push object back to opjpool for reuse */
                objpool_push(on, pool);
                return;
        }

        sop = container_of(pool, struct ot_context, pool);
        WARN_ON(sop != pool->context);

        /* unref objpool with nod removed forever */
        objpool_drop(on, pool);
}

static void ot_bulk_async(struct ot_item *item, int irq)
{
        struct ot_test *test = item->test;
        struct ot_node *nods[OT_NR_MAX_BULK];
        int i, stop;

        for (i = 0; i < item->bulk[irq]; i++)
                nods[i] = objpool_pop(item->pool);

        if (!irq) {
                if (item->delay || !(++(item->niters) & 0x7FFF))
                        msleep(item->delay);
                get_cpu();
        }

        stop = atomic_read_acquire(&test->data.stop);

        /* drop all objects and deref objpool */
        while (i-- > 0) {
                struct ot_node *on = nods[i];

                if (on) {
                        on->refs++;
                        ot_nod_recycle(on, item->pool, stop);
                        item->stat[irq].nhits++;
                } else {
                        item->stat[irq].nmiss++;
                }
        }

        if (!irq)
                put_cpu();
}

static int ot_start_async(struct ot_test *test)
{
        struct ot_context *sop;
        ktime_t start;
        u64 duration;
        unsigned long timeout;
        int cpu;

        /* initialize objpool for syncrhonous testcase */
        sop = g_ot_async_ops[test->mode].init(test);
        if (!sop)
                return -ENOMEM;

        /* grab rwsem to block testing threads */
        down_write(&test->data.start);

        for_each_possible_cpu(cpu) {
                struct ot_item *item = per_cpu_ptr(&ot_pcup_items, cpu);
                struct task_struct *work;

                ot_init_cpu_item(item, test, &sop->pool, ot_bulk_async);

                /* skip offline cpus */
                if (!cpu_online(cpu))
                        continue;

                work = kthread_run_on_cpu(ot_thread_worker, item, cpu, "ot_worker_%d");
                if (IS_ERR(work))
                        pr_err("failed to create thread for cpu %d\n", cpu);
        }

        /* wait a while to make sure all threads waiting at start line */
        msleep(20);

        /* in case no threads were created: memory insufficient ? */
        if (atomic_dec_and_test(&test->data.nthreads))
                complete(&test->data.wait);

        /* start objpool testing threads */
        start = ktime_get();
        up_write(&test->data.start);

        /* yeild cpu to worker threads for duration ms */
        timeout = msecs_to_jiffies(test->duration);
        schedule_timeout_interruptible(timeout);

        /* tell workers threads to quit */
        atomic_set_release(&test->data.stop, 1);

        /* do async-finalization */
        g_ot_async_ops[test->mode].fini(sop);

        /* wait all workers threads finish and quit */
        wait_for_completion(&test->data.wait);
        duration = (u64) ktime_us_delta(ktime_get(), start);

        /* assure rcu callback is triggered */
        wait_for_completion(&test->data.rcu);

        /*
         * now we are sure that objpool is finalized either
         * by rcu callback or by worker threads
         */

        /* report testing summary and performance results */
        ot_perf_report(test, duration);

        /* report memory allocation summary */
        ot_mem_report(test);

        return 0;
}

/*
 * predefined testing cases:
 *   synchronous case / overrun case / async case
 *
 * async: synchronous or asynchronous testing
 * mode: only mode 0 supported
 * objsz: object size
 * duration: int, total test time in ms
 * delay: int, delay (in ms) between each iteration
 * bulk_normal: int, repeat times for thread worker
 * bulk_irq: int, repeat times for irq consumer
 * hrtimer: unsigned long, hrtimer intervnal in ms
 * name: char *, tag for current test ot_item
 */

#define NODE_COMPACT sizeof(struct ot_node)
#define NODE_VMALLOC (512)

static struct ot_test g_testcases[] = {

        /* sync & normal */
        {0, 0, NODE_COMPACT, 1000, 0,  1,  0,  0, "sync: percpu objpool"},
        {0, 0, NODE_VMALLOC, 1000, 0,  1,  0,  0, "sync: percpu objpool from vmalloc"},

        /* sync & hrtimer */
        {0, 0, NODE_COMPACT, 1000, 0,  1,  1,  4, "sync & hrtimer: percpu objpool"},
        {0, 0, NODE_VMALLOC, 1000, 0,  1,  1,  4, "sync & hrtimer: percpu objpool from vmalloc"},

        /* sync & overrun */
        {0, 0, NODE_COMPACT, 1000, 0, 16,  0,  0, "sync overrun: percpu objpool"},
        {0, 0, NODE_VMALLOC, 1000, 0, 16,  0,  0, "sync overrun: percpu objpool from vmalloc"},

        /* async mode */
        {1, 0, NODE_COMPACT, 1000, 100,  1,  0,  0, "async: percpu objpool"},
        {1, 0, NODE_VMALLOC, 1000, 100,  1,  0,  0, "async: percpu objpool from vmalloc"},

        /* async + hrtimer mode */
        {1, 0, NODE_COMPACT, 1000, 0,  4,  4,  4, "async & hrtimer: percpu objpool"},
        {1, 0, NODE_VMALLOC, 1000, 0,  4,  4,  4, "async & hrtimer: percpu objpool from vmalloc"},
};

static int __init ot_mod_init(void)
{
        int i;

        /* perform testings */
        for (i = 0; i < ARRAY_SIZE(g_testcases); i++) {
                ot_init_data(&g_testcases[i].data);
                if (g_testcases[i].async)
                        ot_start_async(&g_testcases[i]);
                else
                        ot_start_sync(&g_testcases[i]);
        }

        /* show tests summary */
        pr_info("\n");
        pr_info("Summary of testcases:\n");
        for (i = 0; i < ARRAY_SIZE(g_testcases); i++) {
                pr_info("    duration: %lluus \thits: %10lu \tmiss: %10lu \t%s\n",
                        g_testcases[i].data.duration, g_testcases[i].data.objects.nhits,
                        g_testcases[i].data.objects.nmiss, g_testcases[i].name);
        }

        return -EAGAIN;
}

static void __exit ot_mod_exit(void)
{
}

module_init(ot_mod_init);
module_exit(ot_mod_exit);

MODULE_DESCRIPTION("Test module for lockless object pool");
MODULE_LICENSE("GPL");