// This is free and unencumbered software released into the public domain.
//
// Anyone is free to copy, modify, publish, use, compile, sell, or
// distribute this software, either in source code form or as a compiled
// binary, for any purpose, commercial or non-commercial, and by any
// means.
//
// In jurisdictions that recognize copyright laws, the author or authors
// of this software dedicate any and all copyright interest in the
// software to the public domain. We make this dedication for the benefit
// of the public at large and to the detriment of our heirs and
// successors. We intend this dedication to be an overt act of
// relinquishment in perpetuity of all present and future rights to this
// software under copyright law.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
// IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
// OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
// ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
// OTHER DEALINGS IN THE SOFTWARE.
//
// For more information, please refer to <http://unlicense.org/>
//
// ------------------------------------------------------------------------
// Taken from: https://github.com/skeeto/lqueue
//
// Atomic, bounded work queue for one writer and multiple readers.
// Threads are spawned internally by the work queue on creation. For
// each job added to the queue with workqueue_push(), the pass function is
// called with the queue-local thread id (0 - n) and the given void
// pointer argument. The thread id can be used to access, for example,
// custom thread-local storage.
//
// When the queue is full, the caller of workqueue_push() (thread 0) will
// assist in completing jobs to open up space in the queue. Similarly,
// in workqueue_wait(), the calling thread will assist in jobs until all
// jobs are consumed *and* completed. Because of this, the queue
// spawns one less thread, because the "main" queuing thread
// participates in job completion.
//
// Calling workqueue_free() will block (a la workqueue_wait()) for all jobs
// to complete before destroying the queue.

#include "config.h"

#include <sys/param.h>
#if HAVE_ERR
# include <err.h>
#endif
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <limits.h>
#include <semaphore.h>
#include <pthread.h>
#include <stdatomic.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "flow.h"
#include "mem.h"
#include "mempool.h"
#include "str.h"
#include "util.h"
#include "workqueue.h"

// C11 lock-free bounded queue. Supports one writer and multiple
// readers. To simplify memory management data offered to the
// queue is copied into the queue's buffers and copied back out
// on retrieval.
//
// Queue functions return false if the queue is full/empty.
struct Lqueue {
	atomic_ulong head;
	atomic_ulong tail;
	unsigned long mask;
	size_t element_size;
	void *buffer;
};

struct WorkqueueData {
	pthread_t thread;
	int id;
	struct Workqueue *q;
	sem_t *pause;
	sem_t *complete;
};

struct Workqueue {
	struct Lqueue *lqueue;
	sem_t *count;
	size_t nthreads;
	struct WorkqueueData *threads;
};

struct WorkqueueJob {
	WorkqueueFn f;
	void *arg;
};

// Prototypes
static struct Lqueue *lqueue_new(size_t, size_t);
static void lqueue_free(struct Lqueue *);
static bool lqueue_offer(struct Lqueue *, void *);
static bool lqueue_poll(struct Lqueue *, void *);
static sem_t *workqueue_sem_new(void);
static void workqueue_sem_free(sem_t *);
static void workqueue_pause(int, void *);
static void *worker(void *);

// Constants
static const WorkqueueFn POISON = (WorkqueueFn)worker;

#ifdef __APPLE__
static int workqueue_sem_id_counter = 0;
#endif

struct Lqueue *
lqueue_new(size_t min_size, size_t element_size)
{
	min_size = next_power_of_2(min_size);
	struct Lqueue *q = xmalloc(sizeof(struct Lqueue));
	q->buffer = xrecallocarray(NULL, 0, min_size, element_size);
	q->mask = min_size - 1;
	q->element_size = element_size;
	q->head = ATOMIC_VAR_INIT(0);
	q->tail = ATOMIC_VAR_INIT(0);
	return q;
}

void
lqueue_free(struct Lqueue *q)
{
	if (q) {
		free(q->buffer);
		free(q);
	}
}

bool
lqueue_offer(struct Lqueue *q, void *v)
{
	unsigned long head = atomic_load(&q->head);
	unsigned long tail = atomic_load(&q->tail);
	if (((tail + 1) & q->mask) == head) {
		return false;
	}
	memcpy(q->buffer + q->element_size * tail, v, q->element_size);
	atomic_store(&q->tail, ((tail + 1) & q->mask));
	return true;
}

bool
lqueue_poll(struct Lqueue *q, void *v)
{
	unsigned long head = atomic_load(&q->head);
	unsigned long tail = atomic_load(&q->tail);
	unsigned long next_head;
	do {
		if (head == tail) {
			return false;
		}
		memcpy(v, q->buffer + q->element_size * head, q->element_size);
		next_head = (head + 1) & q->mask;
	} while (!atomic_compare_exchange_weak(&q->head, &head, next_head));
	return true;
}

void
workqueue_pause(int id, void *arg)
{
	struct Workqueue *q = arg;
	if (sem_post(q->threads[id - 1].complete) < 0) {
		err(1, "sem_post");
	}
	while (sem_wait(q->threads[id - 1].pause) < 0) {
		if (errno == EINTR) {
			continue;
		} else {
			err(1, "sem_wait");
		}
	}
}

