echo.c

/*
 * libecho — kernel-bypass I/O engine on io_uring
 *
 * Async, SQPOLL, buffer pool, linked ops, epoll loop.
 */

#include "echo.h"

#include <errno.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/syscall.h>
#include <unistd.h>

/* ── raw syscalls ───────────────────────────────────────────────────────── */

static int
__echo_setup(unsigned entries, struct io_uring_params *p)
{
	return (int)syscall(__NR_io_uring_setup, entries, p);
}

static int
__echo_enter(int fd, unsigned to_submit, unsigned min_complete, unsigned flags)
{
	return (int)syscall(__NR_io_uring_enter, fd, to_submit,
	    min_complete, flags, NULL, 0);
}

static int
__echo_register(int fd, unsigned opcode, const void *arg, unsigned nr_args)
{
	return (int)syscall(__NR_io_uring_register, fd, opcode, arg, nr_args);
}

/* ── mmap helper ────────────────────────────────────────────────────────── */

static int
__ring_mmap(int fd, off_t offset, size_t len, void **ptr_out)
{
	void *p;
	p = mmap(NULL, len, PROT_READ | PROT_WRITE,
	    MAP_SHARED | MAP_POPULATE, fd, offset);
	if (p == MAP_FAILED) return -errno;
	*ptr_out = p;
	return 0;
}

/* ── ring init ──────────────────────────────────────────────────────────── */

static int
__ring_map(int fd, struct io_uring_params *p, struct le_ring *ring)
{
	int ret;

	ring->sq_ring_sz = p->sq_off.array + p->sq_entries * sizeof(__u32);
	ring->cq_ring_sz = p->cq_off.cqes + p->cq_entries * sizeof(struct io_uring_cqe);
	ring->sqes_sz    = p->sq_entries * sizeof(struct io_uring_sqe);

	ret = __ring_mmap(fd, IORING_OFF_SQ_RING, ring->sq_ring_sz, &ring->sq_ring);
	if (ret) return ret;
	ret = __ring_mmap(fd, IORING_OFF_CQ_RING, ring->cq_ring_sz, &ring->cq_ring);
	if (ret) goto err_sq;
	ret = __ring_mmap(fd, IORING_OFF_SQES, ring->sqes_sz, (void **)&ring->sqes);
	if (ret) goto err_cq;

	ring->sq_head    = (unsigned *)((char *)ring->sq_ring + p->sq_off.head);
	ring->sq_tail    = (unsigned *)((char *)ring->sq_ring + p->sq_off.tail);
	ring->sq_mask    = (unsigned *)((char *)ring->sq_ring + p->sq_off.ring_mask);
	ring->sq_flags   = (unsigned *)((char *)ring->sq_ring + p->sq_off.flags);
	ring->sq_dropped = (unsigned *)((char *)ring->sq_ring + p->sq_off.dropped);
	ring->sq_array   = (__u32     *)((char *)ring->sq_ring + p->sq_off.array);

	ring->cq_head     = (unsigned *)((char *)ring->cq_ring + p->cq_off.head);
	ring->cq_tail     = (unsigned *)((char *)ring->cq_ring + p->cq_off.tail);
	ring->cq_mask     = (unsigned *)((char *)ring->cq_ring + p->cq_off.ring_mask);
	ring->cq_flags    = (unsigned *)((char *)ring->cq_ring + p->cq_off.flags);
	ring->cq_overflow = (unsigned *)((char *)ring->cq_ring + p->cq_off.overflow);
	ring->cqes        = (struct io_uring_cqe *)((char *)ring->cq_ring + p->cq_off.cqes);

	return 0;
err_cq:
	munmap(ring->cq_ring, ring->cq_ring_sz);
err_sq:
	munmap(ring->sq_ring, ring->sq_ring_sz);
	return ret;
}

static int
__ring_init(struct le_ring *ring, unsigned entries, unsigned flags)
{
	struct io_uring_params p;
	int fd, ret;

	memset(&p, 0, sizeof(p));
	p.flags = flags;

	if (flags & IORING_SETUP_SQPOLL) {
		p.sq_thread_idle = 2000; /* 2s idle before sleep */
	}

	fd = __echo_setup(entries, &p);
	if (fd < 0) return fd;

	ring->fd         = fd;
	ring->sq_entries = p.sq_entries;
	ring->cq_entries = p.cq_entries;
	ring->flags      = flags;

	ret = __ring_map(fd, &p, ring);
	if (ret) { close(fd); ring->fd = -1; return ret; }

	ring->sqe_tail = 0;
	ring->cqe_head = 0;
	return 0;
}

