//
// Cforall Version 1.0.0 Copyright (C) 2020 University of Waterloo
//
// The contents of this file are covered under the licence agreement in the
// file "LICENCE" distributed with Cforall.
//
// io.cfa --
//
// Author           : Thierry Delisle
// Created On       : Thu Apr 23 17:31:00 2020
// Last Modified By :
// Last Modified On :
// Update Count     :
//

#define __cforall_thread__

#if defined(__CFA_DEBUG__)
	// #define __CFA_DEBUG_PRINT_IO__
	// #define __CFA_DEBUG_PRINT_IO_CORE__
#endif


#if defined(CFA_HAVE_LINUX_IO_URING_H)
	#define _GNU_SOURCE         /* See feature_test_macros(7) */
	#include <errno.h>
	#include <signal.h>
	#include <stdint.h>
	#include <string.h>
	#include <unistd.h>

	extern "C" {
		#include <sys/syscall.h>

		#include <linux/io_uring.h>
	}

	#include "stats.hfa"
	#include "kernel.hfa"
	#include "kernel/fwd.hfa"
	#include "io/types.hfa"

	__attribute__((unused)) static const char * opcodes[] = {
		"OP_NOP",
		"OP_READV",
		"OP_WRITEV",
		"OP_FSYNC",
		"OP_READ_FIXED",
		"OP_WRITE_FIXED",
		"OP_POLL_ADD",
		"OP_POLL_REMOVE",
		"OP_SYNC_FILE_RANGE",
		"OP_SENDMSG",
		"OP_RECVMSG",
		"OP_TIMEOUT",
		"OP_TIMEOUT_REMOVE",
		"OP_ACCEPT",
		"OP_ASYNC_CANCEL",
		"OP_LINK_TIMEOUT",
		"OP_CONNECT",
		"OP_FALLOCATE",
		"OP_OPENAT",
		"OP_CLOSE",
		"OP_FILES_UPDATE",
		"OP_STATX",
		"OP_READ",
		"OP_WRITE",
		"OP_FADVISE",
		"OP_MADVISE",
		"OP_SEND",
		"OP_RECV",
		"OP_OPENAT2",
		"OP_EPOLL_CTL",
		"OP_SPLICE",
		"OP_PROVIDE_BUFFERS",
		"OP_REMOVE_BUFFERS",
		"OP_TEE",
		"INVALID_OP"
	};

//=============================================================================================
// I/O Syscall
//=============================================================================================
	static int __io_uring_enter( struct $io_context & ctx, unsigned to_submit, bool get ) {
		__STATS__( false, io.calls.count++; )
		bool need_sys_to_submit = false;
		bool need_sys_to_complete = false;
		unsigned flags = 0;

		TO_SUBMIT:
		if( to_submit > 0 ) {
			if( !(ctx.ring_flags & IORING_SETUP_SQPOLL) ) {
				need_sys_to_submit = true;
				break TO_SUBMIT;
			}
			if( (*ctx.sq.flags) & IORING_SQ_NEED_WAKEUP ) {
				need_sys_to_submit = true;
				flags |= IORING_ENTER_SQ_WAKEUP;
			}
		}

		if( get && !(ctx.ring_flags & IORING_SETUP_SQPOLL) ) {
			flags |= IORING_ENTER_GETEVENTS;
			if( (ctx.ring_flags & IORING_SETUP_IOPOLL) ) {
				need_sys_to_complete = true;
			}
		}

		int ret = 0;
		if( need_sys_to_submit || need_sys_to_complete ) {
			__cfadbg_print_safe(io_core, "Kernel I/O : IO_URING enter %d %u %u\n", ctx.fd, to_submit, flags);
			__STATS__( false, io.calls.blocks++; )
			ret = syscall( __NR_io_uring_enter, ctx.fd, to_submit, 0, flags, (sigset_t *)0p, _NSIG / 8);
			__cfadbg_print_safe(io_core, "Kernel I/O : IO_URING %d returned %d\n", ctx.fd, ret);
		}

		// Memory barrier
		__atomic_thread_fence( __ATOMIC_SEQ_CST );
		return ret;
	}

