#define _POSIX_C_SOURCE 200809L

#include <errno.h>
#include <pthread.h>
#include <stdatomic.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#define B19A_WORKERS 64
#define B19A_BODY_SIZE 256
#define B19A_ENTRY_LIMIT 65536
#define B19A_MOUNT_BUDGET 1048576
#define B19A_EINTEGRITY 97

enum cache_state {
	CACHE_EMPTY,
	CACHE_INFLIGHT,
	CACHE_READY,
	CACHE_FAILED,
};

enum cache_lookup {
	CACHE_BYPASS,
	CACHE_HIT,
	CACHE_OWNER,
	CACHE_ERROR,
};

struct mount_budget {
	pthread_mutex_t lock;
	size_t resident;
};

struct cache {
	pthread_mutex_t lock;
	pthread_cond_t cv;
	struct mount_budget *budget;
	unsigned char *body;
	size_t body_size;
	size_t charged;
	unsigned int waiters;
	int error;
	enum cache_state state;
	bool closing;
};

struct wave {
	struct cache *cache;
	pthread_barrier_t barrier;
	unsigned char payload[B19A_BODY_SIZE];
	atomic_uint loads;
	atomic_uint bypasses;
	int expected_error;
	int results[B19A_WORKERS];
};

struct worker_arg {
	struct wave *wave;
	unsigned int index;
};

struct close_arg {
	struct cache *cache;
	atomic_bool done;
};

static void
fail(const char *message)
{
	fprintf(stderr, "B19a cache model failure: %s\n", message);
	exit(1);
}

static void
check_pthread(int error, const char *operation)
{
	if (error != 0) {
		errno = error;
		perror(operation);
		exit(1);
	}
}

static void
budget_init(struct mount_budget *budget)
{
	memset(budget, 0, sizeof(*budget));
	check_pthread(pthread_mutex_init(&budget->lock, NULL),
	    "pthread_mutex_init budget");
}

static bool
budget_reserve(struct mount_budget *budget, size_t amount)
{
	bool reserved;

	check_pthread(pthread_mutex_lock(&budget->lock),
	    "pthread_mutex_lock budget");
	reserved = amount <= B19A_MOUNT_BUDGET - budget->resident;
	if (reserved)
		budget->resident += amount;
	check_pthread(pthread_mutex_unlock(&budget->lock),
	    "pthread_mutex_unlock budget");
	return (reserved);
}

static void
budget_release(struct mount_budget *budget, size_t amount)
{
	check_pthread(pthread_mutex_lock(&budget->lock),
	    "pthread_mutex_lock budget");
	if (amount > budget->resident)
		fail("mount budget underflow");
	budget->resident -= amount;
	check_pthread(pthread_mutex_unlock(&budget->lock),
	    "pthread_mutex_unlock budget");
}

static size_t
budget_resident(struct mount_budget *budget)
{
	size_t resident;

	check_pthread(pthread_mutex_lock(&budget->lock),
	    "pthread_mutex_lock budget");
	resident = budget->resident;
	check_pthread(pthread_mutex_unlock(&budget->lock),
	    "pthread_mutex_unlock budget");
	return (resident);
}

static void
budget_destroy(struct mount_budget *budget)
{
	if (budget_resident(budget) != 0)
		fail("mount budget remained charged at destroy");
	check_pthread(pthread_mutex_destroy(&budget->lock),
	    "pthread_mutex_destroy budget");
}

static void
cache_init(struct cache *cache, struct mount_budget *budget)
{
	memset(cache, 0, sizeof(*cache));
	cache->budget = budget;
	check_pthread(pthread_mutex_init(&cache->lock, NULL),
	    "pthread_mutex_init cache");
	check_pthread(pthread_cond_init(&cache->cv, NULL),
	    "pthread_cond_init cache");
}

