proxy: add request and buffer memory limits

Adds:
mcp.active_req_limit(count)
mcp.buffer_memory_limit(kilobytes)

Divides by the number of worker threads and creates a per-worker-thread
limit for the number of concurrent proxy requests, and how many bytes
used specifically for value bytes. This does not represent total memory
usage but will be close.

Buffer memory for inbound set requests is not accounted for until after
the object has been read from the socket; to be improved in a future
update. This should be fine unless clients send just the SET request and
then hang without sending further data.

Limits should be live-adjustable via configuration reloads.

1 files changed, 49 insertions, 0 deletions

diff --git a/proxy_lua.c b/proxy_lua.c
index 852624e..25209d0 100644
--- a/proxy_lua.c
+++ b/proxy_lua.c
@@ -82,11 +82,15 @@ static int mcplib_response_line(lua_State *L) {
 }
 
 static int mcplib_response_gc(lua_State *L) {
+    LIBEVENT_THREAD *t = PROXY_GET_THR(L);
     mcp_resp_t *r = luaL_checkudata(L, -1, "mcp.response");
 
     // On error/similar we might be holding the read buffer.
     // If the buf is handed off to mc_resp for return, this pointer is NULL
     if (r->buf != NULL) {
+        pthread_mutex_lock(&t->proxy_limit_lock);
+        t->proxy_buffer_memory_used -= r->blen;
+        pthread_mutex_unlock(&t->proxy_limit_lock);
         free(r->buf);
     }
 
@@ -872,6 +876,49 @@ static int mcplib_backend_read_timeout(lua_State *L) {
     return 0;
 }
 
+static int mcplib_active_req_limit(lua_State *L) {
+    proxy_ctx_t *ctx = PROXY_GET_CTX(L);
+    uint64_t limit = luaL_checkinteger(L, -1);
+
+    if (limit == 0) {
+        limit = UINT64_MAX;
+    } else {
+        // FIXME: global
+        int tcount = settings.num_threads;
+        // The actual limit is per-worker-thread, so divide it up.
+        if (limit > tcount * 2) {
+            limit /= tcount;
+        }
+    }
+
+    STAT_L(ctx);
+    ctx->active_req_limit = limit;
+    STAT_UL(ctx);
+
+    return 0;
+}
+
+// limit specified in kilobytes
+static int mcplib_buffer_memory_limit(lua_State *L) {
+    proxy_ctx_t *ctx = PROXY_GET_CTX(L);
+    uint64_t limit = luaL_checkinteger(L, -1);
+
+    if (limit == 0) {
+        limit = UINT64_MAX;
+    } else {
+        limit *= 1024;
+
+        int tcount = settings.num_threads;
+        if (limit > tcount * 2) {
+            limit /= tcount;
+        }
+
+        ctx->buffer_memory_limit = limit;
+    }
+
+    return 0;
+}
+
 // mcp.attach(mcp.HOOK_NAME, function)
 // fill hook structure: if lua function, use luaL_ref() to store the func
 static int mcplib_attach(lua_State *L) {
@@ -1203,6 +1250,8 @@ int proxy_register_libs(void *ctx, LIBEVENT_THREAD *t, void *state) {
         {"backend_read_timeout", mcplib_backend_read_timeout},
         {"backend_failure_limit", mcplib_backend_failure_limit},
         {"tcp_keepalive", mcplib_tcp_keepalive},
+        {"active_req_limit", mcplib_active_req_limit},
+        {"buffer_memory_limit", mcplib_buffer_memory_limit},
         {NULL, NULL}
     };