#include <linux/kernel.h>
#include <linux/sched.h>
#include <linux/kthread.h>
#include <linux/hardirq.h>
#include <linux/module.h>
#include <linux/slab.h>
#include <linux/ktime.h>

MODULE_LICENSE("GPL");

static int nrthreads = 128;
module_param(nrthreads, int, 0644);

static int loopcount = 1024;
module_param(loopcount, int, 0644);

static int usehrtime = 0;
module_param(usehrtime, int, 0644);

static int slack = 50000;
module_param(slack, int, 0644);

static int msecs = 1;
module_param(msecs, int, 0644);

static DECLARE_COMPLETION(done);
static struct task_struct **threads;
static atomic_t nrunning;

static int timeoutbench_test(void *unused)
{
	int i;
	ktime_t expires = ktime_set(0, msecs * NSEC_PER_MSEC);

	atomic_inc(&nrunning);

	for (i = 0; !kthread_should_stop() && i < loopcount; i++) {
		if (usehrtime) {
			set_current_state(TASK_UNINTERRUPTIBLE);
			schedule_hrtimeout_range(&expires, slack, HRTIMER_MODE_REL);
		}
		else
			schedule_timeout_uninterruptible(msecs_to_jiffies(msecs));
	}

	if (atomic_dec_and_test(&nrunning))
		complete(&done);
	return 0;
}

static int __init timeoutbench_init(void)
{
	int i;

	atomic_set(&nrunning, 0);

	threads = kmalloc(nrthreads * sizeof(struct task_struct *), GFP_KERNEL);
	if (!threads)
		return -ENOMEM;

	for (i = 0; i < nrthreads; i++) {
		threads[i] = kthread_create(timeoutbench_test, NULL,
					    "timeoutbench_test/%d", i);
		if (IS_ERR(threads[i])) {
			int j, err = PTR_ERR(threads[i]);

			for (j = 0; j < i; j++)
				kthread_stop(threads[j]);
			kfree(threads);
			return err;
		}
		get_task_struct(threads[i]);
		wake_up_process(threads[i]);
	}
	return 0;
}

static void __exit timeoutbench_exit(void)
{
	int i;

	wait_for_completion(&done);
	for (i = 0; i < nrthreads; i++) {
		kthread_stop(threads[i]);
		put_task_struct(threads[i]);
	}
	kfree(threads);
}

module_init(timeoutbench_init);
module_exit(timeoutbench_exit);