[Minor] Allow to get memory stats from the workers

1 files changed, 35 insertions, 0 deletions

diff --git a/src/lua/lua_worker.c b/src/lua/lua_worker.c
index 3e38fa3f1..6c73736c8 100644
--- a/src/lua/lua_worker.c
+++ b/src/lua/lua_worker.c
@@ -20,6 +20,10 @@
 #include "rspamd_control.h"
 #include "ottery.h"
 
+#ifdef WITH_JEMALLOC
+#include <jemalloc/jemalloc.h>
+#endif
+
 #include <sys/wait.h>
 
 /***
@@ -36,6 +40,7 @@ LUA_FUNCTION_DEF (worker, get_pid);
 LUA_FUNCTION_DEF (worker, is_scanner);
 LUA_FUNCTION_DEF (worker, is_primary_controller);
 LUA_FUNCTION_DEF (worker, spawn_process);
+LUA_FUNCTION_DEF (worker, get_mem_stats);
 
 const luaL_reg worker_reg[] = {
 		LUA_INTERFACE_DEF (worker, get_name),
@@ -45,6 +50,7 @@ const luaL_reg worker_reg[] = {
 		LUA_INTERFACE_DEF (worker, spawn_process),
 		LUA_INTERFACE_DEF (worker, is_scanner),
 		LUA_INTERFACE_DEF (worker, is_primary_controller),
+		LUA_INTERFACE_DEF (worker, get_mem_stats),
 		{"__tostring", rspamd_lua_class_tostring},
 		{NULL, NULL}
 };
@@ -217,6 +223,35 @@ lua_worker_is_primary_controller (lua_State *L)
 	return 1;
 }
 
+#ifdef WITH_JEMALLOC
+static void
+lua_worker_jemalloc_stats_cb (void *ud, const char *msg)
+{
+	lua_State *L = (lua_State *)ud;
+
+	lua_pushstring (L, msg);
+}
+#endif
+
+static gint
+lua_worker_get_mem_stats (lua_State *L)
+{
+	struct rspamd_worker *w = lua_check_worker (L, 1);
+
+	if (w) {
+#ifdef WITH_JEMALLOC
+		malloc_stats_print (lua_worker_jemalloc_stats_cb, (void *)L, NULL);
+#else
+		lua_pushstring (L, "no stats, jemalloc support is required");
+#endif
+	}
+	else {
+		return luaL_error (L, "invalid arguments");
+	}
+
+	return 1;
+}
+
 struct rspamd_lua_process_cbdata {
 	gint sp[2];
 	gint func_cbref;