IoContext.cpp

// SPDX-License-Identifier: LGPL-3.0-or-later
// Copyright © 2020 Florian Fischer
#include "io/IoContext.hpp"

#include <liburing.h>						// for io_uring_get_sqe, io_uring_submit
#include <liburing/io_uring.h>	// for io_uring_cqe
#include <sys/eventfd.h>				// for eventfd
#include <sys/socket.h>					// for socklen_t, recv, send, MSG_DON...
#include <sys/types.h>					// for off_t, ssize_t

#include <atomic>		// for atomic
#include <cassert>	// for assert
#include <cerrno>		// for errno, EAGAIN, EWOULDBLOCK, EINTR
#include <cstdio>		// for perrror
#include <cstdlib>	// for exit
#include <memory>		// for make_unique, unique_ptr
#include <vector>		// for push_back

#include "CallerEnvironment.hpp"	// for CallerEnvironment
#include "Common.hpp"							// for DIE, unlikely
#include "Debug.hpp"							// for LogSubsystem, LogSubsyst...
#include "Emper.hpp"							// for IO_POLLING_IO_THREAD, DEBUG
#include "Runtime.hpp"						// for Runtime::getRuntime
#include "emper-config.h"					// IWYU pragma: keep
#include "io/Future.hpp"					// for emper::io::Future
#include "io/Operation.hpp"				// for emper::io::Operation
#include "io/Stats.hpp"						// for emper::io::Stats

#ifndef EMPER_LOG_OFF
#include <ostream>	// for basic_osteram::operator<<, operator<<
#endif

