#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 B32_WORKERS 64
#define B32_DATA_SIZE 32

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

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

struct map_blocks {
	uint64_t m_pa;
	uint64_t m_la;
	uint64_t m_plen;
	uint64_t m_llen;
	uint16_t m_deviceid;
	uint8_t m_algorithmformat;
	unsigned int m_flags;
};

struct cache_key {
	struct map_blocks map;
	uint64_t nid;
	size_t decoded_size;
};

struct cache {
	pthread_mutex_t lock;
	pthread_cond_t cv;
	struct cache_key key;
	unsigned char *data;
	unsigned int waiters;
	int error;
	enum cache_state state;
	bool closing;
};

struct wave {
	struct cache *cache;
	struct cache_key key;
	pthread_barrier_t barrier;
	unsigned char payload[B32_DATA_SIZE];
	atomic_uint decodes;
	atomic_uint bypasses;
	int expected_error;
	int results[B32_WORKERS];
};

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

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

static void
fail(const char *message)
{
	fprintf(stderr, "B32 oracle 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 bool
key_equal(const struct cache_key *left, const struct cache_key *right)
{
	return (left->nid == right->nid &&
	    left->decoded_size == right->decoded_size &&
	    left->map.m_pa == right->map.m_pa &&
	    left->map.m_la == right->map.m_la &&
	    left->map.m_plen == right->map.m_plen &&
	    left->map.m_llen == right->map.m_llen &&
	    left->map.m_deviceid == right->map.m_deviceid &&
	    left->map.m_algorithmformat == right->map.m_algorithmformat &&
	    left->map.m_flags == right->map.m_flags);
}

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

static enum cache_lookup
cache_claim(struct cache *cache, const struct cache_key *key,
    unsigned char *output, int *errorp)
{
	unsigned char *old;
	enum cache_lookup result;

	*errorp = 0;
	old = NULL;
	check_pthread(pthread_mutex_lock(&cache->lock), "pthread_mutex_lock");
	if (cache->closing)
		goto bypass;
	if (cache->state != CACHE_EMPTY && key_equal(&cache->key, key)) {
		switch (cache->state) {
		case CACHE_READY:
			if (cache->data == NULL)
				fail("READY state has no data");
			memcpy(output, cache->data, key->decoded_size);
			check_pthread(pthread_mutex_unlock(&cache->lock),
			    "pthread_mutex_unlock");
			return (CACHE_HIT);
		case CACHE_INFLIGHT:
			++cache->waiters;
			do {
				check_pthread(pthread_cond_wait(&cache->cv, &cache->lock),
				    "pthread_cond_wait");
			} while (cache->state == CACHE_INFLIGHT);
			if (!key_equal(&cache->key, key))
				fail("inflight key changed before waiter consumed it");
			if (cache->state == CACHE_READY) {
				if (cache->data == NULL)
					fail("published success has no data");
				memcpy(output, cache->data, key->decoded_size);
				result = CACHE_HIT;
			} else {
				if (cache->state != CACHE_FAILED || cache->error <= 0)
					fail("published failure has no positive errno");
				*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");
			check_pthread(pthread_mutex_unlock(&cache->lock),
			    "pthread_mutex_unlock");
			return (result);
		case CACHE_FAILED:
			if (cache->waiters != 0)
				goto bypass;
			cache->error = 0;
			cache->state = CACHE_EMPTY;
			break;
		case CACHE_EMPTY:
			break;
		}
	}
	if (cache->state == CACHE_INFLIGHT || cache->waiters != 0)
		goto bypass;
	old = cache->data;
	cache->data = NULL;
	cache->key = *key;
	cache->error = 0;
	cache->state = CACHE_INFLIGHT;
	check_pthread(pthread_mutex_unlock(&cache->lock), "pthread_mutex_unlock");
	free(old);
	return (CACHE_OWNER);

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

static void
cache_complete(struct cache *cache, const struct cache_key *key,
    unsigned char *data, int error)
{
	if ((error == 0) != (data != NULL))
		fail("owner completion is not typed");
	check_pthread(pthread_mutex_lock(&cache->lock), "pthread_mutex_lock");
	if (cache->state != CACHE_INFLIGHT || !key_equal(&cache->key, key))
		fail("owner lost its inflight key");
	if (error == 0) {
		cache->data = data;
		cache->error = 0;
		cache->state = CACHE_READY;
	} else {
		if (error < 1)
			fail("owner published a non-positive errno");
		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");
	check_pthread(pthread_mutex_unlock(&cache->lock), "pthread_mutex_unlock");
}

static unsigned char *
copy_payload(const unsigned char *payload, size_t size)
{
	unsigned char *copy;

	copy = malloc(size);
	if (copy == NULL)
		fail("malloc failed");
	memcpy(copy, payload, 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");
		waiters = cache->waiters;
		check_pthread(pthread_mutex_unlock(&cache->lock),
		    "pthread_mutex_unlock");
		if (waiters == expected)
			return;
		nanosleep(&delay, NULL);
	}
	fail("workers did not register as same-key waiters");
}

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

	arg = opaque;
	wave = arg->wave;
	barrier_result = pthread_barrier_wait(&wave->barrier);
	if (barrier_result != 0 && barrier_result != PTHREAD_BARRIER_SERIAL_THREAD)
		check_pthread(barrier_result, "pthread_barrier_wait");
	lookup = cache_claim(wave->cache, &wave->key, output, &error);
	if (lookup == CACHE_OWNER) {
		atomic_fetch_add_explicit(&wave->decodes, 1, memory_order_relaxed);
		wait_for_waiters(wave->cache, B32_WORKERS - 1);
		if (wave->expected_error != 0) {
			cache_complete(wave->cache, &wave->key, NULL,
			    wave->expected_error);
			wave->results[arg->index] = wave->expected_error;
		} else {
			memcpy(output, wave->payload, sizeof(output));
			cache_complete(wave->cache, &wave->key,
			    copy_payload(wave->payload, sizeof(wave->payload)), 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] = EBUSY;
	}
	return (NULL);
}

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

	memset(&wave, 0, sizeof(wave));
	wave.cache = cache;
	wave.key = *key;
	wave.expected_error = expected_error;
	for (index = 0; index < sizeof(wave.payload); ++index)
		wave.payload[index] = (unsigned char)(0xa0U + index);
	check_pthread(pthread_barrier_init(&wave.barrier, NULL, B32_WORKERS),
	    "pthread_barrier_init");
	for (index = 0; index < B32_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 < B32_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.decodes, memory_order_relaxed) != 1)
		fail("same-key wave had more than one owner decode");
	if (atomic_load_explicit(&wave.bypasses, memory_order_relaxed) != 0)
		fail("same-key wave bypassed its owner generation");
	for (index = 0; index < B32_WORKERS; ++index) {
		if (wave.results[index] != expected_error)
			fail("same-key waiter received different bytes or errno");
	}
}

static struct cache_key
base_key(void)
{
	return ((struct cache_key) {
	    .map = {
		.m_pa = 0x1000,
		.m_la = 0x2000,
		.m_plen = 4096,
		.m_llen = B32_DATA_SIZE,
		.m_deviceid = 3,
		.m_algorithmformat = 1,
		.m_flags = 1,
	    },
	    .nid = 42,
	    .decoded_size = B32_DATA_SIZE,
	});
}

static void
seed_ready(struct cache *cache, const struct cache_key *key)
{
	unsigned char output[B32_DATA_SIZE];
	unsigned char payload[B32_DATA_SIZE];
	enum cache_lookup lookup;
	int error;

	memset(payload, 0x5a, sizeof(payload));
	lookup = cache_claim(cache, key, output, &error);
	if (lookup != CACHE_OWNER)
		fail("seed miss did not become owner");
	cache_complete(cache, key, copy_payload(payload, sizeof(payload)), 0);
}

static void
expect_distinct_key(struct cache *cache, const struct cache_key *base,
    const struct cache_key *changed)
{
	unsigned char output[B32_DATA_SIZE];
	unsigned char payload[B32_DATA_SIZE];
	enum cache_lookup lookup;
	int error;

	lookup = cache_claim(cache, changed, output, &error);
	if (lookup != CACHE_OWNER)
		fail("changed key incorrectly hit or bypassed");
	memset(payload, 0x33, sizeof(payload));
	cache_complete(cache, changed, copy_payload(payload, sizeof(payload)), 0);
	lookup = cache_claim(cache, base, output, &error);
	if (lookup != CACHE_OWNER)
		fail("key reuse did not start a fresh owner generation");
	memset(payload, 0x5a, sizeof(payload));
	cache_complete(cache, base, copy_payload(payload, sizeof(payload)), 0);
}

static void
test_exact_key(struct cache *cache, const struct cache_key *base)
{
	struct cache_key changed;

	seed_ready(cache, base);
	changed = *base;
	++changed.nid;
	expect_distinct_key(cache, base, &changed);
	changed = *base;
	--changed.decoded_size;
	expect_distinct_key(cache, base, &changed);
	changed = *base;
	++changed.map.m_pa;
	expect_distinct_key(cache, base, &changed);
	changed = *base;
	++changed.map.m_la;
	expect_distinct_key(cache, base, &changed);
	changed = *base;
	++changed.map.m_plen;
	expect_distinct_key(cache, base, &changed);
	changed = *base;
	++changed.map.m_llen;
	expect_distinct_key(cache, base, &changed);
	changed = *base;
	++changed.map.m_deviceid;
	expect_distinct_key(cache, base, &changed);
	changed = *base;
	++changed.map.m_algorithmformat;
	expect_distinct_key(cache, base, &changed);
	changed = *base;
	changed.map.m_flags ^= 2U;
	expect_distinct_key(cache, base, &changed);
}

static void
test_fallback_and_waiter_generation(struct cache *cache,
    const struct cache_key *base)
{
	struct cache_key changed;
	unsigned char output[B32_DATA_SIZE];
	unsigned char payload[B32_DATA_SIZE];
	enum cache_lookup lookup;
	int error;

	changed = *base;
	++changed.nid;
	lookup = cache_claim(cache, base, output, &error);
	if (lookup != CACHE_HIT)
		fail("ready base key was not retained");
	lookup = cache_claim(cache, &changed, output, &error);
	if (lookup != CACHE_OWNER)
		fail("ready eviction did not create a new owner");
	lookup = cache_claim(cache, base, output, &error);
	if (lookup != CACHE_BYPASS)
		fail("different key did not use no-cache inflight fallback");
	memset(payload, 0x7c, sizeof(payload));
	cache_complete(cache, &changed, copy_payload(payload, sizeof(payload)), 0);

	check_pthread(pthread_mutex_lock(&cache->lock), "pthread_mutex_lock");
	cache->waiters = 1;
	check_pthread(pthread_mutex_unlock(&cache->lock), "pthread_mutex_unlock");
	lookup = cache_claim(cache, base, output, &error);
	if (lookup != CACHE_BYPASS)
		fail("published generation was evicted before waiter consumption");
	check_pthread(pthread_mutex_lock(&cache->lock), "pthread_mutex_lock");
	cache->waiters = 0;
	check_pthread(pthread_cond_broadcast(&cache->cv),
	    "pthread_cond_broadcast");
	check_pthread(pthread_mutex_unlock(&cache->lock), "pthread_mutex_unlock");
}

static void *
fini_worker(void *opaque)
{
	struct fini_arg *arg;
	struct cache *cache;
	unsigned char *data;

	arg = opaque;
	cache = arg->cache;
	check_pthread(pthread_mutex_lock(&cache->lock), "pthread_mutex_lock");
	cache->closing = true;
	while (cache->state == CACHE_INFLIGHT || cache->waiters != 0)
		check_pthread(pthread_cond_wait(&cache->cv, &cache->lock),
		    "pthread_cond_wait");
	data = cache->data;
	cache->data = NULL;
	cache->error = 0;
	cache->state = CACHE_EMPTY;
	check_pthread(pthread_mutex_unlock(&cache->lock), "pthread_mutex_unlock");
	free(data);
	atomic_store_explicit(&arg->done, true, memory_order_release);
	return (NULL);
}

static void
test_shutdown(struct cache *cache, const struct cache_key *key)
{
	struct fini_arg arg;
	struct timespec delay = { .tv_sec = 0, .tv_nsec = 20000000 };
	unsigned char output[B32_DATA_SIZE];
	unsigned char payload[B32_DATA_SIZE];
	pthread_t thread;
	enum cache_lookup lookup;
	int error;

	lookup = cache_claim(cache, key, output, &error);
	if (lookup != CACHE_OWNER)
		fail("shutdown setup did not create inflight owner");
	arg.cache = cache;
	atomic_init(&arg.done, false);
	check_pthread(pthread_create(&thread, NULL, fini_worker, &arg),
	    "pthread_create");
	nanosleep(&delay, NULL);
	if (atomic_load_explicit(&arg.done, memory_order_acquire))
		fail("shutdown did not wait for inflight owner");
	memset(payload, 0x91, sizeof(payload));
	cache_complete(cache, key, copy_payload(payload, sizeof(payload)), 0);
	check_pthread(pthread_join(thread, NULL), "pthread_join");
	if (!atomic_load_explicit(&arg.done, memory_order_acquire))
		fail("shutdown did not finish after owner publication");
	check_pthread(pthread_cond_destroy(&cache->cv), "pthread_cond_destroy");
	check_pthread(pthread_mutex_destroy(&cache->lock), "pthread_mutex_destroy");
}

int
main(void)
{
	struct cache cache;
	struct cache_key key, failure_key, retry_key;
	unsigned char output[B32_DATA_SIZE];
	unsigned char payload[B32_DATA_SIZE];
	enum cache_lookup lookup;
	int error;

	cache_init(&cache);
	key = base_key();
	run_wave(&cache, &key, 0);
	failure_key = key;
	++failure_key.nid;
	run_wave(&cache, &failure_key, EIO);
	lookup = cache_claim(&cache, &failure_key, output, &error);
	if (lookup != CACHE_OWNER)
		fail("published failure did not become retryable");
	memset(payload, 0x44, sizeof(payload));
	cache_complete(&cache, &failure_key,
	    copy_payload(payload, sizeof(payload)), 0);

	retry_key = key;
	retry_key.nid += 2;
	test_exact_key(&cache, &retry_key);
	test_fallback_and_waiter_generation(&cache, &retry_key);
	retry_key.nid += 3;
	test_shutdown(&cache, &retry_key);
	printf("TC168-cache-inflight host oracle PASS: one owner, exact key, typed failure, retry, fallback, eviction, shutdown\n");
	return (0);
}