static enum cache_lookup
cache_claim(struct cache *cache, size_t body_size, unsigned char *output,
    int *errorp)
{
	enum cache_lookup result;

	*errorp = 0;
	check_pthread(pthread_mutex_lock(&cache->lock),
	    "pthread_mutex_lock cache");
	if (cache->closing) {
		*errorp = ENXIO;
		goto bypass;
	}
	if (cache->state == CACHE_READY) {
		if (cache->body == NULL || cache->body_size != body_size)
			fail("READY cache body is inconsistent");
		memcpy(output, cache->body, body_size);
		check_pthread(pthread_mutex_unlock(&cache->lock),
		    "pthread_mutex_unlock cache");
		return (CACHE_HIT);
	}
	if (cache->state == CACHE_INFLIGHT) {
		++cache->waiters;
		do {
			check_pthread(pthread_cond_wait(&cache->cv, &cache->lock),
			    "pthread_cond_wait cache");
		} while (cache->state == CACHE_INFLIGHT);
		if (cache->state == CACHE_READY) {
			if (cache->body == NULL || cache->body_size != body_size)
				fail("published cache body is inconsistent");
			memcpy(output, cache->body, body_size);
			result = CACHE_HIT;
		} else {
			if (cache->state != CACHE_FAILED || cache->error <= 0)
				fail("waiter observed an untyped failure");
			*errorp = cache->error;
			result = CACHE_ERROR;
		}
		--cache->waiters;
		if (cache->state == CACHE_FAILED && cache->waiters == 0) {
			cache->error = 0;
			cache->state = CACHE_EMPTY;
		}
		if (cache->waiters == 0)
			check_pthread(pthread_cond_broadcast(&cache->cv),
			    "pthread_cond_broadcast cache");
		check_pthread(pthread_mutex_unlock(&cache->lock),
		    "pthread_mutex_unlock cache");
		return (result);
	}
	if (cache->state == CACHE_FAILED) {
		if (cache->waiters != 0)
			goto bypass;
		cache->error = 0;
		cache->state = CACHE_EMPTY;
	}
	if (body_size > B19A_ENTRY_LIMIT ||
	    !budget_reserve(cache->budget, body_size))
		goto bypass;
	cache->body_size = body_size;
	cache->charged = body_size;
	cache->state = CACHE_INFLIGHT;
	check_pthread(pthread_mutex_unlock(&cache->lock),
	    "pthread_mutex_unlock cache");
	return (CACHE_OWNER);

bypass:
	check_pthread(pthread_mutex_unlock(&cache->lock),
	    "pthread_mutex_unlock cache");
	return (CACHE_BYPASS);
}

static void
cache_complete(struct cache *cache, unsigned char *body, size_t body_size,
    int error)
{
	if ((error == 0) != (body != NULL))
		fail("owner completion is not typed");
	check_pthread(pthread_mutex_lock(&cache->lock),
	    "pthread_mutex_lock cache");
	if (cache->state != CACHE_INFLIGHT || cache->body_size != body_size)
		fail("cache owner lost its inflight body");
	if (error == 0) {
		cache->body = body;
		cache->error = 0;
		cache->state = CACHE_READY;
	} else {
		if (error < 1)
			fail("owner published a non-positive errno");
		budget_release(cache->budget, cache->charged);
		cache->charged = 0;
		cache->body_size = 0;
		cache->error = error;
		cache->state = CACHE_FAILED;
		if (cache->waiters == 0) {
			cache->error = 0;
			cache->state = CACHE_EMPTY;
		}
	}
	check_pthread(pthread_cond_broadcast(&cache->cv),
	    "pthread_cond_broadcast cache");
	check_pthread(pthread_mutex_unlock(&cache->lock),
	    "pthread_mutex_unlock cache");
}

static unsigned char *
body_copy(const unsigned char *payload, size_t body_size)
{
	unsigned char *copy;

	copy = malloc(body_size);
	if (copy == NULL)
		fail("body allocation failed");
	memcpy(copy, payload, body_size);
	return (copy);
}

static void
wait_for_waiters(struct cache *cache, unsigned int expected)
{
	struct timespec delay = { .tv_sec = 0, .tv_nsec = 1000000 };
	unsigned int attempt, waiters;

	for (attempt = 0; attempt < 5000; ++attempt) {
		check_pthread(pthread_mutex_lock(&cache->lock),
		    "pthread_mutex_lock cache");
		waiters = cache->waiters;
		check_pthread(pthread_mutex_unlock(&cache->lock),
		    "pthread_mutex_unlock cache");
		if (waiters == expected)
			return;
		nanosleep(&delay, NULL);
	}
	fail("workers did not register as same-vnode waiters");
}

