FUNCTION FLUSH re-create lua VM, fix flush not gc, fix flush async + load crash (#1826)

There will be two issues in this test:
```
test {FUNCTION - test function flush} {
    for {set i 0} {$i < 10000} {incr i} {
        r function load [get_function_code LUA test_$i test_$i {return 'hello'}]
    }
    set before_flush_memory [s used_memory_vm_functions]
    r function flush sync
    set after_flush_memory [s used_memory_vm_functions]
    puts "flush sync, before_flush_memory: $before_flush_memory, after_flush_memory: $after_flush_memory"

    for {set i 0} {$i < 10000} {incr i} {
        r function load [get_function_code LUA test_$i test_$i {return 'hello'}]
    }
    set before_flush_memory [s used_memory_vm_functions]
    r function flush async
    set after_flush_memory [s used_memory_vm_functions]
    puts "flush async, before_flush_memory: $before_flush_memory, after_flush_memory: $after_flush_memory"

    for {set i 0} {$i < 10000} {incr i} {
        r function load [get_function_code LUA test_$i test_$i {return 'hello'}]
    }
    puts "Test done"
}
```

The first one is the test output, we can see that after executing
FUNCTION FLUSH,
used_memory_vm_functions has not changed at all:
```
flush sync, before_flush_memory: 2962432, after_flush_memory: 2962432
flush async, before_flush_memory: 4504576, after_flush_memory: 4504576
```

The second one is there is a crash when loading the functions during the
async
flush:
```
=== VALKEY BUG REPORT START: Cut & paste starting from here ===
 # valkey 255.255.255 crashed by signal: 11, si_code: 2
 # Accessing address: 0xe0429b7100000a3c
 # Crashed running the instruction at: 0x102e0b09c

------ STACK TRACE ------
EIP:
0   valkey-server                       0x0000000102e0b09c luaH_getstr + 52

Backtrace:
0   libsystem_platform.dylib            0x000000018b066584 _sigtramp + 56
1   valkey-server                       0x0000000102e01054 luaD_precall + 96
2   valkey-server                       0x0000000102e01b10 luaD_call + 104
3   valkey-server                       0x0000000102e00d1c luaD_rawrunprotected + 76
4   valkey-server                       0x0000000102e01e3c luaD_pcall + 60
5   valkey-server                       0x0000000102dfc630 lua_pcall + 300
6   valkey-server                       0x0000000102f77770 luaEngineCompileCode + 708
7   valkey-server                       0x0000000102f71f50 scriptingEngineCallCompileCode + 104
8   valkey-server                       0x0000000102f700b0 functionsCreateWithLibraryCtx + 2088
9   valkey-server                       0x0000000102f70898 functionLoadCommand + 312
10  valkey-server                       0x0000000102e3978c call + 416
11  valkey-server                       0x0000000102e3b5b8 processCommand + 3340
12  valkey-server                       0x0000000102e563cc processInputBuffer + 520
13  valkey-server                       0x0000000102e55808 readQueryFromClient + 92
14  valkey-server                       0x0000000102f696e0 connSocketEventHandler + 180
15  valkey-server                       0x0000000102e20480 aeProcessEvents + 372
16  valkey-server                       0x0000000102e4aad0 main + 26412
17  dyld                                0x000000018acab154 start + 2476

------ STACK TRACE DONE ------
```

The reason is that, in the old implementation (introduced in 7.0),
FUNCTION FLUSH
use lua_unref to remove the script from lua VM. lua_unref does not
trigger the gc,
it causes us to not be able to effectively reclaim memory after the
FUNCTION FLUSH.

The other issue is that, since we don't re-create the lua VM in FUNCTION
FLUSH,
loading the functions during a FUNCTION FLUSH ASYNC will result a crash
because
lua engine state is not thread-safe.

The correct solution is to re-create a new Lua VM to use, just like
SCRIPT FLUSH.

---------

Signed-off-by: Binbin <[email protected]>
Signed-off-by: Ricardo Dias <[email protected]>
Co-authored-by: Ricardo Dias <[email protected]>
(cherry picked from commit b4c93cc)
Signed-off-by: cherukum-amazon <[email protected]>

Author: 2 people

src/db.c

@@ -687,7 +687,7 @@ long long emptyData(int dbnum, int flags, void(callback)(hashtable *)) {
     if (with_functions) {
         serverAssert(dbnum == -1);
         /* TODO: fix this callback incompatibility. The arg is not used. */
-        functionsLibCtxClearCurrent(async, (void (*)(dict *))callback);
+        functionReset(async, (void (*)(dict *))callback);
     }
 
     /* Also fire the end event. Note that this event will fire almost

src/eval.c

@@ -146,9 +146,9 @@ void freeEvalScripts(dict *scripts, list *scripts_lru_list, list *engine_callbac
         listIter *iter = listGetIterator(engine_callbacks, 0);
         listNode *node = NULL;
         while ((node = listNext(iter)) != NULL) {
-            callableLazyEvalReset *callback = listNodeValue(node);
+            callableLazyEnvReset *callback = listNodeValue(node);
             if (callback != NULL) {
-                callback->engineLazyEvalResetCallback(callback->context);
+                callback->engineLazyEnvResetCallback(callback->context);
                 zfree(callback);
             }
         }
@@ -159,7 +159,7 @@ void freeEvalScripts(dict *scripts, list *scripts_lru_list, list *engine_callbac
 
 static void resetEngineEvalEnvCallback(scriptingEngine *engine, void *context) {
     int async = context != NULL;
-    callableLazyEvalReset *callback = scriptingEngineCallResetEvalEnvFunc(engine, async);
+    callableLazyEnvReset *callback = scriptingEngineCallResetEnvFunc(engine, VMSE_EVAL, async);
 
     if (async) {
         list *callbacks = context;
@@ -174,7 +174,6 @@ void evalRelease(int async) {
         list *engine_callbacks = listCreate();
         scriptingEngineManagerForEachEngine(resetEngineEvalEnvCallback, engine_callbacks);
         freeEvalScriptsAsync(evalCtx.scripts, evalCtx.scripts_lru_list, engine_callbacks);
-
     } else {
         freeEvalScripts(evalCtx.scripts, evalCtx.scripts_lru_list, NULL);
         scriptingEngineManagerForEachEngine(resetEngineEvalEnvCallback, NULL);

src/functions.c

@@ -165,32 +165,62 @@ void functionsLibCtxClear(functionsLibCtx *lib_ctx, void(callback)(dict *)) {
     lib_ctx->cache_memory = 0;
 }
 
-void functionsLibCtxClearCurrent(int async, void(callback)(dict *)) {
+static void resetEngineOrCollectResetCallbacks(scriptingEngine *engine, void *context) {
+    int async = context != NULL;
+    callableLazyEnvReset *callback = scriptingEngineCallResetEnvFunc(engine, VMSE_FUNCTION, async);
+
+    if (async) {
+        list *callbacks = context;
+        listAddNodeTail(callbacks, callback);
+    }
+}
+
+void functionsLibCtxReleaseCurrent(int async, void(callback)(dict *)) {
     if (async) {
-        functionsLibCtx *old_l_ctx = curr_functions_lib_ctx;
-        curr_functions_lib_ctx = functionsLibCtxCreate();
-        freeFunctionsAsync(old_l_ctx);
+        list *engine_callbacks = listCreate();
+        scriptingEngineManagerForEachEngine(resetEngineOrCollectResetCallbacks, engine_callbacks);
+        freeFunctionsAsync(curr_functions_lib_ctx, engine_callbacks);
     } else {
-        functionsLibCtxClear(curr_functions_lib_ctx, callback);
+        functionsLibCtxFree(curr_functions_lib_ctx, callback, NULL);
+        scriptingEngineManagerForEachEngine(resetEngineOrCollectResetCallbacks, NULL);
     }
 }
 
 /* Free the given functions ctx */
 static void functionsLibCtxFreeGeneric(functionsLibCtx *functions_lib_ctx, int async) {
     if (async) {
-        freeFunctionsAsync(functions_lib_ctx);
+        freeFunctionsAsync(functions_lib_ctx, NULL);
     } else {
-        functionsLibCtxFree(functions_lib_ctx);
+        functionsLibCtxFree(functions_lib_ctx, NULL, NULL);
     }
 }
 
+void functionReset(int async, void(callback)(dict *)) {
+    functionsLibCtxReleaseCurrent(async, callback);
+    functionsInit();
+}
+
 /* Free the given functions ctx */
-void functionsLibCtxFree(functionsLibCtx *functions_lib_ctx) {
-    functionsLibCtxClear(functions_lib_ctx, NULL);
+void functionsLibCtxFree(functionsLibCtx *functions_lib_ctx, void(callback)(dict *), list *engine_callbacks) {
+    functionsLibCtxClear(functions_lib_ctx, callback);
     dictRelease(functions_lib_ctx->functions);
     dictRelease(functions_lib_ctx->libraries);
     dictRelease(functions_lib_ctx->engines_stats);
     zfree(functions_lib_ctx);
+
+    if (engine_callbacks) {
+        listIter *iter = listGetIterator(engine_callbacks, 0);
+        listNode *node = NULL;
+        while ((node = listNext(iter)) != NULL) {
+            callableLazyEnvReset *engine_callback = listNodeValue(node);
+            if (engine_callback != NULL) {
+                engine_callback->engineLazyEnvResetCallback(engine_callback->context);
+                zfree(engine_callback);
+            }
+        }
+        listReleaseIterator(iter);
+        listRelease(engine_callbacks);
+    }
 }
 
 /* Swap the current functions ctx with the given one.
@@ -824,7 +854,7 @@ void functionFlushCommand(client *c) {
         return;
     }
 
-    functionsLibCtxClearCurrent(async, NULL);
+    functionReset(async, NULL);
 
     /* Indicate that the command changed the data so it will be replicated and
      * counted as a data change (for persistence configuration) */

src/functions.h

@@ -90,10 +90,10 @@ dict *functionsLibGet(void);
 size_t functionsLibCtxFunctionsLen(functionsLibCtx *functions_ctx);
 functionsLibCtx *functionsLibCtxGetCurrent(void);
 functionsLibCtx *functionsLibCtxCreate(void);
-void functionsLibCtxClearCurrent(int async, void(callback)(dict *));
-void functionsLibCtxFree(functionsLibCtx *functions_lib_ctx);
+void functionsLibCtxFree(functionsLibCtx *functions_lib_ctx, void(callback)(dict *), list *engine_callbacks);
 void functionsLibCtxClear(functionsLibCtx *lib_ctx, void(callback)(dict *));
 void functionsLibCtxSwapWithCurrent(functionsLibCtx *new_lib_ctx, int async);
+void functionReset(int async, void(callback)(dict *));
 
 void functionsRemoveLibFromEngine(scriptingEngine *engine);
 

src/lazyfree.c

@@ -66,8 +66,9 @@ void lazyFreeEvalScripts(void *args[]) {
 /* Release the functions ctx. */
 void lazyFreeFunctionsCtx(void *args[]) {
     functionsLibCtx *functions_lib_ctx = args[0];
+    list *engine_callbacks = args[1];
     size_t len = functionsLibCtxFunctionsLen(functions_lib_ctx);
-    functionsLibCtxFree(functions_lib_ctx);
+    functionsLibCtxFree(functions_lib_ctx, NULL, engine_callbacks);
     atomic_fetch_sub_explicit(&lazyfree_objects, len, memory_order_relaxed);
     atomic_fetch_add_explicit(&lazyfreed_objects, len, memory_order_relaxed);
 }
@@ -239,13 +240,13 @@ void freeEvalScriptsAsync(dict *scripts, list *scripts_lru_list, list *engine_ca
 }
 
 /* Free functions ctx, if the functions ctx contains enough functions, free it in async way. */
-void freeFunctionsAsync(functionsLibCtx *functions_lib_ctx) {
+void freeFunctionsAsync(functionsLibCtx *functions_lib_ctx, list *engine_callbacks) {
     if (functionsLibCtxFunctionsLen(functions_lib_ctx) > LAZYFREE_THRESHOLD) {
         atomic_fetch_add_explicit(&lazyfree_objects, functionsLibCtxFunctionsLen(functions_lib_ctx),
                                   memory_order_relaxed);
-        bioCreateLazyFreeJob(lazyFreeFunctionsCtx, 1, functions_lib_ctx);
+        bioCreateLazyFreeJob(lazyFreeFunctionsCtx, 2, functions_lib_ctx, engine_callbacks);
     } else {
-        functionsLibCtxFree(functions_lib_ctx);
+        functionsLibCtxFree(functions_lib_ctx, NULL, engine_callbacks);
     }
 }
 

src/lua/engine_lua.c

@@ -14,11 +14,6 @@
 #define LUA_ENGINE_NAME "LUA"
 #define REGISTRY_ERROR_HANDLER_NAME "__ERROR_HANDLER__"
 
-typedef struct luaFunction {
-    lua_State *lua;   /* Pointer to the lua context where this function was created. Only used in EVAL context. */
-    int function_ref; /* Special ID that allows getting the Lua function object from the Lua registry */
-} luaFunction;
-
 typedef struct luaEngineCtx {
     lua_State *eval_lua;     /* The Lua interpreter for EVAL commands. We use just one for all EVAL calls */
     lua_State *function_lua; /* The Lua interpreter for FCALL commands. We use just one for all FCALL calls */
@@ -260,17 +255,9 @@ static void luaEngineFunctionCall(ValkeyModuleCtx *module_ctx,
     serverAssert(module_ctx == NULL);
 
     luaEngineCtx *lua_engine_ctx = (luaEngineCtx *)engine_ctx;
-    lua_State *lua = NULL;
-    int lua_function_ref = -1;
-
-    if (type == VMSE_EVAL) {
-        lua = lua_engine_ctx->eval_lua;
-        luaFunction *script = compiled_function->function;
-        lua_function_ref = script->function_ref;
-    } else {
-        lua = lua_engine_ctx->function_lua;
-        lua_function_ref = luaFunctionGetLuaFunctionRef(compiled_function);
-    }
+    lua_State *lua = type == VMSE_EVAL ? lua_engine_ctx->eval_lua : lua_engine_ctx->function_lua;
+    luaFunction *script = compiled_function->function;
+    int lua_function_ref = script->function_ref;
 
     /* Push the pcall error handler function on the stack. */
     lua_pushstring(lua, REGISTRY_ERROR_HANDLER_NAME);
@@ -290,10 +277,10 @@ static void luaEngineFunctionCall(ValkeyModuleCtx *module_ctx,
     lua_pop(lua, 1); /* Remove the error handler. */
 }
 
-static void resetEvalContext(void *context) {
-    lua_State *eval_lua = context;
-    lua_gc(eval_lua, LUA_GCCOLLECT, 0);
-    lua_close(eval_lua);
+static void resetLuaContext(void *context) {
+    lua_State *lua = context;
+    lua_gc(lua, LUA_GCCOLLECT, 0);
+    lua_close(lua);
 
 #if !defined(USE_LIBC)
     /* The lua interpreter may hold a lot of memory internally, and lua is
@@ -308,27 +295,30 @@ static void resetEvalContext(void *context) {
 #endif
 }
 
-static callableLazyEvalReset *luaEngineResetEvalEnv(ValkeyModuleCtx *module_ctx,
-                                                    engineCtx *engine_ctx,
-                                                    int async) {
+static callableLazyEnvReset *luaEngineResetEvalEnv(ValkeyModuleCtx *module_ctx,
+                                                   engineCtx *engine_ctx,
+                                                   subsystemType type,
+                                                   int async) {
     /* The lua engine is implemented in the core, and not in a Valkey Module */
     serverAssert(module_ctx == NULL);
 
     luaEngineCtx *lua_engine_ctx = (luaEngineCtx *)engine_ctx;
-    serverAssert(lua_engine_ctx->eval_lua);
-    callableLazyEvalReset *callback = NULL;
+    serverAssert(type == VMSE_EVAL || type == VMSE_FUNCTION);
+    lua_State *lua = type == VMSE_EVAL ? lua_engine_ctx->eval_lua : lua_engine_ctx->function_lua;
+    serverAssert(lua);
+    callableLazyEnvReset *callback = NULL;
 
     if (async) {
         callback = zcalloc(sizeof(*callback));
-        *callback = (callableLazyEvalReset){
-            .context = lua_engine_ctx->eval_lua,
-            .engineLazyEvalResetCallback = resetEvalContext,
+        *callback = (callableLazyEnvReset){
+            .context = lua,
+            .engineLazyEnvResetCallback = resetLuaContext,
         };
     } else {
-        resetEvalContext(lua_engine_ctx->eval_lua);
+        resetLuaContext(lua);
     }
 
-    initializeLuaState(lua_engine_ctx, VMSE_EVAL);
+    initializeLuaState(lua_engine_ctx, type);
 
     return callback;
 }
@@ -350,21 +340,21 @@ static void luaEngineFreeFunction(ValkeyModuleCtx *module_ctx,
                                   compiledFunction *compiled_function) {
     /* The lua engine is implemented in the core, and not in a Valkey Module */
     serverAssert(module_ctx == NULL);
+    serverAssert(type == VMSE_EVAL || type == VMSE_FUNCTION);
 
     luaEngineCtx *lua_engine_ctx = engine_ctx;
-    if (type == VMSE_EVAL) {
-        luaFunction *script = (luaFunction *)compiled_function->function;
-        if (lua_engine_ctx->eval_lua == script->lua) {
-            /* The lua context is still the same, which means that we're not
-             * resetting the whole eval context, and therefore, we need to
-             * delete the function from the lua context.
-             */
-            lua_unref(lua_engine_ctx->eval_lua, script->function_ref);
-        }
-        zfree(script);
-    } else {
-        luaFunctionFreeFunction(lua_engine_ctx->function_lua, compiled_function->function);
+    lua_State *lua = type == VMSE_EVAL ? lua_engine_ctx->eval_lua : lua_engine_ctx->function_lua;
+    serverAssert(lua);
+
+    luaFunction *script = (luaFunction *)compiled_function->function;
+    if (lua == script->lua) {
+        /* The lua context is still the same, which means that we're not
+         * resetting the whole eval context, and therefore, we need to
+         * delete the function from the lua context.
+         */
+        lua_unref(lua, script->function_ref);
     }
+    zfree(script);
 
     if (compiled_function->name) {
         decrRefCount(compiled_function->name);
@@ -379,11 +369,12 @@ int luaEngineInitEngine(void) {
     ldbInit();
 
     engineMethods methods = {
+        .version = VALKEYMODULE_SCRIPTING_ENGINE_ABI_VERSION,
         .compile_code = luaEngineCompileCode,
         .free_function = luaEngineFreeFunction,
         .call_function = luaEngineFunctionCall,
         .get_function_memory_overhead = luaEngineFunctionMemoryOverhead,
-        .reset_eval_env = luaEngineResetEvalEnv,
+        .reset_env = luaEngineResetEvalEnv,
         .get_memory_info = luaEngineGetMemoryInfo,
     };
 

src/lua/engine_lua.h

@@ -4,6 +4,11 @@
 #include "../scripting_engine.h"
 #include <lua.h>
 
+typedef struct luaFunction {
+    lua_State *lua;   /* Pointer to the lua context where this function was created. Only used in EVAL context. */
+    int function_ref; /* Special ID that allows getting the Lua function object from the Lua registry */
+} luaFunction;
+
 int luaEngineInitEngine(void);
 
 #endif /* _ENGINE_LUA_ */