/*
 * Wound/Wait Mutex testing module.
 *  Copyright (C) 2016-2018 VMWare Inc.
 */

#include "ww_mutex_test.h"

#ifdef WW_BUILTIN
#include <linux/ww_mutex.h>
#define ww_acquire_batch_begin(_a)
#define ww_acquire_batch_end(_a)
#else
#include "ww_mutex.h"
#endif

#include <linux/init.h>
#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/printk.h>
#include <linux/random.h>
#include <linux/delay.h>

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Thomas Hellström");
MODULE_DESCRIPTION("Test module for ww_mutex performance");
MODULE_VERSION("1.0.0");

struct ww_mutex_item  {
	struct ww_mutex mutex;
	/* Align to a typical cache line */
	u8 pad[64 - sizeof(struct ww_mutex) % 64];
};

struct wwt_device {
	struct ww_mutex_item *locks;
	int *used;
	int num_locks;
};

struct wwt_thread {
	struct wwt_device *dev;
	struct ww_mutex **p_locks;
	u8 *used;
	int num_locks;
	int thread;
	u32 num_rollbacks;
	struct work_struct work;
};




static struct wwt_device g_dev;
static struct wwt_thread threads[WWT_NUM_THREADS];
static DEFINE_WW_CLASS(wwt_class);

static void wwt_thread_rollback(struct wwt_thread *t, int upto)
{
	while(upto--)
		ww_mutex_unlock(t->p_locks[upto]);
	t->num_rollbacks++;
}

static int wwt_thread_lock_sequence(struct wwt_thread *t)
{
	struct ww_acquire_ctx ctx;
	int ret = 0;
	int i;

	ww_acquire_init(&ctx, &wwt_class);
	ww_acquire_batch_begin(&ctx);
restart:
	for (i = 0; i < t->num_locks; ++i) {
		ret = ww_mutex_lock(t->p_locks[i], &ctx);
		if (ret) {
			WARN_ON(i == 0);
			wwt_thread_rollback(t, i);
			if (ret == -EDEADLK) {
				swap(t->p_locks[0], t->p_locks[i]);
				goto restart;
			} else
				goto out_err;
		}
	}
	ww_acquire_batch_end(&ctx);
	ww_acquire_done(&ctx);

	udelay(WWT_CS_UDELAY);

	ww_acquire_batch_begin(&ctx);
	wwt_thread_rollback(t, t->num_locks);
	t->num_rollbacks--;
out_err:
	ww_acquire_batch_end(&ctx);
	ww_acquire_fini(&ctx);
	return ret;
}

static void wwt_lock_work(struct work_struct *work) {
	struct wwt_thread *t = container_of(work, struct wwt_thread, work);
	int i, ret;

	printk(KERN_INFO "thread %d: Starting work.\n", t->thread);
	for (i = 0; i < WWT_NUM_SUB; ++i) {
		ret = wwt_thread_lock_sequence(t);
		if (ret)
			break;
	}

	printk(KERN_INFO "thread %d; Fininshed with %u rollbacks. ret is %d\n",
	       t->thread, t->num_rollbacks, ret);
}


static int wwt_thread_setup(struct wwt_thread *t, int thread)
{
	struct wwt_device *dev = &g_dev;
	u32 idx;
	int i;

	t->dev = dev;
	t->p_locks = vmalloc(WWT_NUM_T_LOCKS * sizeof(*t->p_locks));
	t->num_rollbacks = 0;
	t->thread = thread;
	INIT_WORK(&t->work, wwt_lock_work);
#ifndef WWT_NO_SHARED
	memset(dev->used, 0, dev->num_locks * sizeof(*dev->used));
#endif
	if (!t->p_locks)
		return -ENOMEM;

	for (i = 0; i < WWT_NUM_T_LOCKS; ++i) {
		get_random_bytes(&idx, sizeof(idx));
		idx %= WWT_NUM_LOCKS;
		while(dev->used[idx]) {
			idx++;
			idx %= WWT_NUM_LOCKS;
		}
		dev->used[idx] = 1;
		t->p_locks[i] = &dev->locks[idx].mutex;
	}
	t->num_locks = WWT_NUM_T_LOCKS;

	return 0;
}

static void wwt_thread_destroy(struct wwt_thread *t)
{
	cancel_work_sync(&t->work);
	if (t->p_locks)
		vfree(t->p_locks);
}

static void wwt_device_fini(struct wwt_device *wdev)
{
	int i;

	if (wdev->used)
		vfree(wdev->used);

	for (i = 0; i < ARRAY_SIZE(threads); ++i)
		wwt_thread_destroy(&threads[i]);

	for (i = 0; i < wdev->num_locks; ++i)
		ww_mutex_destroy(&wdev->locks[i].mutex);

	vfree(wdev->locks);
	wdev->num_locks = 0;
}


static int wwt_device_init(struct wwt_device *wdev, int num_locks)
{
	int i, ret;

	wdev->num_locks = 0;
	wdev->locks = vzalloc(num_locks * sizeof(*wdev->locks));
	if (!wdev->locks)
		return -ENOMEM;

	wdev->used = vzalloc(num_locks * sizeof(*wdev->used));
	if (!wdev->locks)
		goto out_thread_setup;

	for (i = 0; i < num_locks; ++i)
		ww_mutex_init(&wdev->locks[i].mutex, &wwt_class);

	wdev->num_locks = num_locks;

	printk(KERN_INFO "Allocated %d mutexes. Alignment is %lu %lu\n",
	       num_locks, ((unsigned long) wdev->locks) % 64,
	       sizeof(struct ww_mutex));

	for (i = 0; i < ARRAY_SIZE(threads); ++i) {
		ret = wwt_thread_setup(&threads[i], i);
		if (ret)
			goto out_thread_setup;
	}

	printk(KERN_INFO "Set up %lu threads\n", ARRAY_SIZE(threads));
	usleep_range(10000, 20000);
	for (i = 0; i < ARRAY_SIZE(threads); ++i) {
#ifndef WWT_SEQUENTIAL
	  queue_work(system_unbound_wq, &threads[i].work);
#else
	  wwt_lock_work(&threads[i].work);
#endif
	}

	return 0;

out_thread_setup:
	wwt_device_fini(&g_dev);
	return ret;
}

static int __init wwt_init(void) {

	printk(KERN_INFO "Hello, World!\n");
	return wwt_device_init(&g_dev, WWT_NUM_LOCKS);
}

static void __exit wwt_exit(void) {
	wwt_device_fini(&g_dev);
	printk(KERN_INFO "Goodbye, World!\n");
}

module_init(wwt_init);
module_exit(wwt_exit);