From 4b96ade66076aab0d47f3c12351c20490c3fe206 Mon Sep 17 00:00:00 2001
From: Cloud Wu <cloudwu@gmail.com>
Date: Fri, 6 May 2016 23:40:44 +0800
Subject: [PATCH] support sandbox memory limit

---
 lualib/skynet/manager.lua   |  4 ++++
 service-src/service_snlua.c | 29 ++++++++++++++++++++++-------
 test/testmemlimit.lua       | 18 ++++++++++++++++++
 3 files changed, 44 insertions(+), 7 deletions(-)
 create mode 100644 test/testmemlimit.lua

diff --git a/lualib/skynet/manager.lua b/lualib/skynet/manager.lua
index 4365d2e2..627d1e12 100644
--- a/lualib/skynet/manager.lua
+++ b/lualib/skynet/manager.lua
@@ -20,6 +20,10 @@ function skynet.abort()
 	c.command("ABORT")
 end
 
+function skynet.memlimit(bytes)
+	debug.getregistry().memlimit = bytes
+end
+
 local function globalname(name, handle)
 	local c = string.sub(name,1,1)
 	assert(c ~= ':')
diff --git a/service-src/service_snlua.c b/service-src/service_snlua.c
index bd2f6b2c..4870baa9 100644
--- a/service-src/service_snlua.c
+++ b/service-src/service_snlua.c
@@ -16,6 +16,7 @@ struct snlua {
 	struct skynet_context * ctx;
 	size_t mem;
 	size_t mem_report;
+	size_t mem_limit;
 };
 
 // LUA_CACHELIB may defined in patched lua for shared proto
@@ -57,7 +58,7 @@ traceback (lua_State *L) {
 }
 
 static void
-_report_launcher_error(struct skynet_context *ctx) {
+report_launcher_error(struct skynet_context *ctx) {
 	// sizeof "ERROR" == 5
 	skynet_sendname(ctx, 0, ".launcher", PTYPE_TEXT, 0, "ERROR", 5);
 }
@@ -72,7 +73,7 @@ optstring(struct skynet_context *ctx, const char *key, const char * str) {
 }
 
 static int
-_init(struct snlua *l, struct skynet_context *ctx, const char * args, size_t sz) {
+init_cb(struct snlua *l, struct skynet_context *ctx, const char * args, size_t sz) {
 	lua_State *L = l->L;
 	l->ctx = ctx;
 	lua_gc(L, LUA_GCSTOP, 0);
@@ -105,17 +106,23 @@ _init(struct snlua *l, struct skynet_context *ctx, const char * args, size_t sz)
 	int r = luaL_loadfile(L,loader);
 	if (r != LUA_OK) {
 		skynet_error(ctx, "Can't load %s : %s", loader, lua_tostring(L, -1));
-		_report_launcher_error(ctx);
+		report_launcher_error(ctx);
 		return 1;
 	}
 	lua_pushlstring(L, args, sz);
 	r = lua_pcall(L,1,0,1);
 	if (r != LUA_OK) {
 		skynet_error(ctx, "lua loader error : %s", lua_tostring(L, -1));
-		_report_launcher_error(ctx);
+		report_launcher_error(ctx);
 		return 1;
 	}
 	lua_settop(L,0);
+	if (lua_getfield(L, LUA_REGISTRYINDEX, "memlimit") == LUA_TNUMBER) {
+		size_t limit = lua_tointeger(L, -1);
+		l->mem_limit = limit;
+		skynet_error(ctx, "Set memory limit to %.2f M", (float)limit / (1024 * 1024));
+	}
+	lua_pop(L, 1);
 
 	lua_gc(L, LUA_GCRESTART, 0);
 
@@ -123,11 +130,11 @@ _init(struct snlua *l, struct skynet_context *ctx, const char * args, size_t sz)
 }
 
 static int
-_launch(struct skynet_context * context, void *ud, int type, int session, uint32_t source , const void * msg, size_t sz) {
+launch_cb(struct skynet_context * context, void *ud, int type, int session, uint32_t source , const void * msg, size_t sz) {
 	assert(type == 0 && session == 0);
 	struct snlua *l = ud;
 	skynet_callback(context, NULL, NULL);
-	int err = _init(l, context, msg, sz);
+	int err = init_cb(l, context, msg, sz);
 	if (err) {
 		skynet_command(context, "EXIT", NULL);
 	}
@@ -140,7 +147,7 @@ snlua_init(struct snlua *l, struct skynet_context *ctx, const char * args) {
 	int sz = strlen(args);
 	char * tmp = skynet_malloc(sz);
 	memcpy(tmp, args, sz);
-	skynet_callback(ctx, l , _launch);
+	skynet_callback(ctx, l , launch_cb);
 	const char * self = skynet_command(ctx, "REG", NULL);
 	uint32_t handle_id = strtoul(self+1, NULL, 16);
 	// it must be first message
@@ -151,9 +158,16 @@ snlua_init(struct snlua *l, struct skynet_context *ctx, const char * args) {
 static void *
 lalloc(void * ud, void *ptr, size_t osize, size_t nsize) {
 	struct snlua *l = ud;
+	size_t mem = l->mem;
 	l->mem += nsize;
 	if (ptr)
 		l->mem -= osize;
+	if (l->mem_limit != 0 && l->mem > l->mem_limit) {
+		if (ptr == NULL || nsize > osize) {
+			l->mem = mem;
+			return NULL;
+		}
+	}
 	if (l->mem > l->mem_report) {
 		l->mem_report *= 2;
 		skynet_error(l->ctx, "Memory warning %.2f M", (float)l->mem / (1024 * 1024));
@@ -166,6 +180,7 @@ snlua_create(void) {
 	struct snlua * l = skynet_malloc(sizeof(*l));
 	memset(l,0,sizeof(*l));
 	l->mem_report = MEMORY_WARNING_REPORT;
+	l->mem_limit = 0;
 	l->L = lua_newstate(lalloc, l);
 	return l;
 }
diff --git a/test/testmemlimit.lua b/test/testmemlimit.lua
new file mode 100644
index 00000000..532df727
--- /dev/null
+++ b/test/testmemlimit.lua
@@ -0,0 +1,18 @@
+local skynet = require "skynet"
+require "skynet.manager"
+
+-- set sandbox memory limit to 1M, must set here (at start, out of skynet.start)
+skynet.memlimit(1 * 1024 * 1024)
+
+skynet.start(function()
+	local a = {}
+	local limit
+	local ok, err = pcall(function()
+		for i=1, 1000000 do
+			limit = i
+			table.insert(a, {})
+		end
+	end)
+	skynet.error(limit, err)
+	skynet.exit()
+end)