namespace emper::io {
thread_local IoContext *IoContext::workerIo = nullptr;

pthread_t IoContext::globalCompleter;

auto IoContext::prepareFutureChain(Future &future, bool recursion) -> int {
	int prepared = 0;
	if (future.dependency) {
		prepared += prepareFutureChain(*future.dependency, true);
	}

	if constexpr (emper::DEBUG) {
		reqs_in_uring++;
		uringFutureSet.insert(&future);
	}

	struct io_uring_sqe *sqe = io_uring_get_sqe(&ring);
	// The kernel consumes all sqes from the SQ during io_uring_enter
	// If we can't get a sqe the chain was to long for our SQ
	assert(sqe);

	future.prepareSqe(sqe);
	io_uring_sqe_set_data(sqe, &future);

	if (recursion) {
		sqe->flags |= IOSQE_IO_LINK;
	}

	return prepared + 1;
}

void IoContext::submit(Future &future) {
	LOGD("submitting " << future.op << " Future " << &future
										 << (future.dependency ? " and it's dependencies" : ""));
#ifdef NDEBUG
	UNUSED_ARG
#endif
	int prepared = prepareFutureChain(future);

	// submit the Future to the io_uring
	int submitted = io_uring_submit(&ring);

	// We can't submit our sqe because the CQ is full
	// in worker thread -> reapCompletions
	// in globalCompleter thread -> TODO: deterministically handle global full CQ
	// for now hope a jam does not happen or will solve itself
	if (unlikely(submitted == -EBUSY)) {
		if constexpr (emper::DEBUG) {
			std::stringstream sst;
			sst << "io_submit returned EBUSY trying to submit in addition to " << reqs_in_uring
					<< std::endl;
			logI(sst.str());
		} else {
			logI("io_submit returned EBUSY");
		}

		TIME_NS(
				{
					do {
						reapCompletions();
					} while ((submitted = io_uring_submit(&ring)) == -EBUSY);
				},
				stats.record_io_submit_full_cq);
	}

	if (unlikely(submitted < 0)) {
		errno = -submitted;
		DIE_MSG_ERRNO("io_uring_submit failed");
	}

	// submit the Future to the io_uring
	// TODO: handle submit errors in Future chains
	// req1 -> invalid_req -> req3
	// will submit only 2 instead of all 3 prepared sqes
	// See: https://github.com/axboe/liburing/issues/186
	assert(submitted == prepared);

	// io_uring will try to synchronously complete any IO request before
	// offloading it to the async backend. See io_uring_enter(2).
	// Immediate offloading can be enforced by setting the IOSQE_ASYNC flag in the sqe.
	// Try to reap a possible synchronous completion.
	reapCompletions();
}

template <CallerEnvironment callerEnvironment>
void IoContext::reapCompletions() {
	// Someone else is currently reaping completions
	if (unlikely(!cq_mutex.try_lock())) {
		return;
	}

	// never reap completions on the global IoContext
	assert(this != getGlobalIo());

	LOGD("Reaping completions");
	unsigned head;
	struct io_uring_cqe *cqe;
	unsigned count = 0;

	int err = io_uring_peek_cqe(&ring, &cqe);
	if (err) {
		if (err == -EAGAIN) {
			goto unlock;
		}
		errno = -err;
		DIE_MSG_ERRNO("io_uring_peek_cqe failed");
	}

	io_uring_for_each_cqe(&ring, head, cqe) {
		count++;

		auto *future = reinterpret_cast<Future *>(io_uring_cqe_get_data(cqe));
		assert(future);

		if constexpr (emper::DEBUG) {
			// assert that the future was previously in the uringFutureSet
			assert(uringFutureSet.erase(future) > 0);
		}

		future->complete<callerEnvironment>(cqe->res);
	}

	LOGD("got " << count << " cqes from the io_uring");
	io_uring_cq_advance(&ring, count);

	if constexpr (emper::DEBUG) {
		assert(count <= reqs_in_uring);
		reqs_in_uring -= count;
	}

	if constexpr (callerEnvironment == CallerEnvironment::EMPER) {
		stats.record_worker_reaps(count);
	} else {
		stats.record_completer_reaps(count);
	}

unlock:
	cq_mutex.unlock();
}

// Show the compiler our template incarnations
template void IoContext::reapCompletions<CallerEnvironment::ANYWHERE>();
template void IoContext::reapCompletions<CallerEnvironment::EMPER>();

// The globalCompleter must be started after all worker io_uring eventfds are submitted
// so we don't have to synchronize the global sq
auto IoContext::globalCompleterFunc(void *arg) -> void * {
	IoContext &io = *reinterpret_cast<IoContext *>(arg);

	io.logD("submit all worker io_uring eventfds");

	// submit all eventfds in the SQ inserted by IoContext::submit_efd calls
	int submitted = io_uring_submit(&io.ring);
	if (unlikely(submitted < 0)) {
		DIE_MSG_ERRNO("initial global io_uring submit failed");
	}

	// We have submitted all eventfds
	assert(submitted == Runtime::getRuntime()->getWorkerCount());

	io.logD("start global completer loop");
	while (true) {
		struct io_uring_cqe *cqe;

		// wait for completions
		int err = io_uring_wait_cqe(&io.ring, &cqe);
		if (unlikely(err)) {
			if (err == -EINTR) {
				continue;
			}
			errno = -err;
			perror("io_uring_wait_cqe");
		}

		auto *workerIo = reinterpret_cast<IoContext *>(io_uring_cqe_get_data(cqe));
		assert(workerIo);

		io_uring_cqe_seen(&io.ring, cqe);

		// re-add the eventfd read
		struct io_uring_sqe *sqe = io_uring_get_sqe(&io.ring);
		// we initialized the global ring big enough we should always get a free sqe
		assert(sqe);

		io_uring_prep_read(sqe, workerIo->ring_eventfd, &workerIo->ring_eventfd_readbuf,
											 sizeof(workerIo->ring_eventfd_readbuf), 0);
		io_uring_sqe_set_data(sqe, reinterpret_cast<void *>(workerIo));

		submitted = io_uring_submit(&io.ring);

		if (unlikely(submitted < 0)) {
			errno = -submitted;
			DIE_MSG_ERRNO("re-submitting eventfd read to global_ring failed");
		}

		assert(submitted == 1);

		workerIo->reapCompletions<CallerEnvironment::ANYWHERE>();
	}

	return nullptr;
}

void IoContext::startGlobalCompleter(IoContext &globalIo) {
	int err = pthread_create(&globalCompleter, nullptr, globalCompleterFunc, &globalIo);
	if (unlikely(err)) {
		errno = err;
		DIE_MSG_ERRNO("Creating global completer thread failed");
	}
}

IoContext::IoContext(size_t uring_entries) {
	unsigned flags = 0;
	if constexpr (emper::IO_URING_SQPOLL) {
		// check if SQPOLL with unregistered files is supported by the kernel
		static int support = -1;

		if (support == -1) {
			support = compareKernelVersion("5.11") >= 0 ? 1 : 0;
		}

		if (!support) {
			DIE_MSG("SQPOLL with un-registered files only supportet since linux 5.11");
		}

		flags |= IORING_SETUP_SQPOLL;
	}

	auto ret = io_uring_queue_init(uring_entries, &ring, flags);
	if (ret < 0) {
		// skip tests if we are not able to setup an io_uring
		perror("io_uring_queue_init failed");
		exit(77);
	}
	LOGD("Ring fd is " << ring.ring_fd);

	ring_eventfd = eventfd(0, 0);
	if (unlikely(ring_eventfd < 0)) {
		DIE_MSG_ERRNO("creating eventfd for io_uring failed");
	}

	if (unlikely(io_uring_register_eventfd(&ring, ring_eventfd) < 0)) {
		DIE_MSG_ERRNO("io_uring_register_eventfd failed");
	}

	if constexpr (emper::STATS) {
		Stats::workerStats.push_back(&stats);
	}
}

// This function must not be executed in parallel because it does not synchronize
// the global io_uring's SQ
void IoContext::submit_efd() {
	// the global IoContext must be initialized to submit our eventfd
	assert(getGlobalIo());

	IoContext &globalIo = *getGlobalIo();

	struct io_uring_sqe *sqe = io_uring_get_sqe(&globalIo.ring);
	// we initialized the global ring big enough we should always get a free sqe
	assert(sqe);

	io_uring_prep_read(sqe, ring_eventfd, &ring_eventfd_readbuf, sizeof(ring_eventfd_readbuf), 0);
	io_uring_sqe_set_data(sqe, this);

	// The sqe we prepared will be submitted to io_uring when the globalCompleter starts.
}

IoContext::~IoContext() {
	io_uring_queue_exit(&ring);
	// TODO: check if this is safe
	close(ring_eventfd);
}

auto IoContext::connect(int socket, const struct sockaddr *address, socklen_t address_len)
		-> std::unique_ptr<Future> {
	auto future = std::make_unique<ConnectFuture>(socket, address, address_len);
	submit(*future);
	return future;
}

auto IoContext::accept(int socket, struct sockaddr *address, socklen_t *address_len)
		-> std::unique_ptr<Future> {
	auto future = std::make_unique<AcceptFuture>(socket, address, address_len);
	submit(*future);
	return future;
}

auto IoContext::recv(int socket, void *buffer, size_t length, int flags)
		-> std::unique_ptr<Future> {
	auto future = std::make_unique<RecvFuture>(socket, buffer, length, flags);
	submit(*future);
	return future;
}

auto IoContext::send(int socket, const void *buffer, size_t length, int flags)
		-> std::unique_ptr<Future> {
	auto future = std::make_unique<SendFuture>(socket, buffer, length, flags);
	submit(*future);
	return future;
}

auto IoContext::readFile(int fildes, void *buf, size_t nbyte, off_t offset)
		-> std::unique_ptr<Future> {
	auto future = std::make_unique<ReadFuture>(fildes, buf, nbyte, offset);
	submit(*future);
	return future;
}

auto IoContext::writeFile(int fildes, const void *buf, size_t nbyte, off_t offset)
		-> std::unique_ptr<Future> {
	auto future = std::make_unique<WriteFuture>(fildes, buf, nbyte, offset);
	submit(*future);
	return future;
}

auto IoContext::close(int fildes) -> std::unique_ptr<Future> {
	auto future = std::make_unique<CloseFuture>(fildes);
	submit(*future);
	return future;
}

auto IoContext::connectAndWait(int socket, const struct sockaddr *address, socklen_t address_len)
		-> int {
	ConnectFuture future(socket, address, address_len);
	submit(future);
	return future.waitAndSetErrno();
}

auto IoContext::acceptAndWait(int socket, struct sockaddr *address, socklen_t *address_len) -> int {
	AcceptFuture future(socket, address, address_len);
	submit(future);
	return future.waitAndSetErrno();
}

auto IoContext::recvAndWait(int socket, void *buffer, size_t length, int flags) -> ssize_t {
	if constexpr (emper::IO_TRY_SYSCALL) {
		ssize_t res = ::recv(socket, buffer, length, flags | MSG_DONTWAIT);
		if (res != -1 || errno != EAGAIN) {
			stats.record_completion(Operation::RECV, res, length, true);
			return res;
		}
	}

	RecvFuture future(socket, buffer, length, flags);
	submit(future);
	return future.waitAndSetErrno();
}

auto IoContext::sendAndWait(int socket, const void *buffer, size_t length, int flags) -> ssize_t {
	void *mut_buf = const_cast<void *>(buffer);
	if constexpr (emper::IO_TRY_SYSCALL) {
		ssize_t res = ::send(socket, mut_buf, length, flags | MSG_DONTWAIT);
		if (res != -1 || errno != EAGAIN) {
			stats.record_completion(Operation::SEND, res, length, true);
			return res;
		}
	}

	SendFuture future(socket, mut_buf, length, flags);
	submit(future);
	return future.waitAndSetErrno();
}

auto IoContext::readFileAndWait(int fildes, void *buf, size_t nbyte, off_t offset) -> ssize_t {
	ReadFuture future(fildes, buf, nbyte, offset);
	submit(future);
	return future.waitAndSetErrno();
}

auto IoContext::writeFileAndWait(int fildes, const void *buf, size_t nbyte, off_t offset)
		-> ssize_t {
	void *mut_buf = const_cast<void *>(buf);
	WriteFuture future(fildes, mut_buf, nbyte, offset);
	submit(future);
	return future.waitAndSetErrno();
}

auto IoContext::closeAndWait(int fildes) -> ssize_t {
	CloseFuture future(fildes);
	submit(future);
	return future.waitAndSetErrno();
}
}	 // namespace emper::io