//=============================================================================================
// I/O Polling
//=============================================================================================
	static inline unsigned __flush( struct $io_context & );
	static inline __u32 __release_sqes( struct $io_context & );

	static bool __drain_io( struct  $io_context & ctx ) {
		unsigned to_submit = __flush( ctx );
		int ret = __io_uring_enter( ctx, to_submit, true );
		if( ret < 0 ) {
			switch((int)errno) {
			case EAGAIN:
			case EINTR:
			case EBUSY:
				// Update statistics
				__STATS__( false, io.calls.errors.busy ++; )
				return true;
				break;
			default:
				abort( "KERNEL ERROR: IO_URING SYSCALL - (%d) %s\n", (int)errno, strerror(errno) );
			}
		}

		// update statistics
		if (to_submit > 0) {
			__STATS__( false, io.calls.submitted += ret; )
			/* paranoid */ verify( ctx.sq.to_submit <= *ctx.sq.num );

			/* paranoid */ verify( ctx.sq.to_submit >= ret );
			ctx.sq.to_submit -= ret;

			/* paranoid */ verify( ctx.sq.to_submit <= *ctx.sq.num );

			if(ret) {
				__cfadbg_print_safe(io, "Kernel I/O : %u submitted to io_uring\n", ret);
			}
		}

		// Release the consumed SQEs
		__release_sqes( ctx );

		// Drain the queue
		unsigned head = *ctx.cq.head;
		unsigned tail = *ctx.cq.tail;
		const __u32 mask = *ctx.cq.mask;

		// Nothing was new return 0
		if (head == tail) {
			return ctx.sq.to_submit > 0;
		}

		__u32 count = tail - head;
		/* paranoid */ verify( count != 0 );
		__STATS__( false, io.calls.completed += count; )

		for(i; count) {
			unsigned idx = (head + i) & mask;
			volatile struct io_uring_cqe & cqe = ctx.cq.cqes[idx];

			/* paranoid */ verify(&cqe);

			struct io_future_t * future = (struct io_future_t *)(uintptr_t)cqe.user_data;
			__cfadbg_print_safe( io, "Kernel I/O : Syscall completed : cqe %p, result %d for %p\n", &cqe, cqe.res, future );

			fulfil( *future, cqe.res );
		}

		if(count) {
			__cfadbg_print_safe(io, "Kernel I/O : %u completed\n", count);
		}

		// Mark to the kernel that the cqe has been seen
		// Ensure that the kernel only sees the new value of the head index after the CQEs have been read.
		__atomic_store_n( ctx.cq.head, head + count, __ATOMIC_SEQ_CST );

		return count > 0 || to_submit > 0;
	}

	void main( $io_context & this ) {
		__cfadbg_print_safe(io_core, "Kernel I/O : IO poller %d (%p) ready\n", this.fd, &this);

		const int reset_cnt = 5;
		int reset = reset_cnt;
		// Then loop until we need to start
		LOOP:
		while() {
			waitfor( ^?{} : this) {
				break LOOP;
			}
			or else {}

			// Drain the io
			bool again = __drain_io( this );

			if(!again) reset--;

			// If we got something, just yield and check again
			if(reset > 1) {
				yield();
				continue LOOP;
			}

			// We alread failed to find completed entries a few time.
			if(reset == 1) {
				// Rearm the context so it can block
				// but don't block right away
				// we need to retry one last time in case
				// something completed *just now*
				__ioctx_prepare_block( this );
				continue LOOP;
			}

			__STATS__( false,
				io.poller.sleeps += 1;
			)
			__cfadbg_print_safe(io_core, "Kernel I/O : Parking io poller %d (%p)\n", this.fd, &this);

			// block this thread
			wait( this.sem );

			// restore counter
			reset = reset_cnt;
		}

		__cfadbg_print_safe(io_core, "Kernel I/O : Fast poller %d (%p) stopping\n", this.fd, &this);
	}

