feat: fix lua memory deallocation and lua_mem_gc_threshold=0 disable lua force GC calls (#5587)

* fix: lua script memory free

Author: BorysTheDev

src/core/interpreter.cc

@@ -58,7 +58,7 @@ ABSL_FLAG(LuaGcFlag, luagc, {},
 
 ABSL_FLAG(uint64_t, lua_mem_gc_threshold, 10000000,
           "Specifies Lua interpreter's per thread memory limit in bytes after which the GC will be "
-          "called forcefully.");
+          "called forcefully. 0 value remove forced GC calls");
 
 static bool AbslParseFlag(std::string_view in, LuaGcFlag* flag, std::string* err) {
   if (in.empty()) {
@@ -562,19 +562,24 @@ int RedisLogCommand(lua_State* lua) {
 
 // See https://www.lua.org/manual/5.3/manual.html#lua_Alloc
 void* mimalloc_glue(void* ud, void* ptr, size_t osize, size_t nsize) {
-  (void)ud;
+  int64_t& used_bytes = *static_cast<int64_t*>(ud);
+
   if (nsize == 0) {
-    InterpreterManager::tl_stats().used_bytes -= mi_usable_size(ptr);
+    used_bytes -= mi_usable_size(ptr);
     mi_free_size(ptr, osize);
     return nullptr;
   } else if (ptr == nullptr) {
     ptr = mi_malloc(nsize);
-    InterpreterManager::tl_stats().used_bytes += mi_usable_size(ptr);
+    used_bytes += mi_usable_size(ptr);
     return ptr;
   } else {
-    InterpreterManager::tl_stats().used_bytes -= mi_usable_size(ptr);
+    const auto old_size = mi_usable_size(ptr);
     ptr = mi_realloc(ptr, nsize);
-    InterpreterManager::tl_stats().used_bytes += mi_usable_size(ptr);
+    if (ptr) {
+      used_bytes -= old_size;
+      used_bytes += mi_usable_size(ptr);
+    }
+
     return ptr;
   }
 }
@@ -584,7 +589,9 @@ void* mimalloc_glue(void* ud, void* ptr, size_t osize, size_t nsize) {
 Interpreter::Interpreter() {
   InterpreterManager::tl_stats().interpreter_cnt++;
 
-  lua_ = lua_newstate(mimalloc_glue, nullptr);
+  // interpreter can be runnned in different threads so we need to calculate
+  // used memory via &used_bytes_ additional parameter
+  lua_ = lua_newstate(mimalloc_glue, &used_bytes_);
   InitLua(lua_);
   void** ptr = static_cast<void**>(lua_getextraspace(lua_));
   *ptr = this;
@@ -1211,10 +1218,17 @@ void InterpreterManager::Return(Interpreter* ir) {
   const uint64_t max_memory_usage = absl::GetFlag(FLAGS_lua_mem_gc_threshold);
   using namespace chrono;
   ++tl_stats().interpreter_return;
-  if (tl_stats().used_bytes > max_memory_usage) {
+  tl_stats().used_bytes += ir->TakeUsedBytes();
+  if (max_memory_usage != 0 && tl_stats().used_bytes > max_memory_usage) {
     ++tl_stats().force_gc_calls;
     auto before = steady_clock::now();
     tl_stats().gc_freed_memory += ir->RunGC();
+
+    VLOG(2) << "stats_used_bytes: " << tl_stats().used_bytes
+            << " lua_mem_gc_threshold: " << max_memory_usage
+            << " force_gc_calls: " << tl_stats().force_gc_calls
+            << " freed_mem: " << tl_stats().gc_freed_memory;
+
     auto after = steady_clock::now();
     tl_stats().gc_duration_ns += duration_cast<nanoseconds>(after - before).count();
   }

src/core/interpreter.h

@@ -89,6 +89,10 @@ class Interpreter {
   // Add function with sha and body to interpreter.
   AddResult AddFunction(std::string_view sha, std::string_view body, std::string* error);
 
+  int64_t TakeUsedBytes() {
+    return std::exchange(used_bytes_, 0);
+  }
+
   bool Exists(std::string_view sha) const;
 
   enum RunResult {
@@ -150,6 +154,7 @@ class Interpreter {
   unsigned cmd_depth_ = 0;
   RedisFunc redis_func_;
   std::string buffer_;
+  int64_t used_bytes_ = 0;
   char name_buffer_[32];  // backing storage for cmd name
 };
 

src/core/interpreter_test.cc

@@ -13,6 +13,8 @@ extern "C" {
 #include <absl/strings/str_replace.h>
 #include <gmock/gmock.h>
 
+#include <thread>
+
 #include "base/gtest.h"
 #include "base/logging.h"
 
@@ -526,4 +528,68 @@ TEST_F(InterpreterTest, LuaIntOverflow) {
   EXPECT_FALSE(Execute("EVAL \"struct.pack('>I2147483648', '10')\" 0"));
 }
 
+TEST_F(InterpreterTest, LuaGcStatistic) {
+  InterpreterManager im(1);
+  auto* interpreter = im.Get();
+
+  std::string_view keys[] = {"key1", "key2", "key3", "key4", "key5", "key6", "key7"};
+  interpreter->SetGlobalArray("KEYS", SliceSpan{keys});
+
+  auto cb = [](Interpreter::CallArgs ca) {
+    auto* reply = ca.translator;
+    reply->OnInt(1);
+  };
+  interpreter->SetRedisFunc(cb);
+  // next script generate several big values and set them to the keys
+  // after the script is finished, GM isn't called for all values and
+  // in the most cases we have more than 300k allocated memory
+  // that will be cleaned later in the separate thread
+  std::string script = R"(
+        for i = 1, 7 do
+          local str = string.rep(i, 1024 * 100)
+          redis.call('SET', KEYS[1], str .. str)
+        end
+       )";
+
+  char sha_buf[64];
+  Interpreter::FuncSha1(script, sha_buf);
+  string_view sha{sha_buf, std::strlen(sha_buf)};
+
+  string result;
+  Interpreter::AddResult add_res = interpreter->AddFunction(sha, script, &result);
+  EXPECT_EQ(Interpreter::ADD_OK, add_res);
+
+  // When script is executed in the most cases we see that not all memory was deallocated
+  // immediately and can be deallocated later
+  Interpreter::RunResult run_res = interpreter->RunFunction(sha, &error_);
+  EXPECT_EQ(Interpreter::RUN_OK, run_res);
+
+  // check that after script is finished not the all memory was deallocated
+  uint64_t used_bytes = InterpreterManager::tl_stats().used_bytes;
+  EXPECT_GE(used_bytes, 0);
+
+  auto force_gc_calls = InterpreterManager::tl_stats().force_gc_calls;
+  // we need return interpreter to update statistic
+  // force_gc_calls shouldn't be called
+  im.Return(interpreter);
+  EXPECT_EQ(force_gc_calls, InterpreterManager::tl_stats().force_gc_calls);
+  EXPECT_LE(used_bytes, InterpreterManager::tl_stats().used_bytes);
+
+  used_bytes = InterpreterManager::tl_stats().used_bytes;
+
+  // we get the same interpeter again to call GC in separate thread
+  auto* new_interpreter = im.Get();
+  EXPECT_EQ(interpreter, new_interpreter);
+
+  // check that even if memory is deallocated in separate thread our statistic is correct
+  std::thread t([&] {
+    interpreter->RunGC();
+    EXPECT_EQ(InterpreterManager::tl_stats().used_bytes, 0);
+  });
+  t.join();
+
+  im.Return(interpreter);
+  EXPECT_GE(used_bytes, InterpreterManager::tl_stats().used_bytes);
+}
+
 }  // namespace dfly