int
le_ring_init(struct le_ring *ring, unsigned entries, unsigned flags)
{
	if (!ring || entries < 2 || entries > 32768) return -EINVAL;
	memset(ring, 0, sizeof(*ring));
	ring->fd = -1;
	return __ring_init(ring, entries, flags);
}

int
le_ring_init_async(struct le_ring *ring, unsigned entries,
    unsigned flags, unsigned max_tasks)
{
	int ret;
	struct le_task *tasks;

	if (max_tasks == 0 || max_tasks > entries) return -EINVAL;

	ret = le_ring_init(ring, entries, flags);
	if (ret < 0) return ret;

	tasks = calloc(max_tasks, sizeof(struct le_task));
	if (!tasks) { le_ring_destroy(ring); return -ENOMEM; }

	ring->tasks     = tasks;
	ring->max_tasks = max_tasks;
	ring->running   = 0;
	ring->event_fd  = -1;

	return 0;
}

void
le_ring_destroy(struct le_ring *ring)
{
	if (!ring) return;

	if (ring->event_fd >= 0) close(ring->event_fd);
	free(ring->tasks);

	munmap(ring->sq_ring, ring->sq_ring_sz);
	munmap(ring->cq_ring, ring->cq_ring_sz);
	munmap(ring->sqes, ring->sqes_sz);

	if (ring->fd >= 0) close(ring->fd);
	memset(ring, 0, sizeof(*ring));
	ring->fd = -1;
}

/* ── submission ─────────────────────────────────────────────────────────── */

static int
__submit(struct le_ring *ring, unsigned wait_nr)
{
	unsigned to_submit, head, tail;
	int ret, flags;

	head = __atomic_load_n(ring->sq_head, __ATOMIC_ACQUIRE);
	tail = ring->sqe_tail;
	to_submit = tail - head;

	if (!to_submit && !wait_nr) return 0;

	flags = (wait_nr || (ring->flags & IORING_SETUP_SQPOLL))
	    ? IORING_ENTER_GETEVENTS : 0;

	__atomic_store_n(ring->sq_tail, tail, __ATOMIC_RELEASE);
	le_write_barrier();

	ret = __echo_enter(ring->fd, to_submit, wait_nr, flags);
	return ret < 0 ? -errno : ret;
}

int le_submit(struct le_ring *ring) { return __submit(ring, 0); }
int le_submit_and_wait(struct le_ring *ring, unsigned wait_nr) { return __submit(ring, wait_nr); }

/* ── SQE-only flush (SQPOLL mode) ───────────────────────────────────────── */

int
le_flush(struct le_ring *ring)
{
	unsigned to_submit, head, tail;

	head = __atomic_load_n(ring->sq_head, __ATOMIC_ACQUIRE);
	tail = ring->sqe_tail;
	to_submit = tail - head;
	if (!to_submit) return 0;

	__atomic_store_n(ring->sq_tail, tail, __ATOMIC_RELEASE);
	le_write_barrier();

	if (ring->flags & IORING_SETUP_SQPOLL) {
		/* kernel polls it, just wake if needed */
		return __echo_enter(ring->fd, to_submit, 0,
		    IORING_ENTER_SQ_WAKEUP);
	}

	return __echo_enter(ring->fd, to_submit, 0, 0);
}

/* ── async ──────────────────────────────────────────────────────────────── */

int
le_submit_task(struct le_ring *ring, struct io_uring_sqe *sqe,
    le_cb cb, void *udata)
{
	/* pack cb index into user_data; index = (sqe - ring->sqes) */
	unsigned idx = (unsigned)(sqe - ring->sqes);

	ring->tasks[idx].cb    = cb;
	ring->tasks[idx].udata = udata;
	sqe->user_data = idx;

	/* submit immediately if SQPOLL, else just flush tail */
	if (ring->flags & IORING_SETUP_SQPOLL)
		return le_flush(ring);

	return le_submit(ring);
}

