util.poll: Import poll library with epoll and select support

2 files changed, 425 insertions, 1 deletions

diff --git a/util-src/GNUmakefile b/util-src/GNUmakefile
index dfe3950f..f2f9b1bc 100644
--- a/util-src/GNUmakefile
+++ b/util-src/GNUmakefile
@@ -6,7 +6,7 @@ CFLAGS+=-I$(LUA_INCDIR)
 INSTALL_DATA=install -m644
 TARGET?=../util/
 
-ALL=encodings.so hashes.so net.so pposix.so signal.so table.so ringbuffer.so time.so
+ALL=encodings.so hashes.so net.so pposix.so signal.so table.so ringbuffer.so time.so poll.so
 
 ifdef RANDOM
 ALL+=crand.so
diff --git a/util-src/poll.c b/util-src/poll.c
new file mode 100644
index 00000000..d0894c54
--- /dev/null
+++ b/util-src/poll.c
@@ -0,0 +1,424 @@
+
+/*
+ * Lua polling library
+ * Copyright (C) 2017-2018 Kim Alvefur
+ *
+ * This project is MIT licensed. Please see the
+ * COPYING file in the source package for more information.
+ *
+ */
+
+#include <unistd.h>
+#include <string.h>
+#include <errno.h>
+
+#ifdef __linux__
+#define USE_EPOLL
+#endif
+
+#ifdef USE_EPOLL
+#include <sys/epoll.h>
+#ifndef MAX_EVENTS
+#define MAX_EVENTS 10
+#endif
+#else
+#include <sys/select.h>
+#endif
+
+#include <lualib.h>
+#include <lauxlib.h>
+
+#ifdef USE_EPOLL
+#define STATE_MT "util.poll.epoll"
+#else
+#define STATE_MT "util.poll.select"
+#endif
+
+/*
+ * Structure to keep state for each type of API
+ */
+typedef struct Lpoll_state {
+	int processed;
+#ifdef USE_EPOLL
+	int epoll_fd;
+	struct epoll_event events[MAX_EVENTS];
+#else
+	fd_set wantread;
+	fd_set wantwrite;
+	fd_set readable;
+	fd_set writable;
+	fd_set all;
+	fd_set err;
+#endif
+} Lpoll_state;
+
+/*
+ * Add an FD to be watched
+ */
+int Ladd(lua_State *L) {
+	struct Lpoll_state *state = luaL_checkudata(L, 1, STATE_MT);
+	int fd = luaL_checkinteger(L, 2);
+
+	int wantread = lua_toboolean(L, 3);
+	int wantwrite = lua_toboolean(L, 4);
+
+#ifdef USE_EPOLL
+	struct epoll_event event;
+	event.data.fd = fd;
+	event.events = (wantread ? EPOLLIN : 0) | (wantwrite ? EPOLLOUT : 0);
+
+	event.events |= EPOLLERR | EPOLLHUP | EPOLLRDHUP;
+
+	int ret = epoll_ctl(state->epoll_fd, EPOLL_CTL_ADD, fd, &event);
+
+	if(ret < 0) {
+		ret = errno;
+		lua_pushnil(L);
+		lua_pushstring(L, strerror(ret));
+		lua_pushinteger(L, ret);
+		return 3;
+	}
+
+	lua_pushboolean(L, 1);
+	return 1;
+
+#else
+
+	if(FD_ISSET(fd, &state->all)) {
+		lua_pushnil(L);
+		lua_pushstring(L, strerror(EEXIST));
+		lua_pushinteger(L, EEXIST);
+	}
+
+	FD_CLR(fd, &state->readable);
+	FD_CLR(fd, &state->writable);
+	FD_CLR(fd, &state->err);
+
+	FD_SET(fd, &state->all);
+
+	if(wantread) {
+		FD_SET(fd, &state->wantread);
+	}
+	else {
+		FD_CLR(fd, &state->wantread);
+	}
+
+	if(wantwrite) {
+		FD_SET(fd, &state->wantwrite);
+	}
+	else {
+		FD_CLR(fd, &state->wantwrite);
+	}
+
+	lua_pushboolean(L, 1);
+	return 1;
+#endif
+}
+
+/*
+ * Set events to watch for, readable and/or writable
+ */
+int Lset(lua_State *L) {
+	struct Lpoll_state *state = luaL_checkudata(L, 1, STATE_MT);
+	int fd = luaL_checkinteger(L, 2);
+
+#ifdef USE_EPOLL
+
+	int wantread = lua_toboolean(L, 3);
+	int wantwrite = lua_toboolean(L, 4);
+
+	struct epoll_event event;
+	event.data.fd = fd;
+	event.events = (wantread ? EPOLLIN : 0) | (wantwrite ? EPOLLOUT : 0);
+
+	event.events |= EPOLLERR | EPOLLHUP | EPOLLRDHUP;
+
+	int ret = epoll_ctl(state->epoll_fd, EPOLL_CTL_MOD, fd, &event);
+
+	if(ret == 0) {
+		lua_pushboolean(L, 1);
+		return 1;
+	}
+	else {
+		ret = errno;
+		lua_pushnil(L);
+		lua_pushstring(L, strerror(ret));
+		lua_pushinteger(L, ret);
+		return 3;
+	}
+
+#else
+
+	if(!FD_ISSET(fd, &state->all)) {
+		lua_pushnil(L);
+		lua_pushstring(L, strerror(ENOENT));
+		lua_pushinteger(L, ENOENT);
+	}
+
+	if(!lua_isnoneornil(L, 3)) {
+		if(lua_toboolean(L, 3)) {
+			FD_SET(fd, &state->wantread);
+		}
+		else {
+			FD_CLR(fd, &state->wantread);
+		}
+	}
+
+	if(!lua_isnoneornil(L, 4)) {
+		if(lua_toboolean(L, 4)) {
+			FD_SET(fd, &state->wantwrite);
+		}
+		else {
+			FD_CLR(fd, &state->wantwrite);
+		}
+	}
+
+	lua_pushboolean(L, 1);
+	return 1;
+#endif
+}
+
+/*
+ * Remove FDs
+ */
+int Ldel(lua_State *L) {
+	struct Lpoll_state *state = luaL_checkudata(L, 1, STATE_MT);
+	int fd = luaL_checkinteger(L, 2);
+
+#ifdef USE_EPOLL
+
+	struct epoll_event event;
+	event.data.fd = fd;
+
+	int ret = epoll_ctl(state->epoll_fd, EPOLL_CTL_DEL, fd, &event);
+
+	if(ret == 0) {
+		lua_pushboolean(L, 1);
+		return 1;
+	}
+	else {
+		ret = errno;
+		lua_pushnil(L);
+		lua_pushstring(L, strerror(ret));
+		lua_pushinteger(L, ret);
+		return 3;
+	}
+
+#else
+
+	if(!FD_ISSET(fd, &state->all)) {
+		lua_pushnil(L);
+		lua_pushstring(L, strerror(ENOENT));
+		lua_pushinteger(L, ENOENT);
+	}
+
+	FD_CLR(fd, &state->wantread);
+	FD_CLR(fd, &state->wantwrite);
+	FD_CLR(fd, &state->readable);
+	FD_CLR(fd, &state->writable);
+	FD_CLR(fd, &state->all);
+	FD_CLR(fd, &state->err);
+
+	lua_pushboolean(L, 1);
+	return 1;
+#endif
+}
+
+
+/*
+ * Check previously manipulated event state for FDs ready for reading or writing
+ */
+inline int Lpushevent(lua_State *L, struct Lpoll_state *state) {
+#ifdef USE_EPOLL
+
+	if(state->processed > 0) {
+		state->processed--;
+		struct epoll_event event = state->events[state->processed];
+		lua_pushinteger(L, event.data.fd);
+		lua_pushboolean(L, event.events & (EPOLLIN | EPOLLHUP | EPOLLRDHUP | EPOLLERR));
+		lua_pushboolean(L, event.events & EPOLLOUT);
+		return 3;
+	}
+
+#else
+
+	for(int fd = state->processed + 1; fd < FD_SETSIZE; fd++) {
+		if(FD_ISSET(fd, &state->readable) || FD_ISSET(fd, &state->writable) || FD_ISSET(fd, &state->err)) {
+			lua_pushinteger(L, fd);
+			lua_pushboolean(L, FD_ISSET(fd, &state->readable) | FD_ISSET(fd, &state->err));
+			lua_pushboolean(L, FD_ISSET(fd, &state->writable));
+			FD_CLR(fd, &state->readable);
+			FD_CLR(fd, &state->writable);
+			FD_CLR(fd, &state->err);
+			state->processed = fd;
+			return 3;
+		}
+	}
+
+#endif
+	return 0;
+}
+
+/*
+ * Wait for event
+ */
+int Lwait(lua_State *L) {
+	struct Lpoll_state *state = luaL_checkudata(L, 1, STATE_MT);
+
+	int ret = Lpushevent(L, state);
+
+	if(ret != 0) {
+		return ret;
+	}
+
+	lua_Number timeout = luaL_checknumber(L, 2);
+	luaL_argcheck(L, timeout >= 0, 1, "positive number expected");
+
+#ifdef USE_EPOLL
+	ret = epoll_wait(state->epoll_fd, state->events, MAX_EVENTS, timeout * 1000);
+#else
+	/*
+	 * select(2) mutates the fd_sets passed to it so in order to not
+	 * have to recreate it manually every time a copy is made.
+	 */
+	memcpy(&state->readable, &state->wantread, sizeof(fd_set));
+	memcpy(&state->writable, &state->wantwrite, sizeof(fd_set));
+	memcpy(&state->err, &state->all, sizeof(fd_set));
+
+	struct timeval tv;
+	tv.tv_sec = (time_t)timeout;
+	tv.tv_usec = ((suseconds_t)(timeout * 1000000)) % 1000000;
+
+	ret = select(FD_SETSIZE, &state->readable, &state->writable, NULL, &tv);
+#endif
+
+	if(ret == 0) {
+		lua_pushnil(L);
+		lua_pushstring(L, "timeout");
+		return 2;
+	}
+	else if(ret < 0) {
+		ret = errno;
+		lua_pushnil(L);
+		lua_pushstring(L, strerror(ret));
+		lua_pushinteger(L, ret);
+		return 3;
+	}
+
+	/*
+	 * Search for the first ready FD and return it
+	 */
+#ifdef USE_EPOLL
+	state->processed = ret;
+#else
+	state->processed = -1;
+#endif
+	return Lpushevent(L, state);
+}
+
+#ifdef USE_EPOLL
+/*
+ * Return Epoll FD
+ */
+int Lgetfd(lua_State *L) {
+	struct Lpoll_state *state = luaL_checkudata(L, 1, STATE_MT);
+	lua_pushinteger(L, state->epoll_fd);
+	return 1;
+}
+
+/*
+ * Close epoll FD
+ */
+int Lgc(lua_State *L) {
+	struct Lpoll_state *state = luaL_checkudata(L, 1, STATE_MT);
+
+	if(close(state->epoll_fd) == 0) {
+		state->epoll_fd = -1;
+	}
+	else {
+		lua_pushstring(L, strerror(errno));
+		lua_error(L);
+	}
+
+	return 0;
+}
+#endif
+
+/*
+ * String representation
+ */
+int Ltos(lua_State *L) {
+	struct Lpoll_state *state = luaL_checkudata(L, 1, STATE_MT);
+	lua_pushfstring(L, "%s: %p", STATE_MT, state);
+	return 1;
+}
+
+/*
+ * Create a new context
+ */
+int Lnew(lua_State *L) {
+	/* Allocate state */
+	Lpoll_state *state = lua_newuserdata(L, sizeof(Lpoll_state));
+	luaL_setmetatable(L, STATE_MT);
+
+	/* Initialize state */
+#ifdef USE_EPOLL
+	state->epoll_fd = epoll_create1(0);
+	state->processed = 0;
+#else
+	FD_ZERO(&state->wantread);
+	FD_ZERO(&state->wantwrite);
+	FD_ZERO(&state->readable);
+	FD_ZERO(&state->writable);
+	state->processed = FD_SETSIZE;
+#endif
+
+	return 1;
+}
+
+/*
+ * Open library
+ */
+int luaopen_util_poll(lua_State *L) {
+	luaL_checkversion(L);
+
+	luaL_newmetatable(L, STATE_MT);
+	{
+
+		lua_pushliteral(L, STATE_MT);
+		lua_setfield(L, -2, "__name");
+
+		lua_pushcfunction(L, Ltos);
+		lua_setfield(L, -2, "__tostring");
+
+		lua_createtable(L, 0, 2);
+		{
+			lua_pushcfunction(L, Ladd);
+			lua_setfield(L, -2, "add");
+			lua_pushcfunction(L, Lset);
+			lua_setfield(L, -2, "set");
+			lua_pushcfunction(L, Ldel);
+			lua_setfield(L, -2, "del");
+			lua_pushcfunction(L, Lwait);
+			lua_setfield(L, -2, "wait");
+#ifdef USE_EPOLL
+			lua_pushcfunction(L, Lgetfd);
+			lua_setfield(L, -2, "getfd");
+#endif
+		}
+		lua_setfield(L, -2, "__index");
+
+#ifdef USE_EPOLL
+		lua_pushcfunction(L, Lgc);
+		lua_setfield(L, -2, "__gc");
+#endif
+	}
+
+	lua_createtable(L, 0, 1);
+	{
+		lua_pushcfunction(L, Lnew);
+		lua_setfield(L, -2, "new");
+	}
+	return 1;
+}
+