External memory-usage tracking

Author: Allen Short

docs/sandbox.md

@@ -4,6 +4,7 @@
 
 * **output_limit** - the largest output string an input or analysis plugin can inject into the host (bytes, default 64KiB)
 * **memory_limit** - the maximum amount of memory a plugin can use before being terminated (bytes, default 8MiB)
+* **external_memory_limit** - the maximum amount of memory external libraries (exposed to the sandbox as userdata) can use before plugin is terminated; set to 0 for unlimited (bytes, default unlimited)
 * **instruction_limit** - the maximum number of Lua instructions a plugin can execute in a single API function call (count, default 1MM)
 * **path** - The path used by require to search for a Lua loader. See [package loaders](http://www.lua.org/manual/5.1/manual.html#pdf-package.loaders)
   for the path syntax.  By default no paths are set in the sandbox and everything has been moved to a sandbox configuration table.

include/luasandbox.h

@@ -27,15 +27,16 @@
 
 #define LSB_ERROR_SIZE 256
 
-#define LSB_SHUTTING_DOWN     "shutting down"
-#define LSB_CONFIG_TABLE      "lsb_config"
-#define LSB_THIS_PTR          "lsb_this_ptr"
-#define LSB_MEMORY_LIMIT      "memory_limit"
-#define LSB_INSTRUCTION_LIMIT "instruction_limit"
-#define LSB_OUTPUT_LIMIT      "output_limit"
-#define LSB_LUA_PATH          "path"
-#define LSB_LUA_CPATH         "cpath"
-#define LSB_NIL_ERROR         "<nil error message>"
+#define LSB_SHUTTING_DOWN         "shutting down"
+#define LSB_CONFIG_TABLE          "lsb_config"
+#define LSB_THIS_PTR              "lsb_this_ptr"
+#define LSB_MEMORY_LIMIT          "memory_limit"
+#define LSB_EXTERNAL_MEMORY_LIMIT "external_memory_limit"
+#define LSB_INSTRUCTION_LIMIT     "instruction_limit"
+#define LSB_OUTPUT_LIMIT          "output_limit"
+#define LSB_LUA_PATH              "path"
+#define LSB_LUA_CPATH             "cpath"
+#define LSB_NIL_ERROR             "<nil error message>"
 
 typedef enum {
   LSB_UNKNOWN     = 0,
@@ -53,9 +54,10 @@ typedef enum {
 } lsb_usage_stat;
 
 typedef enum {
-  LSB_UT_MEMORY       = 0,
-  LSB_UT_INSTRUCTION  = 1,
-  LSB_UT_OUTPUT       = 2,
+  LSB_UT_MEMORY          = 0,
+  LSB_UT_INSTRUCTION     = 1,
+  LSB_UT_OUTPUT          = 2,
+  LSB_UT_EXTERNAL_MEMORY = 3,
 
   LSB_UT_MAX
 } lsb_usage_type;

include/luasandbox/heka/sandbox.h

@@ -54,6 +54,8 @@ typedef struct lsb_heka_sandbox lsb_heka_sandbox;
 typedef struct lsb_heka_stats {
   unsigned long long mem_cur;
   unsigned long long mem_max;
+  unsigned long long ext_mem_cur;
+  unsigned long long ext_mem_max;
   unsigned long long ins_max;
   unsigned long long out_max;
   unsigned long long im_cnt;
@@ -338,6 +340,16 @@ LSB_HEKA_EXPORT const char* lsb_heka_get_lua_file(lsb_heka_sandbox *hsb);
  */
 LSB_HEKA_EXPORT lsb_heka_stats lsb_heka_get_stats(lsb_heka_sandbox *hsb);
 
+/**
+ * Retrieve the limit on external memory allocation by native modules, or 0 if none is set.
+ * @param hsb Heka sandbox
+ *
+ * @return size_t Memory allocation limit
+ */
+LSB_HEKA_EXPORT size_t lsb_heka_get_ext_memory_limit(lsb_heka_sandbox *hsb);
+
+LSB_HEKA_EXPORT void lsb_heka_adjust_ext_memory_usage(lsb_heka_sandbox *hsb, int sizechange);
+
 /**
  * Queries the state of the sandbox.
  *

src/heka/sandbox.c

@@ -10,6 +10,7 @@
 
 #include <errno.h>
 #include <math.h>
+#include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
 
@@ -1058,24 +1059,41 @@ const char* lsb_heka_get_lua_file(lsb_heka_sandbox *hsb)
 
 lsb_heka_stats lsb_heka_get_stats(lsb_heka_sandbox *hsb)
 {
-  if (!hsb) return (struct lsb_heka_stats){ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
+  if (!hsb) return (struct lsb_heka_stats){ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
 
   return (struct lsb_heka_stats){
-    .mem_cur      = lsb_usage(hsb->lsb, LSB_UT_MEMORY, LSB_US_CURRENT),
-    .mem_max      = lsb_usage(hsb->lsb, LSB_UT_MEMORY, LSB_US_MAXIMUM),
-    .out_max      = lsb_usage(hsb->lsb, LSB_UT_OUTPUT, LSB_US_MAXIMUM),
-    .ins_max      = lsb_usage(hsb->lsb, LSB_UT_INSTRUCTION, LSB_US_MAXIMUM),
-    .im_cnt       = hsb->stats.im_cnt,
-    .im_bytes     = hsb->stats.im_bytes,
-    .pm_cnt       = hsb->stats.pm_cnt,
-    .pm_failures  = hsb->stats.pm_failures,
-    .pm_avg       = hsb->stats.pm.mean,
-    .pm_sd        = lsb_sd_running_stats(&hsb->stats.pm),
-    .te_avg       = hsb->stats.te.mean,
-    .te_sd        = lsb_sd_running_stats(&hsb->stats.te)
+    .mem_cur     = lsb_usage(hsb->lsb, LSB_UT_MEMORY, LSB_US_CURRENT),
+    .mem_max     = lsb_usage(hsb->lsb, LSB_UT_MEMORY, LSB_US_MAXIMUM),
+    .ext_mem_cur = lsb_usage(hsb->lsb, LSB_UT_EXTERNAL_MEMORY, LSB_US_CURRENT),
+    .ext_mem_max = lsb_usage(hsb->lsb, LSB_UT_EXTERNAL_MEMORY, LSB_US_MAXIMUM),
+    .out_max     = lsb_usage(hsb->lsb, LSB_UT_OUTPUT, LSB_US_MAXIMUM),
+    .ins_max     = lsb_usage(hsb->lsb, LSB_UT_INSTRUCTION, LSB_US_MAXIMUM),
+    .im_cnt      = hsb->stats.im_cnt,
+    .im_bytes    = hsb->stats.im_bytes,
+    .pm_cnt      = hsb->stats.pm_cnt,
+    .pm_failures = hsb->stats.pm_failures,
+    .pm_avg      = hsb->stats.pm.mean,
+    .pm_sd       = lsb_sd_running_stats(&hsb->stats.pm),
+    .te_avg      = hsb->stats.te.mean,
+    .te_sd       = lsb_sd_running_stats(&hsb->stats.te)
   };
 }
 
+size_t lsb_heka_get_ext_memory_limit(lsb_heka_sandbox *hsb)
+{
+  return hsb ? lsb_usage(hsb->lsb, LSB_UT_EXTERNAL_MEMORY, LSB_US_LIMIT) : 0;
+}
+
+void lsb_heka_adjust_ext_memory_usage(lsb_heka_sandbox *hsb, int sizechange)
+{
+  size_t oldsize = hsb->lsb->usage[LSB_UT_EXTERNAL_MEMORY][LSB_US_CURRENT];
+  size_t newsize = hsb->lsb->usage[LSB_UT_EXTERNAL_MEMORY][LSB_US_CURRENT] + sizechange;
+
+  if (newsize > hsb->lsb->usage[LSB_UT_EXTERNAL_MEMORY][LSB_US_MAXIMUM]) {
+    hsb->lsb->usage[LSB_UT_EXTERNAL_MEMORY][LSB_US_MAXIMUM] = newsize;
+  }
+  hsb->lsb->usage[LSB_UT_EXTERNAL_MEMORY][LSB_US_CURRENT] = newsize;
+}
 
 bool lsb_heka_is_running(lsb_heka_sandbox *hsb)
 {

src/heka/test/test_heka_sandbox.c

@@ -350,6 +350,8 @@ static char* test_timer_event()
   lsb_heka_stats stats = lsb_heka_get_stats(hsb);
   mu_assert(0 < stats.mem_cur, "received %llu", stats.mem_cur);
   mu_assert(0 < stats.mem_max, "received %llu", stats.mem_max);
+  mu_assert(0 == stats.ext_mem_cur, "received %llu", stats.ext_mem_cur);
+  mu_assert(0 == stats.ext_mem_max, "received %llu", stats.ext_mem_max);
   mu_assert(0 == stats.out_max, "received %llu", stats.out_max);
   mu_assert(0 < stats.ins_max, "received %llu", stats.ins_max);
   mu_assert(0 == stats.pm_cnt, "received %llu", stats.pm_cnt);

src/luasandbox.c

@@ -297,6 +297,9 @@ static lua_State* load_sandbox_config(const char *cfg, lsb_logger *logger)
   ret = check_int(L, LUA_GLOBALSINDEX, LSB_MEMORY_LIMIT, 8 * 1024 * 1024);
   if (ret) goto cleanup;
 
+  ret = check_int(L, LUA_GLOBALSINDEX, LSB_EXTERNAL_MEMORY_LIMIT, 0);
+  if (ret) goto cleanup;
+
   ret = check_int(L, LUA_GLOBALSINDEX, LSB_INSTRUCTION_LIMIT, 1000000);
   if (ret) goto cleanup;
 
@@ -477,6 +480,7 @@ lsb_lua_sandbox* lsb_create(void *parent,
   lua_pop(lua_cfg, 2);
   lua_close(lua_cfg);
   size_t ml = get_int(lsb->lua, -1, "memory_limit");
+  size_t eml = get_int(lsb->lua, -1, "external_memory_limit");
   size_t il = get_int(lsb->lua, -1, "instruction_limit");
   size_t ol = get_int(lsb->lua, -1, "output_limit");
   int log_level = get_int(lsb->lua, -1, "log_level");
@@ -492,6 +496,7 @@ lsb_lua_sandbox* lsb_create(void *parent,
 
   lsb->parent = parent;
   lsb->usage[LSB_UT_MEMORY][LSB_US_LIMIT] = ml;
+  lsb->usage[LSB_UT_EXTERNAL_MEMORY][LSB_US_LIMIT] = eml;
   lsb->usage[LSB_UT_INSTRUCTION][LSB_US_LIMIT] = il;
   lsb->usage[LSB_UT_OUTPUT][LSB_US_LIMIT] = ol;
   lsb->state = LSB_UNKNOWN;

src/test/lua/read_config.lua

@@ -1,5 +1,6 @@
 assert(type(read_config) == "function")
 assert(read_config("memory_limit") == 65765)
+assert(read_config("external_memory_limit") == 1610612736)
 assert(read_config("instruction_limit") == 1000)
 assert(read_config("output_limit") == 1024)
 

src/test/test_generic_sandbox.c

@@ -223,6 +223,7 @@ static char* test_create_error()
 static char* test_read_config()
 {
   const char *cfg = "memory_limit = 65765\n"
+      "external_memory_limit = 1610612736\n"
       "instruction_limit = 1000\n"
       "output_limit = 1024\n"
       "array = {'foo', 99}\n"