static void *
wave_worker(void *opaque)
{
	struct worker_arg *arg;
	struct wave *wave;
	unsigned char output[B19A_BODY_SIZE];
	enum cache_lookup lookup;
	int barrier_error, error;

	arg = opaque;
	wave = arg->wave;
	barrier_error = pthread_barrier_wait(&wave->barrier);
	if (barrier_error != 0 && barrier_error != PTHREAD_BARRIER_SERIAL_THREAD)
		check_pthread(barrier_error, "pthread_barrier_wait");
	lookup = cache_claim(wave->cache, sizeof(output), output, &error);
	if (lookup == CACHE_OWNER) {
		atomic_fetch_add_explicit(&wave->loads, 1, memory_order_relaxed);
		wait_for_waiters(wave->cache, B19A_WORKERS - 1);
		if (wave->expected_error != 0) {
			cache_complete(wave->cache, NULL, sizeof(output),
			    wave->expected_error);
			wave->results[arg->index] = wave->expected_error;
		} else {
			memcpy(output, wave->payload, sizeof(output));
			cache_complete(wave->cache,
			    body_copy(wave->payload, sizeof(wave->payload)),
			    sizeof(output), 0);
			wave->results[arg->index] =
			    memcmp(output, wave->payload, sizeof(output)) == 0 ? 0 : EIO;
		}
	} else if (lookup == CACHE_HIT) {
		wave->results[arg->index] =
		    memcmp(output, wave->payload, sizeof(output)) == 0 ? 0 : EIO;
	} else if (lookup == CACHE_ERROR) {
		wave->results[arg->index] = error;
	} else {
		atomic_fetch_add_explicit(&wave->bypasses, 1, memory_order_relaxed);
		wave->results[arg->index] = error != 0 ? error : EBUSY;
	}
	return (NULL);
}

static void
run_wave(struct cache *cache, int expected_error)
{
	struct wave wave;
	struct worker_arg args[B19A_WORKERS];
	pthread_t threads[B19A_WORKERS];
	unsigned int index;

	memset(&wave, 0, sizeof(wave));
	wave.cache = cache;
	wave.expected_error = expected_error;
	for (index = 0; index < sizeof(wave.payload); ++index)
		wave.payload[index] = (unsigned char)(index ^ 0x5a);
	check_pthread(pthread_barrier_init(&wave.barrier, NULL, B19A_WORKERS),
	    "pthread_barrier_init");
	for (index = 0; index < B19A_WORKERS; ++index) {
		args[index].wave = &wave;
		args[index].index = index;
		check_pthread(pthread_create(&threads[index], NULL, wave_worker,
		    &args[index]), "pthread_create");
	}
	for (index = 0; index < B19A_WORKERS; ++index)
		check_pthread(pthread_join(threads[index], NULL), "pthread_join");
	check_pthread(pthread_barrier_destroy(&wave.barrier),
	    "pthread_barrier_destroy");
	if (atomic_load_explicit(&wave.loads, memory_order_relaxed) != 1)
		fail("same-vnode wave did not have exactly one owner");
	if (atomic_load_explicit(&wave.bypasses, memory_order_relaxed) != 0)
		fail("same-vnode wave bypassed the cache");
	for (index = 0; index < B19A_WORKERS; ++index) {
		if (wave.results[index] != expected_error)
			fail("same-vnode waiter observed a different result");
	}
}

static void
cache_invalidate(struct cache *cache)
{
	unsigned char *body;
	size_t charged;

	check_pthread(pthread_mutex_lock(&cache->lock),
	    "pthread_mutex_lock cache");
	while (cache->state == CACHE_INFLIGHT || cache->waiters != 0)
		check_pthread(pthread_cond_wait(&cache->cv, &cache->lock),
		    "pthread_cond_wait cache");
	body = cache->body;
	charged = cache->charged;
	cache->body = NULL;
	cache->body_size = 0;
	cache->charged = 0;
	cache->error = 0;
	cache->state = CACHE_EMPTY;
	check_pthread(pthread_mutex_unlock(&cache->lock),
	    "pthread_mutex_unlock cache");
	free(body);
	if (charged != 0)
		budget_release(cache->budget, charged);
}

static void
cache_close(struct cache *cache)
{
	check_pthread(pthread_mutex_lock(&cache->lock),
	    "pthread_mutex_lock cache");
	cache->closing = true;
	check_pthread(pthread_mutex_unlock(&cache->lock),
	    "pthread_mutex_unlock cache");
	cache_invalidate(cache);
	check_pthread(pthread_cond_destroy(&cache->cv),
	    "pthread_cond_destroy cache");
	check_pthread(pthread_mutex_destroy(&cache->lock),
	    "pthread_mutex_destroy cache");
}

static void *
close_worker(void *opaque)
{
	struct close_arg *arg;

	arg = opaque;
	cache_close(arg->cache);
	atomic_store_explicit(&arg->done, true, memory_order_release);
	return (NULL);
}

