Simple io_uring echo server example

liburing_server_echo.c

#include "liburing.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <unistd.h>

#define QUEUE_DEPTH 32
#define READ_SZ     4096

/* Define event types */
enum {
	EVENT_TYPE_ACCEPT = 0,
	EVENT_TYPE_READ = 1,
	EVENT_TYPE_WRITE = 2
};

struct request {
	int socket;
	int event_type;
	int iovec_count;
	struct iovec iov[];
};

// Global variable to hold ring
struct io_uring ring;

void sigint_handler(int signo)
{
	(void)signo;
	io_uring_queue_exit(&ring);
	exit(EXIT_SUCCESS);
}

void *xmalloc(size_t size)
{
	void *ret = malloc(size);
	if (ret == NULL) {
		perror("malloc");
		exit(EXIT_FAILURE);
	}
	return ret;
}

void add_accept_request(int server_socket, struct sockaddr_in *client_addr, socklen_t *client_addr_len)
{
	struct io_uring_sqe *sqe;
	struct request *req;

	// Get pointer to submission queue
	sqe = io_uring_get_sqe(&ring);

	// Prepare accept request
	req = xmalloc(sizeof(*req));
	io_uring_prep_accept(sqe, server_socket, (struct sockaddr *)client_addr, client_addr_len, 0);
	req->event_type = EVENT_TYPE_ACCEPT;
	io_uring_sqe_set_data(sqe, req);

	// Add accept request to the submission queue
	io_uring_submit(&ring);
}

/* Linux kernel 5.5 has support for readv, but not for recv() or read() */
void add_read_request(int socket)
{
	struct io_uring_sqe *sqe;
	struct request *req;

	// Get pointer to submission queue
	sqe = io_uring_get_sqe(&ring);

	// Prepare read request
	req = xmalloc(sizeof(*req) + sizeof(struct iovec));
	req->iov[0].iov_base = calloc(1, READ_SZ);
	req->iov[0].iov_len = READ_SZ;
	req->event_type = EVENT_TYPE_READ;
	req->socket = socket;
	io_uring_prep_readv(sqe, socket, &req->iov[0], 1, 0);
	io_uring_sqe_set_data(sqe, req);

	// Add accept request to the submission queue
	io_uring_submit(&ring);
}

void add_write_request(struct request *req)
{
	struct io_uring_sqe *sqe;

	// Get pointer to submission queue
	sqe = io_uring_get_sqe(&ring);

	// Prepare read request
	req->event_type = EVENT_TYPE_WRITE;
	io_uring_prep_writev(sqe, req->socket, req->iov, req->iovec_count, 0);
	io_uring_sqe_set_data(sqe, req);

	// Add accept request to the submission queue
	io_uring_submit(&ring);
}

int main(int argc, char *argv[])
{
	int sd;
	struct sockaddr_in srv_addr;
	int on = 1;
	struct io_uring_cqe *cqe;
	struct sockaddr_in client_addr;
	socklen_t client_addr_len;
	in_port_t port;

	if (argc < 2) {
		fputs("Usage: liburing_server_echo port", stderr);
		exit(EXIT_FAILURE);
	}

	port = atoi(argv[1]);

	sd = socket(AF_INET, SOCK_STREAM, 0);
	if (sd < 0) {
		perror("socket");
		exit(EXIT_FAILURE);
	}

	if (setsockopt(sd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(int)) < 0) {
		perror("setsockopt");
		exit(EXIT_FAILURE);
	}

	memset((void *)&srv_addr, 0, sizeof(srv_addr));
	srv_addr.sin_family = AF_INET;
	srv_addr.sin_port = htons(port);
	srv_addr.sin_addr.s_addr = htonl(INADDR_ANY);

	if (bind(sd, (const struct sockaddr *)&srv_addr, sizeof(srv_addr)) < 0) {
		perror("bind");
		exit(EXIT_FAILURE);
	}

	if (listen(sd, SOMAXCONN) < 0) {
		perror("listen");
		exit(EXIT_FAILURE);
	}

	if (io_uring_queue_init(QUEUE_DEPTH, &ring, 0) != 0) {
		perror("io_uring_queue_init");
		exit(EXIT_FAILURE);
	}

	signal(SIGINT, sigint_handler);

	client_addr_len = sizeof(client_addr);
	add_accept_request(sd, &client_addr, &client_addr_len);

	while (1) {
		int ret;
		struct request *req;

		// Wait for an I/O event
		ret = io_uring_wait_cqe(&ring, &cqe);
		if (ret < 0) {
			perror("io_uring_wait_cqe");
			exit(EXIT_FAILURE);
		}

		req = (struct request *)cqe->user_data;
		if (cqe->res < 0) {
			fprintf(stderr, "Async request failed: %s for event: %d\n", strerror(-cqe->res), req->event_type);
			exit(EXIT_FAILURE);
		}

		switch (req->event_type) {
			case EVENT_TYPE_ACCEPT:
				// Re-add the socket for listening
				add_accept_request(sd, &client_addr, &client_addr_len);

				// Add the client socket for reading
				add_read_request(cqe->res);

				// Free the request
				free(req);

				break;
			case EVENT_TYPE_READ:
				// Skip empty requests
				if (cqe->res == 0) {
					fprintf(stderr, "Empty request!\n");
					break;
				}

				// End of input
				if (cqe->res < 0) {
					// Close client socket
					close(req->socket);
					
					// Free the buffer
					free(req->iov[0].iov_base);

					// Free the request
					free(req);
				} else {
					/* When we are here, we need to handle
					 * the request. We do that by simply
					 * re-writing it as is back to the
					 * socket. */
					add_write_request(req);
				}

				break;
			case EVENT_TYPE_WRITE:
				/* When we are here, the write request has
				 * finished and we have to clean up the used
				 * buffers and struct request metadata */

				// Free all buffers (theoretically we have only one)
				for (int i = 0; i < req->iovec_count; i++) {
					free(req->iov[i].iov_base);
				}

				// Free the request
				free(req);

				break;
		}

		/* Mark this request as processed and move on */
		io_uring_cqe_seen(&ring, cqe);
	}

	return 0;
}