int
le_process(struct le_ring *ring, unsigned min_wait)
{
	struct io_uring_cqe *cqe;
	unsigned processed = 0;
	int ret;

	for (;;) {
		while (!le_peek_cqe(ring, &cqe)) {
			unsigned idx = (unsigned)cqe->user_data;
			struct le_task *t;

			if (idx < ring->max_tasks) {
				t = &ring->tasks[idx];
				if (t->cb)
					t->cb(ring, cqe, t->udata);
				t->cb = NULL;
			}
			le_cqe_seen(ring, cqe);
			processed++;
		}
		if (min_wait == 0 || processed >= min_wait) break;

		ret = __submit(ring, min_wait - processed);
		if (ret < 0) return ret;
	}
	return (int)processed;
}

/* ── chain ──────────────────────────────────────────────────────────────── */

int
le_submit_chain(struct le_ring *ring, struct io_uring_sqe **sqes,
    unsigned n, le_cb cb, void *udata)
{
	unsigned i, idx;
	int ret;

	if (!n || n > ring->sq_entries) return -EINVAL;

	idx = (unsigned)(sqes[0] - ring->sqes);
	for (i = 0; i < n; i++) {
		unsigned pos = (idx + i) % ring->sq_entries;
		/* last SQE gets the callback */
		if (i == n - 1) {
			ring->tasks[pos].cb    = cb;
			ring->tasks[pos].udata = udata;
		}
		sqes[i]->user_data = pos;
	}

	ret = le_flush(ring);
	if (ret < 0) return ret;
	return (int)n;
}

/* ── event loop ─────────────────────────────────────────────────────────── */

int
le_loop_init(struct le_ring *ring)
{
	int efd;

	if (ring->event_fd >= 0) return 0;

	efd = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC);
	if (efd < 0) return -errno;

	if (le_register_eventfd(ring, efd) < 0) {
		close(efd);
		return -errno;
	}

	ring->event_fd = efd;
	return 0;
}

int
le_loop_run(struct le_ring *ring, int epfd, int timeout_ms)
{
	struct epoll_event ev, evs[16];
	int nfds, i, ret;

	if (ring->event_fd < 0) {
		ret = le_loop_init(ring);
		if (ret < 0) return ret;
	}

	/* register eventfd with epoll */
	memset(&ev, 0, sizeof(ev));
	ev.events  = EPOLLIN;
	ev.data.fd = ring->event_fd;
	if (epoll_ctl(epfd, EPOLL_CTL_ADD, ring->event_fd, &ev) < 0)
		return -errno;

	ring->running = 1;

	while (ring->running) {
		nfds = epoll_wait(epfd, evs, 16, timeout_ms);
		if (nfds < 0) {
			if (errno == EINTR) continue;
			ret = -errno;
			goto out;
		}

		/* drain eventfd */
		if (nfds > 0) {
			uint64_t dummy;
			for (i = 0; i < nfds; i++)
				read(evs[i].data.fd, &dummy, sizeof(dummy));
		}

		/* process completions */
		ret = le_process(ring, 0);
		if (ret < 0) goto out;

		if (nfds == 0 && timeout_ms >= 0) break; /* timeout */
	}

	ret = 0;
out:
	epoll_ctl(epfd, EPOLL_CTL_DEL, ring->event_fd, NULL);
	return ret;
}

void
le_loop_stop(struct le_ring *ring)
{
	ring->running = 0;
	/* wake the epoll_wait via a dummy eventfd write */
	if (ring->event_fd >= 0) {
		uint64_t one = 1;
		write(ring->event_fd, &one, sizeof(one));
	}
}

/* ── buffer pool ────────────────────────────────────────────────────────── */

int
le_bufpool_init(struct le_ring *ring, struct le_bufpool *pool,
    unsigned nr, size_t buf_sz)
{
	unsigned i;
	int ret;

	memset(pool, 0, sizeof(*pool));
	pool->nr     = nr;
	pool->buf_sz = buf_sz;

	pool->iovecs = calloc(nr, sizeof(struct iovec));
	pool->bufs   = calloc(nr, sizeof(void *));
	pool->free_list = malloc(nr * sizeof(int));

	if (!pool->iovecs || !pool->bufs || !pool->free_list) {
		le_bufpool_destroy(ring, pool);
		return -ENOMEM;
	}

	for (i = 0; i < nr; i++) {
		pool->bufs[i] = calloc(1, buf_sz);
		if (!pool->bufs[i]) { le_bufpool_destroy(ring, pool); return -ENOMEM; }
		pool->iovecs[i].iov_base = pool->bufs[i];
		pool->iovecs[i].iov_len  = buf_sz;
		pool->free_list[i] = (int)i;
	}
	pool->free_head = nr;

	ret = le_register_buffers(ring, pool->iovecs, nr);
	if (ret < 0) { le_bufpool_destroy(ring, pool); return ret; }

	return 0;
}