//=============================================================================================
// I/O Submissions
//=============================================================================================

// Submition steps :
// 1 - Allocate a queue entry. The ring already has memory for all entries but only the ones
//     listed in sq.array are visible by the kernel. For those not listed, the kernel does not
//     offer any assurance that an entry is not being filled by multiple flags. Therefore, we
//     need to write an allocator that allows allocating concurrently.
//
// 2 - Actually fill the submit entry, this is the only simple and straightforward step.
//
// 3 - Append the entry index to the array and adjust the tail accordingly. This operation
//     needs to arrive to two concensus at the same time:
//     A - The order in which entries are listed in the array: no two threads must pick the
//         same index for their entries
//     B - When can the tail be update for the kernel. EVERY entries in the array between
//         head and tail must be fully filled and shouldn't ever be touched again.
//

	static $io_context * __ioarbiter_allocate( $io_arbiter & mutex this, processor *, __u32 idxs[], __u32 want );
	static void __ioarbiter_submit  ( $io_arbiter & mutex this, $io_context * , __u32 idxs[], __u32 have );
	static void __ioarbiter_flush   ( $io_arbiter & mutex this, $io_context * );
	static inline void __ioarbiter_notify( $io_context & ctx );

	//=============================================================================================
	// Allocation
	// for user's convenience fill the sqes from the indexes
	static inline void __fill(struct io_uring_sqe * out_sqes[], __u32 want, __u32 idxs[], struct $io_context * ctx)  {
		struct io_uring_sqe * sqes = ctx->sq.sqes;
		for(i; want) {
			out_sqes[i] = &sqes[idxs[i]];
		}
	}

	// Try to directly allocate from the a given context
	// Not thread-safe
	static inline bool __alloc(struct $io_context * ctx, __u32 idxs[], __u32 want) {
		__sub_ring_t & sq = ctx->sq;
		const __u32 mask  = *sq.mask;
		__u32 fhead = sq.free_ring.head;    // get the current head of the queue
		__u32 ftail = sq.free_ring.tail;    // get the current tail of the queue

		// If we don't have enough sqes, fail
		if((ftail - fhead) < want) { return false; }

		// copy all the indexes we want from the available list
		for(i; want) {
			idxs[i] = sq.free_ring.array[(fhead + i) & mask];
		}

		// Advance the head to mark the indexes as consumed
		__atomic_store_n(&sq.free_ring.head, fhead + want, __ATOMIC_RELEASE);

		// return success
		return true;
	}

	// Allocate an submit queue entry.
	// The kernel cannot see these entries until they are submitted, but other threads must be
	// able to see which entries can be used and which are already un used by an other thread
	// for convenience, return both the index and the pointer to the sqe
	// sqe == &sqes[idx]
	struct $io_context * cfa_io_allocate(struct io_uring_sqe * sqes[], __u32 idxs[], __u32 want) {
		__cfadbg_print_safe(io, "Kernel I/O : attempting to allocate %u\n", want);

		disable_interrupts();
		processor * proc = __cfaabi_tls.this_processor;
		/* paranoid */ verify( __cfaabi_tls.this_processor );
		/* paranoid */ verify( proc->io.lock == false );

		__atomic_store_n( &proc->io.lock, true, __ATOMIC_SEQ_CST );
		$io_context * ctx = proc->io.ctx;
		$io_arbiter * ioarb = proc->cltr->io.arbiter;
		/* paranoid */ verify( ioarb );

		// Can we proceed to the fast path
		if(  ctx				// We alreay have an instance?
		&&  !ctx->revoked )		// Our instance is still valid?
		{
			__cfadbg_print_safe(io, "Kernel I/O : attempting to fast allocation\n");

			// We can proceed to the fast path
			if( __alloc(ctx, idxs, want) ) {
				// Allocation was successful
				// Mark the instance as no longer in-use and re-enable interrupts
				__atomic_store_n( &proc->io.lock, false, __ATOMIC_RELEASE );
				__STATS__( true, io.alloc.fast += 1; )
				enable_interrupts( __cfaabi_dbg_ctx );

				__cfadbg_print_safe(io, "Kernel I/O : fast allocation successful\n");

				__fill( sqes, want, idxs, ctx );
				return ctx;
			}
			// The fast path failed, fallback
			__STATS__( true, io.alloc.fail += 1; )
		}

		// Fast path failed, fallback on arbitration
		__atomic_store_n( &proc->io.lock, false, __ATOMIC_RELEASE );
		__STATS__( true, io.alloc.slow += 1; )
		enable_interrupts( __cfaabi_dbg_ctx );

		__cfadbg_print_safe(io, "Kernel I/O : falling back on arbiter for allocation\n");

		struct $io_context * ret = __ioarbiter_allocate(*ioarb, proc, idxs, want);

		__cfadbg_print_safe(io, "Kernel I/O : slow allocation completed\n");

		__fill( sqes, want, idxs,ret );
		return ret;
	}


	//=============================================================================================
	// submission
	static inline void __submit( struct $io_context * ctx, __u32 idxs[], __u32 have) {
		// We can proceed to the fast path
		// Get the right objects
		__sub_ring_t & sq = ctx->sq;
		const __u32 mask  = *sq.mask;
		__u32 tail = sq.kring.ready;

		// Add the sqes to the array
		for( i; have ) {
			sq.kring.array[ (tail + i) & mask ] = idxs[i];
		}

		// Make the sqes visible to the submitter
		__atomic_store_n(&sq.kring.ready, tail + have, __ATOMIC_RELEASE);

		// Make sure the poller is awake
		__cfadbg_print_safe(io, "Kernel I/O : waking the poller\n");
		post( ctx->sem );
	}

	void cfa_io_submit( struct $io_context * inctx, __u32 idxs[], __u32 have ) __attribute__((nonnull (1))) {
		__cfadbg_print_safe(io, "Kernel I/O : attempting to submit %u\n", have);

		disable_interrupts();
		processor * proc = __cfaabi_tls.this_processor;
		/* paranoid */ verify( __cfaabi_tls.this_processor );
		/* paranoid */ verify( proc->io.lock == false );

		__atomic_store_n( &proc->io.lock, true, __ATOMIC_SEQ_CST );
		$io_context * ctx = proc->io.ctx;

		// Can we proceed to the fast path
		if(  ctx				// We alreay have an instance?
		&&  !ctx->revoked		// Our instance is still valid?
		&&   ctx == inctx )		// We have the right instance?
		{
			__submit(ctx, idxs, have);

			// Mark the instance as no longer in-use, re-enable interrupts and return
			__atomic_store_n( &proc->io.lock, false, __ATOMIC_RELEASE );
			__STATS__( true, io.submit.fast += 1; )
			enable_interrupts( __cfaabi_dbg_ctx );

			__cfadbg_print_safe(io, "Kernel I/O : submitted on fast path\n");
			return;
		}

		// Fast path failed, fallback on arbitration
		__atomic_store_n( &proc->io.lock, false, __ATOMIC_RELEASE );
		__STATS__( true, io.submit.slow += 1; )
		enable_interrupts( __cfaabi_dbg_ctx );

		__cfadbg_print_safe(io, "Kernel I/O : falling back on arbiter for submission\n");

		__ioarbiter_submit(*inctx->arbiter, inctx, idxs, have);
	}

	//=============================================================================================
	// Flushing
	static unsigned __flush( struct $io_context & ctx ) {
		// First check for external
		if( !__atomic_load_n(&ctx.ext_sq.empty, __ATOMIC_SEQ_CST) ) {
			// We have external submissions, delegate to the arbiter
			__ioarbiter_flush( *ctx.arbiter, &ctx );
		}

		__u32 tail  = *ctx.sq.kring.tail;
		__u32 ready = ctx.sq.kring.ready;

		/* paranoid */ verify( ctx.sq.to_submit <= *ctx.sq.num );
		ctx.sq.to_submit += (ready - tail);
		/* paranoid */ verify( ctx.sq.to_submit <= *ctx.sq.num );

		if(ctx.sq.to_submit) {
			__cfadbg_print_safe(io, "Kernel I/O : %u ready to submit\n", ctx.sq.to_submit);
		}

		__atomic_store_n(ctx.sq.kring.tail, ready, __ATOMIC_RELEASE);

		return ctx.sq.to_submit;
	}


	// Go through the ring's submit queue and release everything that has already been consumed
	// by io_uring
	// This cannot be done by multiple threads
	static __u32 __release_sqes( struct $io_context & ctx ) {
		const __u32 mask = *ctx.sq.mask;

		__attribute__((unused))
		__u32 ctail = *ctx.sq.kring.tail;    // get the current tail of the queue
		__u32 chead = *ctx.sq.kring.head;	 // get the current head of the queue
		__u32 phead = ctx.sq.kring.released; // get the head the last time we were here

		__u32 ftail = ctx.sq.free_ring.tail;  // get the current tail of the queue

		// the 3 fields are organized like this diagram
		// except it's are ring
		// ---+--------+--------+----
		// ---+--------+--------+----
		//    ^        ^        ^
		// phead    chead    ctail

		// make sure ctail doesn't wrap around and reach phead
		/* paranoid */ verify(
			   (ctail >= chead && chead >= phead)
			|| (chead >= phead && phead >= ctail)
			|| (phead >= ctail && ctail >= chead)
		);

		// find the range we need to clear
		__u32 count = chead - phead;

		if(count == 0) {
			return 0;
		}

		// We acquired an previous-head/current-head range
		// go through the range and release the sqes
		for( i; count ) {
			__u32 idx = ctx.sq.kring.array[ (phead + i) & mask ];
			ctx.sq.free_ring.array[ (ftail + i) & mask ] = idx;
		}

		ctx.sq.kring.released = chead;		// note up to were we processed
		__atomic_store_n(&ctx.sq.free_ring.tail, ftail + count, __ATOMIC_SEQ_CST);

		__ioarbiter_notify(ctx);

		return count;
	}

