From c478e101f52171e5cd31f216a5db3a99987f6851 Mon Sep 17 00:00:00 2001
From: Paul Kulchenko <paul@kulchenko.com>
Date: Thu, 28 Jul 2022 21:04:24 -0700
Subject: [PATCH] Add max worker processing

---
 tool/net/help.txt  |  6 ++++++
 tool/net/redbean.c | 12 ++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/tool/net/help.txt b/tool/net/help.txt
index 128a0feb7..fe9b83f5f 100644
--- a/tool/net/help.txt
+++ b/tool/net/help.txt
@@ -1403,6 +1403,12 @@ FUNCTIONS
           increased to 1450, since that's the size of ethernet frames.
           This function can only be called from .init.lua.
 
+  ProgramMaxWorkers(int)
+          Limits the number of workers forked by redbean. If that number
+          is reached, the server continues polling until the number of
+          workers is reduced or the value is updated. Setting it to 0
+          removes the limit (this is the default).
+
   ProgramPrivateKey(pem:str)
           Same as the -K flag if called from .init.lua, e.g.
           ProgramPrivateKey(LoadAsset("/.sign.key")) for zip loading or
diff --git a/tool/net/redbean.c b/tool/net/redbean.c
index 360bc3ef5..f2fc02097 100644
--- a/tool/net/redbean.c
+++ b/tool/net/redbean.c
@@ -411,6 +411,7 @@ static int sandboxed;
 static int changeuid;
 static int changegid;
 static int isyielding;
+static int maxworkers;
 static int statuscode;
 static int shutdownsig;
 static int sslpskindex;
@@ -4747,6 +4748,15 @@ static int LuaProgramHeartbeatInterval(lua_State *L) {
   return 1;
 }
 
+static int LuaProgramMaxWorkers(lua_State *L) {
+  if (!lua_isinteger(L, 1) && !lua_isnoneornil(L, 1)) {
+    return luaL_argerror(L, 1, "invalid number of workers; integer expected");
+  }
+  lua_pushinteger(L, maxworkers);
+  if (lua_isinteger(L, 1)) maxworkers = lua_tointeger(L, 1);
+  return 1;
+}
+
 static int LuaProgramAddr(lua_State *L) {
   uint32_t ip;
   OnlyCallFromInitLua(L, "ProgramAddr");
@@ -5239,6 +5249,7 @@ static const luaL_Reg kLuaFuncs[] = {
     {"ProgramTimeout", LuaProgramTimeout},                      //
     {"ProgramUid", LuaProgramUid},                              //
     {"ProgramUniprocess", LuaProgramUniprocess},                //
+    {"ProgramMaxWorkers", LuaProgramMaxWorkers},                //
     {"Rand64", LuaRand64},                                      //
     {"Rdrand", LuaRdrand},                                      //
     {"Rdseed", LuaRdseed},                                      //
@@ -6927,6 +6938,7 @@ static int HandlePoll(int ms) {
         if (!polls[pollid].revents) continue;
         if (polls[pollid].fd < 0) continue;
         if (polls[pollid].fd) {
+          if (maxworkers && shared->workers >= maxworkers) continue;
           // handle listen socket
           lua_repl_lock();
           serverid = pollid - 1;