void
le_bufpool_destroy(struct le_ring *ring, struct le_bufpool *pool)
{
	unsigned i;
	if (!pool) return;
	le_unregister_buffers(ring);
	for (i = 0; i < pool->nr; i++) free(pool->bufs[i]);
	free(pool->bufs);
	free(pool->iovecs);
	free(pool->free_list);
	memset(pool, 0, sizeof(*pool));
}

int
le_bufpool_get(struct le_bufpool *pool)
{
	if (pool->free_head == 0) return -1;
	return pool->free_list[--pool->free_head];
}

void
le_bufpool_put(struct le_bufpool *pool, int idx)
{
	pool->free_list[pool->free_head++] = idx;
}

/* ── completion ─────────────────────────────────────────────────────────── */

int
le_wait_cqe(struct le_ring *ring, struct io_uring_cqe **cqe_out)
{
	struct io_uring_cqe *cqe;
	int ret;

	while (le_peek_cqe(ring, &cqe)) {
		ret = __submit(ring, 1);
		if (ret < 0) return ret;
	}
	*cqe_out = cqe;
	return 0;
}

int
le_wait_cqes(struct le_ring *ring, struct io_uring_cqe **cqes, unsigned count)
{
	unsigned head, completed, i, mask;
	int ret;

	while (1) {
		le_read_barrier();
		head = __atomic_load_n(ring->cq_head, __ATOMIC_ACQUIRE);
		completed = __atomic_load_n(ring->cq_tail, __ATOMIC_ACQUIRE) - head;

		if (completed >= count) {
			mask = *ring->cq_mask;
			for (i = 0; i < count; i++)
				cqes[i] = &ring->cqes[(head + i) & mask];
			return (int)count;
		}
		ret = __submit(ring, count - completed);
		if (ret < 0) return ret;
	}
}

int
le_wait_cqe_timeout(struct le_ring *ring, struct __kernel_timespec *ts,
    struct io_uring_cqe **cqe_out)
{
	struct io_uring_sqe *timeout_sqe;
	struct io_uring_cqe *cqe;
	unsigned head, tail, mask, i;
	unsigned timeout_idx;

	/* fast path: already have a CQE */
	if (!le_peek_cqe(ring, cqe_out))
		return 0;

	/* arm a timeout SQE */
	timeout_sqe = le_get_sqe(ring);
	if (!timeout_sqe) return -EBUSY;

	le_prep_timeout(timeout_sqe, ts, 1, 0);
	timeout_idx = (unsigned)(timeout_sqe - ring->sqes);
	timeout_sqe->user_data = timeout_idx;

	/* submit and wait for at least one completion */
	le_submit_and_wait(ring, 1);

	/* scan all CQEs — return first non-timeout, or -ETIME */
	le_read_barrier();
	head = __atomic_load_n(ring->cq_head, __ATOMIC_ACQUIRE);
	tail = __atomic_load_n(ring->cq_tail, __ATOMIC_ACQUIRE);
	mask = *ring->cq_mask;

	for (i = head; i < tail; i++) {
		cqe = &ring->cqes[i & mask];
		if (cqe->user_data == timeout_idx)
			continue;
		/* found a real CQE */
		__atomic_store_n(ring->cq_head, i + 1, __ATOMIC_RELEASE);
		*cqe_out = cqe;
		return 0;
	}

	/* no real CQE, timeout must have fired */
	__atomic_store_n(ring->cq_head, tail, __ATOMIC_RELEASE);
	return -ETIME;
}

int
le_cancel(struct le_ring *ring, void *user_data)
{
	struct io_uring_sqe *sqe = le_get_sqe(ring);
	if (!sqe) return -EBUSY;
	le_prep_cancel(sqe, user_data, 0);
	return __submit(ring, 1);
}

/* ── registration ───────────────────────────────────────────────────────── */

int le_register_buffers(struct le_ring *ring, const struct iovec *iovecs,
    unsigned nr_iovecs)
{ int r = __echo_register(ring->fd, IORING_REGISTER_BUFFERS, iovecs, nr_iovecs);
  return r < 0 ? -errno : 0; }