//=============================================================================================
// I/O Arbiter
//=============================================================================================
	static inline void __revoke( $io_arbiter & this, $io_context * ctx ) {
		if(ctx->revoked) return;

		/* paranoid */ verify( ctx->proc );
		remove( this.assigned, *ctx );

		// Mark as revoked
		__atomic_store_n(&ctx->revoked, true, __ATOMIC_SEQ_CST);

		// Wait for the processor to no longer use it
		while(ctx->proc->io.lock) Pause();

		// Remove the coupling with the processor
		ctx->proc->io.ctx = 0p;
		ctx->proc = 0p;

		// add to available contexts
		addHead( this.available, *ctx );
	}

	static inline void __assign( $io_arbiter & this, $io_context * ctx, processor * proc ) {
		remove( this.available, *ctx );

		ctx->revoked = false;
		ctx->proc = proc;
		__atomic_store_n(&proc->io.ctx, ctx, __ATOMIC_SEQ_CST);

		// add to assigned contexts
		addTail( this.assigned, *ctx );
	}

	static $io_context * __ioarbiter_allocate( $io_arbiter & mutex this, processor * proc, __u32 idxs[], __u32 want ) {
		__cfadbg_print_safe(io, "Kernel I/O : arbiter allocating\n");

		SeqIter($io_context) iter;
		$io_context & ci;
		// Do we already have something available?
		for( over( iter, this.available ); iter | ci;) {
			__cfadbg_print_safe(io, "Kernel I/O : attempting available context\n");

			$io_context * c = &ci;
			if(__alloc(c, idxs, want)) {
				__assign( this, c, proc);
				return c;
			}
		}


		// Otherwise, we have no choice but to revoke everyone to check if other instance have available data
		for( over( iter, this.assigned ); iter | ci; ) {
			__cfadbg_print_safe(io, "Kernel I/O : revoking context for allocation\n");

			$io_context * c = &ci;
			__revoke( this, c );

			__STATS__( false, io.alloc.revoke += 1; )

			if(__alloc(c, idxs, want)) {
				__assign( this, c, proc);
				return c;
			}
		}

		__cfadbg_print_safe(io, "Kernel I/O : waiting for available resources\n");

		__STATS__( false, io.alloc.block += 1; )

		// No one has any resources left, wait for something to finish
		// Mark as pending
		__atomic_store_n( &this.pending.flag, true, __ATOMIC_SEQ_CST );

		// Wait for our turn to submit
		wait( this.pending.blocked, want );

		__attribute((unused)) bool ret =
		__alloc( this.pending.ctx, idxs, want);
		/* paranoid */ verify( ret );

		__assign( this, this.pending.ctx, proc);
		return this.pending.ctx;
	}

	static void __ioarbiter_notify( $io_arbiter & mutex this, $io_context * ctx ) {
		/* paranoid */ verify( !is_empty(this.pending.blocked) );
		this.pending.ctx = ctx;

		while( !is_empty(this.pending.blocked) ) {
			__u32 have = ctx->sq.free_ring.tail - ctx->sq.free_ring.head;
			__u32 want = front( this.pending.blocked );

			if( have > want ) return;

			signal_block( this.pending.blocked );
		}

		this.pending.flag = false;
	}

	static void __ioarbiter_notify( $io_context & ctx ) {
		if(__atomic_load_n( &ctx.arbiter->pending.flag, __ATOMIC_SEQ_CST)) {
			__ioarbiter_notify( *ctx.arbiter, &ctx );
		}
	}

	// Simply append to the pending
	static void __ioarbiter_submit( $io_arbiter & mutex this, $io_context * ctx, __u32 idxs[], __u32 have ) {
		__cfadbg_print_safe(io, "Kernel I/O : submitting %u from the arbiter to context %u\n", have, ctx->fd);

		/* paranoid */ verify( &this == ctx->arbiter );

		// Mark as pending
		__atomic_store_n( &ctx->ext_sq.empty, false, __ATOMIC_SEQ_CST );

		// Wake-up the poller
		post( ctx->sem );

		__cfadbg_print_safe(io, "Kernel I/O : waiting to submit %u\n", have);

		// Wait for our turn to submit
		wait( ctx->ext_sq.blocked );

		// Submit our indexes
		__submit(ctx, idxs, have);

		__cfadbg_print_safe(io, "Kernel I/O : %u submitted from arbiter\n", have);
	}

	static void __ioarbiter_flush( $io_arbiter & mutex this, $io_context * ctx ) {
		/* paranoid */ verify( &this == ctx->arbiter );

		__STATS__( false, io.flush.external += 1; )

		__revoke( this, ctx );

		__cfadbg_print_safe(io, "Kernel I/O : arbiter flushing\n");

		condition & blcked = ctx->ext_sq.blocked;
		/* paranoid */ verify( ctx->ext_sq.empty == is_empty( blcked ) );
		while(!is_empty( blcked )) {
			signal_block( blcked );
		}

		ctx->ext_sq.empty = true;
	}

	void __ioarbiter_register( $io_arbiter & mutex this, $io_context & ctx ) {
		__cfadbg_print_safe(io, "Kernel I/O : registering new context\n");

		ctx.arbiter = &this;

		// add to available contexts
		addHead( this.available, ctx );

		// Check if this solves pending allocations
		if(this.pending.flag) {
			__ioarbiter_notify( ctx );
		}
	}

	void __ioarbiter_unregister( $io_arbiter & mutex this, $io_context & ctx ) {
		/* paranoid */ verify( &this == ctx.arbiter );

		__revoke( this, &ctx );

		remove( this.available, ctx );
	}
#endif