diff --git a/src/function_lua.c b/src/function_lua.c index d865f31bb3..b535528906 100644 --- a/src/function_lua.c +++ b/src/function_lua.c @@ -228,13 +228,17 @@ static void luaEngineCall(ValkeyModuleCtx *module_ctx, lua_pop(lua, 1); /* Pop error handler */ } -static size_t luaEngineGetUsedMemoy(ValkeyModuleCtx *module_ctx, - engineCtx *engine_ctx) { +static engineMemoryInfo luaEngineGetMemoryInfo(ValkeyModuleCtx *module_ctx, + engineCtx *engine_ctx) { /* The lua engine is implemented in the core, and not in a Valkey Module */ serverAssert(module_ctx == NULL); luaEngineCtx *lua_engine_ctx = engine_ctx; - return luaMemory(lua_engine_ctx->lua); + + return (engineMemoryInfo){ + .used_memory = luaMemory(lua_engine_ctx->lua), + .engine_memory_overhead = zmalloc_size(lua_engine_ctx), + }; } static size_t luaEngineFunctionMemoryOverhead(ValkeyModuleCtx *module_ctx, @@ -245,15 +249,6 @@ static size_t luaEngineFunctionMemoryOverhead(ValkeyModuleCtx *module_ctx, return zmalloc_size(compiled_function); } -static size_t luaEngineMemoryOverhead(ValkeyModuleCtx *module_ctx, - engineCtx *engine_ctx) { - /* The lua engine is implemented in the core, and not in a Valkey Module */ - serverAssert(module_ctx == NULL); - - luaEngineCtx *lua_engine_ctx = engine_ctx; - return zmalloc_size(lua_engine_ctx); -} - static void luaEngineFreeFunction(ValkeyModuleCtx *module_ctx, engineCtx *engine_ctx, void *compiled_function) { @@ -557,10 +552,9 @@ int luaEngineInitEngine(void) { .version = VALKEYMODULE_SCRIPTING_ENGINE_ABI_VERSION, .create_functions_library = luaEngineCreate, .call_function = luaEngineCall, - .get_used_memory = luaEngineGetUsedMemoy, .get_function_memory_overhead = luaEngineFunctionMemoryOverhead, - .get_engine_memory_overhead = luaEngineMemoryOverhead, .free_function = luaEngineFreeFunction, + .get_memory_info = luaEngineGetMemoryInfo, }; return functionsRegisterEngine(LUA_ENGINE_NAME, diff --git a/src/functions.c b/src/functions.c index 6d5f423f77..9a2765c2c1 100644 --- a/src/functions.c +++ b/src/functions.c @@ -476,10 +476,9 @@ int functionsRegisterEngine(const char *engine_name, .engine_ctx = engine_ctx, .create = engine_methods->create_functions_library, .call = engine_methods->call_function, - .get_used_memory = engine_methods->get_used_memory, .get_function_memory_overhead = engine_methods->get_function_memory_overhead, - .get_engine_memory_overhead = engine_methods->get_engine_memory_overhead, .free_function = engine_methods->free_function, + .get_memory_info = engine_methods->get_memory_info, }; client *c = createClient(NULL); @@ -500,9 +499,13 @@ int functionsRegisterEngine(const char *engine_name, functionsAddEngineStats(ei); setupEngineModuleCtx(ei, NULL); - engine_cache_memory += zmalloc_size(ei) + sdsAllocSize(ei->name) + + engineMemoryInfo mem_info = eng->get_memory_info(ei->module_ctx, + eng->engine_ctx); + engine_cache_memory += zmalloc_size(ei) + + sdsAllocSize(ei->name) + zmalloc_size(eng) + - eng->get_engine_memory_overhead(ei->module_ctx, eng->engine_ctx); + mem_info.engine_memory_overhead; + teardownEngineModuleCtx(ei); return C_OK; @@ -1265,8 +1268,9 @@ unsigned long functionsMemory(void) { engineInfo *ei = dictGetVal(entry); engine *engine = ei->engine; setupEngineModuleCtx(ei, NULL); - engines_memory += engine->get_used_memory(ei->module_ctx, - engine->engine_ctx); + engineMemoryInfo mem_info = engine->get_memory_info(ei->module_ctx, + engine->engine_ctx); + engines_memory += mem_info.used_memory; teardownEngineModuleCtx(ei); } dictReleaseIterator(iter); diff --git a/src/functions.h b/src/functions.h index eab3566668..89e39fdc56 100644 --- a/src/functions.h +++ b/src/functions.h @@ -58,6 +58,7 @@ typedef struct functionLibInfo functionLibInfo; typedef ValkeyModuleScriptingEngineCtx engineCtx; typedef ValkeyModuleScriptingEngineFunctionCtx functionCtx; typedef ValkeyModuleScriptingEngineCompiledFunction compiledFunction; +typedef ValkeyModuleScriptingEngineMemoryInfo engineMemoryInfo; typedef ValkeyModuleScriptingEngineMethods engineMethods; typedef struct engine { @@ -90,9 +91,10 @@ typedef struct engine { robj **args, size_t nargs); - /* get current used memory by the engine */ - size_t (*get_used_memory)(ValkeyModuleCtx *module_ctx, - engineCtx *engine_ctx); + /* free the given function */ + void (*free_function)(ValkeyModuleCtx *module_ctx, + engineCtx *engine_ctx, + void *compiled_function); /* Return memory overhead for a given function, * such memory is not counted as engine memory but as general @@ -100,14 +102,10 @@ typedef struct engine { size_t (*get_function_memory_overhead)(ValkeyModuleCtx *module_ctx, void *compiled_function); - /* Return memory overhead for engine (struct size holding the engine)*/ - size_t (*get_engine_memory_overhead)(ValkeyModuleCtx *module_ctx, - engineCtx *engine_ctx); + /* Get the current used memory by the engine */ + engineMemoryInfo (*get_memory_info)(ValkeyModuleCtx *module_ctx, + engineCtx *engine_ctx); - /* free the given function */ - void (*free_function)(ValkeyModuleCtx *module_ctx, - engineCtx *engine_ctx, - void *compiled_function); } engine; /* Hold information about an engine. diff --git a/src/valkeymodule.h b/src/valkeymodule.h index 5ba04e5e08..1d99d2ff7a 100644 --- a/src/valkeymodule.h +++ b/src/valkeymodule.h @@ -817,6 +817,15 @@ typedef struct ValkeyModuleScriptingEngineCompiledFunction { uint64_t f_flags; /* Function flags */ } ValkeyModuleScriptingEngineCompiledFunction; +/* This struct is used to return the memory information of the scripting + * engine. */ +typedef struct ValkeyModuleScriptingEngineMemoryInfo { + /* The memory used by the scripting engine runtime. */ + size_t used_memory; + /* The memory used by the scripting engine data structures. */ + size_t engine_memory_overhead; +} ValkeyModuleScriptingEngineMemoryInfo; + typedef ValkeyModuleScriptingEngineCompiledFunction **(*ValkeyModuleScriptingEngineCreateFunctionsLibraryFunc)( ValkeyModuleCtx *module_ctx, ValkeyModuleScriptingEngineCtx *engine_ctx, @@ -835,23 +844,19 @@ typedef void (*ValkeyModuleScriptingEngineCallFunctionFunc)( ValkeyModuleString **args, size_t nargs); -typedef size_t (*ValkeyModuleScriptingEngineGetUsedMemoryFunc)( - ValkeyModuleCtx *module_ctx, - ValkeyModuleScriptingEngineCtx *engine_ctx); - typedef size_t (*ValkeyModuleScriptingEngineGetFunctionMemoryOverheadFunc)( ValkeyModuleCtx *module_ctx, void *compiled_function); -typedef size_t (*ValkeyModuleScriptingEngineGetEngineMemoryOverheadFunc)( - ValkeyModuleCtx *module_ctx, - ValkeyModuleScriptingEngineCtx *engine_ctx); - typedef void (*ValkeyModuleScriptingEngineFreeFunctionFunc)( ValkeyModuleCtx *module_ctx, ValkeyModuleScriptingEngineCtx *engine_ctx, void *compiled_function); +typedef ValkeyModuleScriptingEngineMemoryInfo (*ValkeyModuleScriptingEngineGetMemoryInfoFunc)( + ValkeyModuleCtx *module_ctx, + ValkeyModuleScriptingEngineCtx *engine_ctx); + typedef struct ValkeyModuleScriptingEngineMethodsV1 { uint64_t version; /* Version of this structure for ABI compat. */ @@ -870,11 +875,8 @@ typedef struct ValkeyModuleScriptingEngineMethodsV1 { /* Function callback to return memory overhead for a given function. */ ValkeyModuleScriptingEngineGetFunctionMemoryOverheadFunc get_function_memory_overhead; - /* Function callback to get current used memory by the engine. */ - ValkeyModuleScriptingEngineGetUsedMemoryFunc get_used_memory; - - /* Function callback to return memory overhead of the engine. */ - ValkeyModuleScriptingEngineGetEngineMemoryOverheadFunc get_engine_memory_overhead; + /* Function callback to get the used memory by the engine. */ + ValkeyModuleScriptingEngineGetMemoryInfoFunc get_memory_info; } ValkeyModuleScriptingEngineMethodsV1; diff --git a/tests/modules/helloscripting.c b/tests/modules/helloscripting.c index 77aa7ebe12..fdca6c8e91 100644 --- a/tests/modules/helloscripting.c +++ b/tests/modules/helloscripting.c @@ -235,29 +235,28 @@ static uint32_t executeHelloLangFunction(HelloFunc *func, return 0; } -static size_t engineGetUsedMemoy(ValkeyModuleCtx *module_ctx, - ValkeyModuleScriptingEngineCtx *engine_ctx) { +static ValkeyModuleScriptingEngineMemoryInfo engineGetMemoryInfo(ValkeyModuleCtx *module_ctx, + ValkeyModuleScriptingEngineCtx *engine_ctx) { VALKEYMODULE_NOT_USED(module_ctx); HelloLangCtx *ctx = (HelloLangCtx *)engine_ctx; - size_t memory = ValkeyModule_MallocSize(ctx); - memory += ValkeyModule_MallocSize(ctx->program); - for (uint32_t i = 0; i < ctx->program->num_functions; i++) { - HelloFunc *func = ctx->program->functions[i]; - memory += ValkeyModule_MallocSize(func); - memory += ValkeyModule_MallocSize(func->name); + ValkeyModuleScriptingEngineMemoryInfo mem_info = {0}; + + if (ctx->program != NULL) { + mem_info.used_memory += ValkeyModule_MallocSize(ctx->program); + + for (uint32_t i = 0; i < ctx->program->num_functions; i++) { + HelloFunc *func = ctx->program->functions[i]; + mem_info.used_memory += ValkeyModule_MallocSize(func); + mem_info.used_memory += ValkeyModule_MallocSize(func->name); + } } - return memory; -} -static size_t engineMemoryOverhead(ValkeyModuleCtx *module_ctx, - ValkeyModuleScriptingEngineCtx *engine_ctx) { - VALKEYMODULE_NOT_USED(module_ctx); - HelloLangCtx *ctx = (HelloLangCtx *)engine_ctx; - size_t overhead = ValkeyModule_MallocSize(engine_ctx); + mem_info.engine_memory_overhead = ValkeyModule_MallocSize(ctx); if (ctx->program != NULL) { - overhead += ValkeyModule_MallocSize(ctx->program); + mem_info.engine_memory_overhead += ValkeyModule_MallocSize(ctx->program); } - return overhead; + + return mem_info; } static size_t engineFunctionMemoryOverhead(ValkeyModuleCtx *module_ctx, @@ -356,10 +355,9 @@ int ValkeyModule_OnLoad(ValkeyModuleCtx *ctx, ValkeyModuleString **argv, .version = VALKEYMODULE_SCRIPTING_ENGINE_ABI_VERSION, .create_functions_library = createHelloLangEngine, .call_function = callHelloLangFunction, - .get_used_memory = engineGetUsedMemoy, .get_function_memory_overhead = engineFunctionMemoryOverhead, - .get_engine_memory_overhead = engineMemoryOverhead, .free_function = engineFreeFunction, + .get_memory_info = engineGetMemoryInfo, }; ValkeyModule_RegisterScriptingEngine(ctx,