int le_unregister_buffers(struct le_ring *ring)
{ int r = __echo_register(ring->fd, IORING_UNREGISTER_BUFFERS, NULL, 0);
  return r < 0 ? -errno : 0; }

int le_register_files(struct le_ring *ring, const int *files, unsigned nr_files)
{ int r = __echo_register(ring->fd, IORING_REGISTER_FILES, files, nr_files);
  return r < 0 ? -errno : 0; }

int le_unregister_files(struct le_ring *ring)
{ int r = __echo_register(ring->fd, IORING_UNREGISTER_FILES, NULL, 0);
  return r < 0 ? -errno : 0; }

int le_register_eventfd(struct le_ring *ring, int event_fd)
{ int r = __echo_register(ring->fd, IORING_REGISTER_EVENTFD, &event_fd, 1);
  return r < 0 ? -errno : 0; }

int le_unregister_eventfd(struct le_ring *ring)
{ int r = __echo_register(ring->fd, IORING_UNREGISTER_EVENTFD, NULL, 0);
  return r < 0 ? -errno : 0; }

/* ── sync wrappers ──────────────────────────────────────────────────────── */

static int
__echo_one(struct le_ring *ring)
{
	struct io_uring_cqe *cqe;
	int ret;

	ret = le_submit_and_wait(ring, 1);
	if (ret < 0) return ret;
	ret = le_wait_cqe(ring, &cqe);
	if (ret < 0) return ret;
	ret = (int)cqe->res;
	le_cqe_seen(ring, cqe);
	return ret;
}

ssize_t le_read(struct le_ring *ring, int fd, void *buf, size_t count, __u64 off)
{ struct io_uring_sqe *sqe = le_get_sqe(ring); if (!sqe) return -EBUSY;
  le_prep_read(sqe, fd, buf, (unsigned)count, off); return __echo_one(ring); }

ssize_t le_write(struct le_ring *ring, int fd, const void *buf, size_t count, __u64 off)
{ struct io_uring_sqe *sqe = le_get_sqe(ring); if (!sqe) return -EBUSY;
  le_prep_write(sqe, fd, buf, (unsigned)count, off); return __echo_one(ring); }

ssize_t le_send(struct le_ring *ring, int fd, const void *buf, size_t len, int flags)
{ struct io_uring_sqe *sqe = le_get_sqe(ring); if (!sqe) return -EBUSY;
  le_prep_send(sqe, fd, buf, len, flags); return __echo_one(ring); }

ssize_t le_recv(struct le_ring *ring, int fd, void *buf, size_t len, int flags)
{ struct io_uring_sqe *sqe = le_get_sqe(ring); if (!sqe) return -EBUSY;
  le_prep_recv(sqe, fd, buf, len, flags); return __echo_one(ring); }

int le_accept(struct le_ring *ring, int fd, struct sockaddr *addr,
    socklen_t *addrlen, int flags)
{ struct io_uring_sqe *sqe = le_get_sqe(ring); if (!sqe) return -EBUSY;
  le_prep_accept(sqe, fd, addr, addrlen, flags); return __echo_one(ring); }

int le_connect(struct le_ring *ring, int fd, const struct sockaddr *addr,
    socklen_t addrlen)
{ struct io_uring_sqe *sqe = le_get_sqe(ring); if (!sqe) return -EBUSY;
  le_prep_connect(sqe, fd, addr, addrlen); return __echo_one(ring); }

int le_close(struct le_ring *ring, int fd)
{ struct io_uring_sqe *sqe = le_get_sqe(ring); if (!sqe) return -EBUSY;
  le_prep_close(sqe, fd); return __echo_one(ring); }

int le_fsync(struct le_ring *ring, int fd)
{ struct io_uring_sqe *sqe = le_get_sqe(ring); if (!sqe) return -EBUSY;
  le_prep_fsync(sqe, fd, 0); return __echo_one(ring); }

int le_nop(struct le_ring *ring)
{ struct io_uring_sqe *sqe = le_get_sqe(ring); if (!sqe) return -EBUSY;
  le_prep_nop(sqe); return __echo_one(ring); }

/* ── utility ────────────────────────────────────────────────────────────── */

const char *le_strerror(int err) { if (err < 0) err = -err; return strerror(err); }

const char *le_version(void) { return "0.1.0"; }