#include <linux/module.h>
#include <linux/moduleparam.h>
#include <asm/atomic.h>
#include <linux/percpu-refcount.h>
#include <linux/workqueue.h>
#include <linux/completion.h>

static struct workqueue_struct *my_wq;

static int test_all_cpus = 0;
module_param_named(all_cpus, test_all_cpus, int, 0444);

static int test_duration = 15;
module_param_named(duration, test_duration, int, 0444);

static int test_copy_bytes = 32;
module_param_named(copy_bytes, test_copy_bytes, int, 0444);

static char test_src_buf[1024];
static DEFINE_PER_CPU(char [1024], test_dst_buf);
static DEFINE_PER_CPU(struct work_struct, test_work);

static atomic_t test_atomic_ref = ATOMIC_INIT(1);

static void test_atomic_workfn(struct work_struct *work)
{
	unsigned long end = jiffies + test_duration * HZ;
	unsigned long cnt = 0;

	do {
		int bytes = test_copy_bytes;

		atomic_inc(&test_atomic_ref);
		while (bytes) {
			int todo = min_t(int, bytes, 1024);

			memcpy(*this_cpu_ptr(&test_dst_buf), test_src_buf, todo);
			bytes -= todo;
		}
		cnt++;
		atomic_dec(&test_atomic_ref);
	} while (time_before(jiffies, end));

	printk("XXX atomic on CPU %d completed %lu loops\n",
	       smp_processor_id(), cnt);
}

static struct percpu_ref test_pcpu_ref;
static DECLARE_COMPLETION(test_cpu_ref_done);

static void test_pcpu_release(struct percpu_ref *ref)
{
	complete(&test_cpu_ref_done);
}

static void test_pcpu_workfn(struct work_struct *work)
{
	unsigned long end = jiffies + test_duration * HZ;
	unsigned long cnt = 0;

	do {
		int bytes = test_copy_bytes;

		percpu_ref_get(&test_pcpu_ref);
		while (bytes) {
			int todo = min_t(int, bytes, 1024);

			memcpy(*this_cpu_ptr(&test_dst_buf), test_src_buf, todo);
			bytes -= todo;
		}
		cnt++;
		percpu_ref_put(&test_pcpu_ref);
	} while (time_before(jiffies, end));

	printk("XXX percpu on CPU %d completed %lu loops\n",
	       smp_processor_id(), cnt);
}

static int test_pcpuref_init(void)
{
	int cpu;

	if (percpu_ref_init(&test_pcpu_ref, test_pcpu_release))
		return -ENOMEM;

	my_wq = alloc_workqueue("test-pcpuref", WQ_CPU_INTENSIVE, 0);
	if (!my_wq) {
		percpu_ref_cancel_init(&test_pcpu_ref);
		return -ENOMEM;
	}

	printk("XXX testing atomic_t for %d secs\n", test_duration);

	for_each_online_cpu(cpu) {
		INIT_WORK(per_cpu_ptr(&test_work, cpu), test_atomic_workfn);
		queue_work_on(cpu, my_wq, per_cpu_ptr(&test_work, cpu));
		if (!test_all_cpus)
			break;
	}

	flush_workqueue(my_wq);

	printk("XXX testing percpu_ref for %d secs\n", test_duration);

	for_each_online_cpu(cpu) {
		INIT_WORK(per_cpu_ptr(&test_work, cpu), test_pcpu_workfn);
		queue_work_on(cpu, my_wq, per_cpu_ptr(&test_work, cpu));
		if (!test_all_cpus)
			break;
	}

	flush_workqueue(my_wq);

	percpu_ref_kill(&test_pcpu_ref);
	wait_for_completion(&test_cpu_ref_done);

	return 0;
}
module_init(test_pcpuref_init);

static void test_pcpuref_exit(void)
{
	destroy_workqueue(my_wq);
}
module_exit(test_pcpuref_exit);

MODULE_LICENSE("GPL v2");