From b12592a5d5e562781216e7e3e33d0baca47042f1 Mon Sep 17 00:00:00 2001
From: Aleteoryx <alyx@aleteoryx.me>
Date: Mon, 2 Sep 2024 23:21:50 -0400
Subject: [PATCH] configurable connection cap

---
 README.md    |  6 ++++--
 nex.nelua    |  2 +-
 nps.nelua    |  2 +-
 socket.nelua | 10 +++++-----
 4 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 712ee95..dcf25bc 100644
--- a/README.md
+++ b/README.md
@@ -33,19 +33,21 @@ protocols, so long as they are not reliant on multiple sockets.
 
 it exposes the following methods from its return value:
 
-### `socket.listen_tcp(addr: string, handler: function(): void)`
+### `socket.listen_tcp(addr: string, conn_cap: uinteger, handler: function(): void)`
 
 listens on the host/port in `addr`. panics if it can't connect or if
 `addr` is malformed. `addr` is either of the format `IPV4:PORT` or
 `{IPV6}:PORT`. passes the connection to `socket.listen_sock`.
 
-### `socket.listen_sock(fd: cint, handler: function(): void)`
+### `socket.listen_sock(fd: cint, conn_cap: uinteger, handler: function(): void)`
 
 `fd` must be a valid socket, bound to some address and ready to
 `accept` incoming connections. as connections are made, `handler` will
 be executed in the context of a coroutine. it is expected to call one
 of the provided [yielding functions](#yielding-functions) eventually.
 
+will begin refusing incoming requests at `conn_cap` open connections
+
 if there is any error reading from or writing to the socket, the
 handler coroutine is deleted and the connection is dropped.
 
diff --git a/nex.nelua b/nex.nelua
index 0f01dc1..7399973 100644
--- a/nex.nelua
+++ b/nex.nelua
@@ -8,7 +8,7 @@ if #arg ~= 1 then
   print("  `{<ipv6 addr>}:<port>'")
 end
 
-socket.listen_tcp(arg[1], function()
+socket.listen_tcp(arg[1], 1024, function()
   local line = socket.recv_line():gsub("\r", "")
 
   if line:subview(1, 1) == "/" then line = line:subview(2) end
diff --git a/nps.nelua b/nps.nelua
index 114250c..ecc4221 100644
--- a/nps.nelua
+++ b/nps.nelua
@@ -10,7 +10,7 @@ if #arg ~= 1 then
   print("  `{<ipv6 addr>}:<port>'")
 end
 
-socket.listen_tcp(arg[1], function()
+socket.listen_tcp(arg[1], 1024, function()
   local target = socket.recv_line():gsub("\r", "")
 
   local lines: vector(string)
diff --git a/socket.nelua b/socket.nelua
index 2459593..266cb41 100644
--- a/socket.nelua
+++ b/socket.nelua
@@ -192,7 +192,7 @@ function handler_state:step(epfd: cint, key: uint32): (boolean)
   return true
 end
 
-function export.listen_sock(sock: cint, handler: function(): void): void
+function export.listen_sock(sock: cint, conn_cap: uinteger, handler: function(): void): void
   die_errno(c_listen(sock, 32) ~= 0, "couldn't listen on TCP socket")
   local epfd = c_epoll_create(0)
   die_errno(epfd == -1, "couldn't create epoll instance")
@@ -214,7 +214,7 @@ function export.listen_sock(sock: cint, handler: function(): void): void
     for i = 0, < event_count do
       if events[i].data.u32 == 0 then
         local fd = c_accept(sock, nilptr, nilptr)
-        if #handlers >= 1024 then -- drop connections if allocation fails
+        if #handlers >= conn_cap then -- drop connections if allocation fails
           local errmsg = "The server is overloaded, please try again later.\n"
           c_send(fd, errmsg.data, errmsg.size, 0)
           c_close(fd)
@@ -254,7 +254,7 @@ end
 local inet_re = '^(%d?%d?%d%.%d?%d?%d%.%d?%d?%d%.%d?%d?%d):(%d?%d?%d?%d?%d)$'
 -- {::1}:2000 -> ::1,2000
 local inet6_re = '^{([^}]+)}:(%d?%d?%d?%d?%d)$'
-function export.listen_tcp(address: string, handler: function(): void): void
+function export.listen_tcp(address: string, conn_cap: uinteger, handler: function(): void): void
   local matched, matches
   if (do matched, matches = address:match(inet_re); in matched end) then
     local c_err: cint = 0
@@ -283,7 +283,7 @@ function export.listen_tcp(address: string, handler: function(): void): void
 
     defer c_close(fd) end
 
-    export.listen_sock(fd, handler)
+    export.listen_sock(fd, conn_cap, handler)
   elseif (do matched, matches = address:match(inet6_re); in matched end) then
     local c_err: cint = 0
 
@@ -311,7 +311,7 @@ function export.listen_tcp(address: string, handler: function(): void): void
 
     defer c_close(fd) end
 
-    export.listen_sock(fd, handler)
+    export.listen_sock(fd, conn_cap, handler)
   end
 end
 
-- 
2.43.4