static void
test_hits_and_failure(struct mount_budget *budget)
{
	struct cache cache;
	unsigned char output[B19A_BODY_SIZE];
	enum cache_lookup lookup;
	unsigned int index;
	int error;

	cache_init(&cache, budget);
	run_wave(&cache, 0);
	if (budget_resident(budget) != B19A_BODY_SIZE)
		fail("success wave did not charge one body");
	for (index = 0; index < 100; ++index) {
		lookup = cache_claim(&cache, sizeof(output), output, &error);
		if (lookup != CACHE_HIT || error != 0)
			fail("ready cache did not hit");
	}
	cache_invalidate(&cache);
	if (budget_resident(budget) != 0 || cache.state != CACHE_EMPTY)
		fail("cache invalidation did not release its body");
	run_wave(&cache, B19A_EINTEGRITY);
	if (budget_resident(budget) != 0 || cache.state != CACHE_EMPTY)
		fail("failed publication leaked body state");
	run_wave(&cache, 0);
	cache_close(&cache);
	if (budget_resident(budget) != 0)
		fail("cache close did not release the retry body");
}

static void
test_limits(struct mount_budget *budget)
{
	struct cache caches[17];
	unsigned char *body, output[B19A_ENTRY_LIMIT + 1];
	enum cache_lookup lookup;
	unsigned int index;
	int error;

	cache_init(&caches[0], budget);
	lookup = cache_claim(&caches[0], B19A_ENTRY_LIMIT + 1, output, &error);
	if (lookup != CACHE_BYPASS || budget_resident(budget) != 0)
		fail("oversized body did not bypass without charge");
	cache_close(&caches[0]);

	for (index = 0; index < 17; ++index)
		cache_init(&caches[index], budget);
	for (index = 0; index < 16; ++index) {
		lookup = cache_claim(&caches[index], B19A_ENTRY_LIMIT, output, &error);
		if (lookup != CACHE_OWNER)
			fail("mount budget rejected an in-budget body");
		body = calloc(1, B19A_ENTRY_LIMIT);
		if (body == NULL)
			fail("budget body allocation failed");
		cache_complete(&caches[index], body, B19A_ENTRY_LIMIT, 0);
	}
	if (budget_resident(budget) != B19A_MOUNT_BUDGET)
		fail("mount budget did not reach its exact hard limit");
	lookup = cache_claim(&caches[16], B19A_ENTRY_LIMIT, output, &error);
	if (lookup != CACHE_BYPASS || budget_resident(budget) != B19A_MOUNT_BUDGET)
		fail("mount budget exhaustion did not bypass");
	for (index = 0; index < 17; ++index)
		cache_close(&caches[index]);
	if (budget_resident(budget) != 0)
		fail("mount budget did not drain after reclaim");
}

static void
test_close_inflight(struct mount_budget *budget)
{
	struct cache cache;
	struct close_arg close_arg;
	pthread_t closer;
	unsigned char output[B19A_BODY_SIZE];
	enum cache_lookup lookup;
	struct timespec delay = { .tv_sec = 0, .tv_nsec = 10000000 };
	int error;

	cache_init(&cache, budget);
	lookup = cache_claim(&cache, sizeof(output), output, &error);
	if (lookup != CACHE_OWNER)
		fail("close test did not establish an inflight owner");
	memset(&close_arg, 0, sizeof(close_arg));
	close_arg.cache = &cache;
	check_pthread(pthread_create(&closer, NULL, close_worker, &close_arg),
	    "pthread_create close");
	nanosleep(&delay, NULL);
	if (atomic_load_explicit(&close_arg.done, memory_order_acquire))
		fail("cache close did not wait for inflight owner");
	cache_complete(&cache, body_copy(output, sizeof(output)), sizeof(output), 0);
	check_pthread(pthread_join(closer, NULL), "pthread_join close");
	if (!atomic_load_explicit(&close_arg.done, memory_order_acquire) ||
	    budget_resident(budget) != 0)
		fail("cache close did not drain inflight completion");
}

int
main(void)
{
	struct mount_budget budget;

	budget_init(&budget);
	test_hits_and_failure(&budget);
	test_limits(&budget);
	test_close_inflight(&budget);
	budget_destroy(&budget);
	printf("{\"body_limit\":%u,\"failure_publication\":\"PASS\","
	    "\"inflight_close\":\"PASS\",\"mount_budget\":%u,"
	    "\"owner_waiter\":\"PASS\",\"reclaim\":\"PASS\","
	    "\"resident_after\":0,\"status\":\"PASS\"}\n",
	    B19A_ENTRY_LIMIT, B19A_MOUNT_BUDGET);
	return (0);
}