void *
worker(void *arg)
{
	struct WorkqueueData *data = arg;
	for (;;) {
		while (sem_wait(data->q->count) < 0) {
			if (errno == EINTR) {
				continue;
			} else {
				err(1, "sem_wait");
			}
		}
		struct WorkqueueJob job;
		if (lqueue_poll(data->q->lqueue, &job)) {
			if (job.f == POISON) {
				break;
			}
			job.f(data->id, job.arg);
		}
	}
	return NULL;
}

sem_t *
workqueue_sem_new()
{
#ifdef __APPLE__
	// macOS does not support unnamed POSIX semaphores, so try to create a named one
	do {
		SCOPE_MEMPOOL(pool);
		char *name = str_printf(pool, "/ias%d_%d", getpid(), workqueue_sem_id_counter++);
		sem_t *sem = sem_open(name, O_CREAT | O_EXCL, 0700, 0);
		if (sem != SEM_FAILED) {
			if (sem_unlink(name) < 0) {
				err(1, "sem_unlink");
			}
			return sem;
		}
	} while (errno == EEXIST || errno == EACCES || errno == EINTR);
	err(1, "sem_open");
#else
	sem_t *sem = xmalloc(sizeof(sem_t));
	if (sem_init(sem, 0, 0) < 0) {
		err(1, "sem_init");
	} else {
		return sem;
	}
#endif
}

void
workqueue_sem_free(sem_t *sem)
{
	unless (sem) {
		return;
	}
#ifdef __APPLE__
	if (sem_close(sem) < 0) {
		err(1, "sem_close");
	}
#else
	sem_destroy(sem);
	free(sem);
#endif
}

struct Workqueue *
workqueue_new(size_t nthreads)
{
	ssize_t ncpu = sysconf(_SC_NPROCESSORS_ONLN);
	if (ncpu == 0) {
		ncpu = 1;
	} else if (ncpu < 0) {
		err(1, "sysconf");
	}
	if (nthreads == 0) {
		nthreads = ncpu;
	}
	nthreads--;
	struct Workqueue *q = xmalloc(sizeof(struct Workqueue));
	q->threads = xrecallocarray(NULL, 0, nthreads, sizeof(q->threads[0]));
	q->lqueue = lqueue_new(next_power_of_2(next_power_of_2(nthreads) + 1) + 1, sizeof(struct WorkqueueJob));
	q->count = workqueue_sem_new();
	q->nthreads = nthreads;
	for (size_t i = 0; i < nthreads; i++) {
		q->threads[i].q = q;
		q->threads[i].id = i + 1;
		q->threads[i].pause = workqueue_sem_new();
		q->threads[i].complete = workqueue_sem_new();
		if (pthread_create(&q->threads[i].thread, NULL, worker, &q->threads[i]) != 0) {
			err(1, "pthread_create");
		}
	}
	return q;
}

void
workqueue_free(struct Workqueue *q)
{
	unless (q) {
		return;
	}

	workqueue_wait(q);
	for (size_t i = 0; i < q->nthreads; i++) {
		workqueue_push(q, POISON, NULL);
	}
	for (size_t i = 0; i < q->nthreads; i++) {
		if (pthread_join(q->threads[i].thread, NULL) != 0) {
			err(1, "pthread_join");
		}
		workqueue_sem_free(q->threads[i].pause);
		workqueue_sem_free(q->threads[i].complete);
	}
	lqueue_free(q->lqueue);
	workqueue_sem_free(q->count);
	free(q);
}

void
workqueue_push(struct Workqueue *q, WorkqueueFn f, void *v)
{
	unless (f) {
		return;
	}
	struct WorkqueueJob job = {f, v};
	while (!lqueue_offer(q->lqueue, &job)) {
		// Help finish jobs until a spot opens in the queue
		struct WorkqueueJob job;
		if (lqueue_poll(q->lqueue, &job)) {
			job.f(0, job.arg);
		}
	}
	if (sem_post(q->count) < 0) {
		err(1, "sem_post");
	}
}

size_t
workqueue_threads(struct Workqueue *q)
{
	// including main thread
	return q->nthreads + 1;
}

void
workqueue_wait(struct Workqueue *q)
{
	// Help finish running jobs
	struct WorkqueueJob job;
	while (lqueue_poll(q->lqueue, &job)) {
		job.f(0, job.arg);
	}

	// Ask all threads to pause
	for (size_t i = 0; i < q->nthreads; i++) {
		workqueue_push(q, workqueue_pause, q);
	}

	// Wait for all threads to complete
	for (size_t i = 0; i < q->nthreads; i++) {
		while (sem_wait(q->threads[i].complete) < 0) {
			if (errno == EINTR) {
				continue;
			} else {
				err(1, "sem_wait");
			}
		}
	}

	// Unpause threads
	for (size_t i = 0; i < q->nthreads; i++) {
		if (sem_post(q->threads[i].pause) < 0) {
			err(1, "sem_post");
		}
	}
}
