examples

io_uring.cpp (5812B)

---

 #include "coroutine_owner.hpp"
 #include "fd.hpp"
 
 #include <liburing.h>
 #include <sys/timerfd.h>
 
 #include <cstring>
 #include <print>
 #include <span>
 #include <map>
 #include <memory>
 #include <cassert>
 
 struct task
 {
 	bool done() const noexcept
 	{
 		return coro->done();
 	}
 
 	struct promise_type
 	{
 		task get_return_object(void) noexcept
 		{
 			return task{make_coroutine_owner(*this)};
 		}
 		auto initial_suspend(void) const noexcept
 		{
 			return std::suspend_never{};
 		}
 		auto final_suspend(void) const noexcept
 		{
 			return std::suspend_always{};
 		}
 		void unhandled_exception(void) const noexcept
 		{
 			std::terminate();
 		}
 		void return_void(void) const noexcept
 		{}
 	};
 
 	coroutine_owner<promise_type> coro;
 };
 
 struct uring
 {
 	struct awaitable;
 private:
 	struct rendezvous;
 public:
 	uring(unsigned entries, unsigned flags)
 	{
 		if (io_uring_queue_init(entries, &ring, flags))
 			throw std::system_error(errno, std::system_category());
 	}
 
 	~uring(void)
 	{
 		assert(inflight.empty());
 		io_uring_queue_exit(&ring);
 	}
 
 	uring(const uring&) = delete;
 	uring& operator =(const uring&) = delete;
 
 	uint64_t wait_and_complete(void);
 	bool process_completion(void);
 	awaitable read(int, std::span<char>, off_t);
 	awaitable write(int, std::span<char const>, off_t);
 
 private:
 	std::pair<io_uring_sqe*, awaitable> get_sqe(void)
 	{
 		auto sqe = io_uring_get_sqe(&ring);
 		assert(sqe);
 		io_uring_sqe_set_data64(sqe, ++last_id);
 		return {sqe, {*this , std::make_shared<rendezvous>(), last_id}};
 	}
 
 	void submit(std::shared_ptr<rendezvous> r)
 	{
 		int n = io_uring_submit(&ring);
 		if (n < 0)
 			throw std::system_error(-n, std::system_category());
 		assert(n == 1);
 		assert(not inflight.contains(last_id));
 		inflight[last_id] = std::move(r);
 	}
 
 	struct rendezvous
 	{
 		std::optional<int> result = std::nullopt;
 		std::coroutine_handle<> continuation = {};
 	};
 
 public:
 	struct awaitable
 	{
 		awaitable(uring& r, std::shared_ptr<rendezvous> p, uint64_t i) noexcept
 		: ring{r}
 		, rp{std::move(p)}
 		, id{i}
 		{}
 		~awaitable() noexcept
 		{
 			if (!rp || await_ready())
 				return;
 
 			rp->continuation = {};
 
 			io_uring_sync_cancel_reg reg = {
 				.addr = id,
 				.fd = 0,
 				.flags = 0,
 				.timeout = {-1, -1},
 				.opcode = 0,
 				.pad = {},
 				.pad2 = {},
 			};
 			int result = io_uring_register_sync_cancel(&ring.ring, &reg);
 			switch (result)
 			{
 			case 0: // Job was cancelled
 				ring.inflight.erase(id);
 				break;
 			case -ENOENT: // No job was found: result in completion queue. Continuation already removed, so nothing to do.
 				break;
 			case -EINVAL: // bad args
 				assert(false);
 				break;
 			case -EALREADY: // not sure what to do here, can't cancel, can't finish destructor until complete?
 				while (ring.wait_and_complete() != id)
 				{}
 				break;
 			default:
 				std::println("cancel: {}", result);
 				assert(false);
 			}
 		}
 		awaitable(awaitable const&) = delete;
 		awaitable& operator =(awaitable const&) = delete;
 		awaitable(awaitable&& old)
 		: ring{old.ring}
 		, rp{std::move(old.rp)}
 		, id{old.id}
 		{
 			//std::println("Awaitable {} moved", id);
 		}
 		bool await_ready(void) const noexcept
 		{
 			return rp->result.has_value();
 		}
 		void await_suspend(std::coroutine_handle<> h) const noexcept
 		{
 			rp->continuation = h;
 		}
 		int await_resume() const
 		{
 			//std::println("Job {} resumed", id);
 			return rp->result.value();
 		}
 //	private:
 		uring& ring;
 		std::shared_ptr<rendezvous> rp;
 		uint64_t id;
 	};
 
 private:
 	io_uring ring;
 	std::map<uint64_t, std::shared_ptr<rendezvous>> inflight;
 	uint64_t last_id = 0;
 };
 
 
 uint64_t uring::wait_and_complete(void)
 {
 	assert(not inflight.empty());
 
 	io_uring_cqe* cqe;
 	if (io_uring_wait_cqe(&ring, &cqe))
 		throw std::system_error(errno, std::system_category());
 
 	uint64_t id = io_uring_cqe_get_data64(cqe);
 	auto it = inflight.find(id);
 	assert(it != inflight.end());
 	auto s = std::move(it->second);
 	inflight.erase(it);
 
 	assert(not s->result);
 	s->result = cqe->res;
 	io_uring_cq_advance(&ring, 1);
 
 	if (s->continuation)
 		s->continuation();
 
 	return id;
 }
 
 bool uring::process_completion(void)
 {
 	if (inflight.empty())
 		return false;
 
 	wait_and_complete();
 	return true;
 }
 
 uring::awaitable uring::read(int fd, std::span<char> buf, off_t offset = 0)
 {
 	auto [sqe, sub] = get_sqe();
 	io_uring_prep_read(sqe, fd, buf.data(), buf.size(), offset);
 	submit(sub.rp);
 	return sub;
 }
 
 uring::awaitable uring::write(int fd, std::span<char const> buf, off_t offset = 0)
 {
 	auto [sqe, sub] = get_sqe();
 	io_uring_prep_write(sqe, fd, buf.data(), buf.size(), offset);
 	submit(sub.rp);
 	return sub;
 }
 
 task read_routine(uring& u)
 {
 	char b[128];
 	for (;;)
 	{
 		std::span<char> buf(b);
 		int n = co_await u.read(0, buf);
 		if (n == 0)
 			co_return;
 		if (n < 0)
 			throw std::system_error(-n, std::system_category());
 
 		buf = buf.subspan(0, n);
 
 		write(1, buf.data(), buf.size());
 		continue;
 
 		while (!buf.empty())
 		{
 			n = /*co_await u.*/write(1, buf.data(), buf.size());
 			if (n <= 0)
 				throw std::system_error(-n, std::system_category());
 
 			buf = buf.subspan(n);
 		}
 	}
 }
 
 task timer_routine(uring& u)
 {
 	auto fd = file_descriptor{timerfd_create(CLOCK_MONOTONIC, 0)};
 	assert(fd);
 
 	itimerspec ts
 	{
 		{2, 0},
 		{2, 0},
 	};
 
 	timerfd_settime(fd, 0, &ts, nullptr);
 
 	for (uint64_t i = 0;;)
 	{
 		uint64_t timeouts;
 		auto s = std::span{reinterpret_cast<char*>(&timeouts), sizeof timeouts};
 		assert(co_await u.read(fd, s) == sizeof timeouts);
 
 		i += timeouts;
 
 		auto x = std::format("Timeout {}\n", i);
 		int n = co_await u.write(1, std::span<char>{x.data(), x.size()});
 		if (n <= 0)
 			throw std::system_error(-n, std::system_category());
 	}
 }
 
 int main()
 {
 	uring r(128, 0);
 
 	task rr = read_routine(r);
 	task tr = timer_routine(r);
 
 	while (not rr.done())
 		r.process_completion();
 
 	return 0;
 }