/*-
 * Copyright (c) 2014 Andrey Zonov <zont@FreeBSD.org>
 * All rights reserved.
 */

#include <sys/param.h>
#include <sys/systm.h>

#include <sys/kernel.h>
#include <sys/kthread.h>
#include <sys/lock.h>
#include <sys/malloc.h>
#include <sys/module.h>
#include <sys/mutex.h>
#include <sys/proc.h>
#include <sys/rmlock.h>
#include <sys/rwlock.h>
#include <sys/sched.h>
#include <sys/types.h>

/*
 * Configuration.
 */
//#define	MTX
//#define	RWL
//#define	RML
#define	PSZ
#define	LOOPS	(10 * hz)
#define	LENGTH	10

#if 0
#define DBG(format, args...) printf(format , ##args)
#else
#define DBG(format, args...)
#endif

/*
 * List implementation.
 */
struct node {
	struct node *next;
	struct node **prev;
};

struct list {
	struct node *head;
};

#define	list_first(list)	((list)->head)
#define	list_next(elm)		((elm)->next)
#define	list_init(list)		do { list_first((list)) = NULL; } while (0)
#define	list_foreach(node, list)	for ((node) = list_first((list)); (node); (node) = list_next((node)))
#define	list_entry(ptr, name, field)	__containerof((ptr), struct name, field)

static __inline void
list_insert(struct list *list, struct node *node)
{
	struct node *head;

	head = list->head;
	node->next = head;
	node->prev = &list->head;
	atomic_store_rel_ptr((volatile uintptr_t *)&list->head, (uintptr_t)node);
	if (head != NULL)
		head->prev = &node->next;
}

static __inline void
list_remove(struct node *node)
{
	struct node *next;

	next = node->next;
	*node->prev = next;
	if (next != NULL)
		next->prev = node->prev;
}

#ifdef PSZ
#include <sys/psz.h>
#endif

/*
 * Hash table implementation.
 */
struct hnode {
	struct node h_link;
	long h_data;
#ifdef PSZ
	struct psz_cb h_psz_cb;
#endif
};

/* XXX instead of per slot lock lock_test uses one slot and one global lock */
struct hslot {
	/* lock */
	long length;
	struct list hnodes;
};

struct hash {
	long nslots;		/* unused */
	long hmask;		/* unused */
	struct hslot hslots[0];
};

static struct hash *
hash_init(long nslots)
{
	struct hash *hash;
	struct hslot *hslot;
	long i;

	nslots = 1 << flsl(nslots - 1); /* round up to the nearest power of 2 */
	hash = malloc(sizeof(*hash) + nslots * sizeof(struct hslot), M_TEMP, M_WAITOK);
	hash->nslots = nslots;
	hash->hmask = nslots - 1;
	for (i = 0; i < nslots; i++) {
		hslot = &hash->hslots[i];
		hslot->length = 0;
		list_init(&hslot->hnodes);
	}

	return (hash);
}

static void
hash_destroy(struct hash *hash)
{

	free(hash, M_TEMP);
}

/*
 * Variables.
 */
static int lock_test_done;
static long consumed, produced;
static struct hash *lock_test_hash;

#ifdef MTX
static struct mtx test_lock;
#define	TEST_LOCK_INIT()	mtx_init(&test_lock, "mutex", NULL, MTX_DEF)
#define	TEST_LOCK_DESTROY()	mtx_destroy(&test_lock)
#define	TEST_RLOCK(x)		mtx_lock(&test_lock)
#define	TEST_RUNLOCK(x)		mtx_unlock(&test_lock)
#define	TEST_WLOCK()		mtx_lock(&test_lock)
#define	TEST_WUNLOCK()		mtx_unlock(&test_lock)
#define	TEST_LOCK_TRACKER(x)
#endif

#ifdef RWL
static struct rwlock test_lock;
#define	TEST_LOCK_INIT()	rw_init(&test_lock, "rwlock")
#define	TEST_LOCK_DESTROY()	rw_destroy(&test_lock)
#define	TEST_RLOCK(x)		rw_rlock(&test_lock)
#define	TEST_RUNLOCK(x)		rw_runlock(&test_lock)
#define	TEST_WLOCK()		rw_wlock(&test_lock)
#define	TEST_WUNLOCK()		rw_wunlock(&test_lock)
#define	TEST_LOCK_TRACKER(x)
#endif

#ifdef RML
static struct rmlock test_lock;
#define	TEST_LOCK_INIT()	rm_init(&test_lock, "rmlock")
#define	TEST_LOCK_DESTROY()	rm_destroy(&test_lock)
#define	TEST_RLOCK(x)		rm_rlock(&test_lock, x)
#define	TEST_RUNLOCK(x)		rm_runlock(&test_lock, x)
#define	TEST_WLOCK()		rm_wlock(&test_lock)
#define	TEST_WUNLOCK()		rm_wunlock(&test_lock)
#define	TEST_LOCK_TRACKER(x)	struct rm_priotracker x
#endif

#ifdef PSZ
static struct mtx test_lock;
#define	TEST_LOCK_INIT()	mtx_init(&test_lock, "pszlock", NULL, MTX_DEF)
#define	TEST_LOCK_DESTROY()	mtx_destroy(&test_lock)
#define	TEST_RLOCK(x)		psz_read_lock()
#define	TEST_RUNLOCK(x)		psz_read_unlock()
#define	TEST_WLOCK()		mtx_lock(&test_lock)
#define	TEST_WUNLOCK()		mtx_unlock(&test_lock)
#define	TEST_LOCK_TRACKER(x)

static void
hnode_free_cb(struct psz_cb *cb)
{
	struct hnode *hnode = psz2struct(cb, struct hnode, h_psz_cb);

	free(hnode, M_TEMP);
	atomic_add_long(&consumed, 1);
}
#endif

static void
lock_test_producer(void *arg)
{
	struct hash *hash = arg;
	struct hslot *hslot = &hash->hslots[0];
	struct hnode *hnode;
	int i, loop;

	for (loop = 0; loop < LOOPS; loop++) {
		for (i = 0; i < LENGTH; i++) {
			hnode = malloc(sizeof(*hnode), M_TEMP, M_WAITOK);
			hnode->h_data = i;
			TEST_WLOCK();
			list_insert(&hslot->hnodes, &hnode->h_link);
			hslot->length++;
			TEST_WUNLOCK();
			DBG("insert: data=%ld\n", hnode->h_data);
		}
		tsleep(curthread, 0, "-", 1);
		DBG("%s: %d\n", __func__, loop);
	}

	kthread_exit();
}

static void
lock_test_consumer(void *arg)
{
	struct hash *hash = arg;
	struct hslot *hslot = &hash->hslots[0];
	struct hnode *hnode;
	struct node *node;

	for ( ;; ) {
		TEST_WLOCK();
		if (hslot->length > 2) {
			long i = 0;
			list_foreach(node, &hslot->hnodes) {
				if (i++ > hslot->length / 2)
					break;
			}
		} else
			node = list_first(&hslot->hnodes);
		if (node != NULL) {
			list_remove(node);
			hslot->length--;
			TEST_WUNLOCK();
			hnode = list_entry(node, hnode, h_link);
			DBG("remove: data=%ld\n", hnode->h_data);
#ifdef PSZ
			psz_enqueue(hnode_free_cb, &hnode->h_psz_cb);
#else
			free(hnode, M_TEMP);
			consumed++;
#endif
			continue;
		}
		TEST_WUNLOCK();
		if (consumed == produced)
			break;
		tsleep(curthread, 0, "-", 1);
	}

	atomic_add_int(&lock_test_done, 1);

	kthread_exit();
}

static void
lock_test_reader(void *arg)
{
	struct hash *hash = arg;
	struct hslot *hslot = &hash->hslots[0];
	TEST_LOCK_TRACKER(t);
	struct hnode *hnode;
	struct node *node;
	long loops = 0, num = 0, sum = 0;

	while (!lock_test_done) {
		TEST_RLOCK(&t);
		list_foreach(node, &hslot->hnodes) {
			hnode = list_entry(node, hnode, h_link);
			sum += hnode->h_data;
			num++;
		}
		TEST_RUNLOCK(&t);
		/* yield sometimes */
		if ((loops % (1<<10)) == 0)
			sched_relinquish(curthread);
		loops++;
	}
	printf("read %ld elements (sum=%ld)\n", num, sum);

	atomic_add_int(&lock_test_done, 1);

	kthread_exit();
}

static int
lock_test_load(void)
{
	int error;

	lock_test_done = 0;
	consumed = 0;
	produced = LOOPS * LENGTH;
	lock_test_hash = hash_init(1);
	TEST_LOCK_INIT();

	error = kthread_add(lock_test_producer, lock_test_hash, NULL, NULL, 0, 0, "lock_test: producer");
	if (error != 0)
		return (error);
	error = kthread_add(lock_test_consumer, lock_test_hash, NULL, NULL, 0, 0, "lock_test: consumer");
	if (error != 0)
		return (error);
	error = kthread_add(lock_test_reader, lock_test_hash, NULL, NULL, 0, 0, "lock_test: reader");

	return (error);
}

static int
lock_test_unload(void)
{
	int error;

	while (lock_test_done != 2) {
		error = tsleep(curthread, PCATCH, "-", hz/10);
		if (error != EWOULDBLOCK)
			return (error);
	}

	TEST_LOCK_DESTROY();
	hash_destroy(lock_test_hash);

	return (0);
}

static int
lock_test_handler(module_t mod, int what, void *arg)
{
	int error = 0;

	switch (what) {
	case MOD_LOAD:
		error = lock_test_load();
		break;
	case MOD_UNLOAD:
		error = lock_test_unload();
		break;
	default:
		error = EOPNOTSUPP;
		break;
	}

	return (error);
}

static moduledata_t lock_test_data = {
	"lock_test",
	lock_test_handler,
	NULL
};

DECLARE_MODULE(lock_test, lock_test_data, SI_SUB_EXEC, SI_ORDER_ANY);