Skip to content

Commit

Permalink
wrapped scripting engine callbacks calls with setup and teardown of t…
Browse files Browse the repository at this point in the history
…he module context object

Signed-off-by: Ricardo Dias <ricardo.dias@percona.com>
  • Loading branch information
rjd15372 committed Dec 19, 2024
1 parent 1687db1 commit 3cfea1c
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 14 deletions.
43 changes: 37 additions & 6 deletions src/functions.c
Original file line number Diff line number Diff line change
Expand Up @@ -118,11 +118,28 @@ static dict *engines = NULL;
/* Libraries Ctx. */
static functionsLibCtx *curr_functions_lib_ctx = NULL;

static void setupEngineModuleCtx(engineInfo *ei, client *c) {
if (ei->engineModule != NULL) {
serverAssert(ei->module_ctx != NULL);
moduleScriptingEngineInitContext(ei->module_ctx, ei->engineModule, c);
}
}

static void teardownEngineModuleCtx(engineInfo *ei) {
if (ei->engineModule != NULL) {
serverAssert(ei->module_ctx != NULL);
moduleFreeContext(ei->module_ctx);
}
}

static size_t functionMallocSize(functionInfo *fi) {
return zmalloc_size(fi) +
sdsAllocSize(fi->name) +
(fi->desc ? sdsAllocSize(fi->desc) : 0) +
fi->li->ei->engine->get_function_memory_overhead(fi->li->ei->module_ctx, fi->function);
setupEngineModuleCtx(fi->li->ei, NULL);
size_t size = zmalloc_size(fi) +
sdsAllocSize(fi->name) +
(fi->desc ? sdsAllocSize(fi->desc) : 0) +
fi->li->ei->engine->get_function_memory_overhead(fi->li->ei->module_ctx, fi->function);
teardownEngineModuleCtx(fi->li->ei);
return size;
}

static size_t libraryMallocSize(functionLibInfo *li) {
Expand All @@ -144,10 +161,12 @@ static void engineFunctionDispose(void *obj) {
if (fi->desc) {
sdsfree(fi->desc);
}
setupEngineModuleCtx(fi->li->ei, NULL);
engine *engine = fi->li->ei->engine;
engine->free_function(fi->li->ei->module_ctx,
engine->engine_ctx,
fi->function);
teardownEngineModuleCtx(fi->li->ei);
zfree(fi);
}

Expand Down Expand Up @@ -471,7 +490,7 @@ int functionsRegisterEngine(const char *engine_name,
*ei = (engineInfo){
.name = engine_name_sds,
.engineModule = engine_module,
.module_ctx = engine_module ? moduleAllocateScriptingEngineContext(engine_module) : NULL,
.module_ctx = engine_module ? moduleAllocateContext() : NULL,
.engine = eng,
.c = c,
};
Expand All @@ -480,9 +499,11 @@ int functionsRegisterEngine(const char *engine_name,

functionsAddEngineStats(ei);

setupEngineModuleCtx(ei, NULL);
engine_cache_memory += zmalloc_size(ei) + sdsAllocSize(ei->name) +
zmalloc_size(eng) +
eng->get_engine_memory_overhead(ei->module_ctx, eng->engine_ctx);
teardownEngineModuleCtx(ei);

return C_OK;
}
Expand Down Expand Up @@ -517,7 +538,6 @@ int functionsUnregisterEngine(const char *engine_name) {
freeClient(ei->c);
if (ei->engineModule != NULL) {
serverAssert(ei->module_ctx != NULL);
moduleFreeContext(ei->module_ctx);
zfree(ei->module_ctx);
}
zfree(ei);
Expand Down Expand Up @@ -742,6 +762,7 @@ static void fcallCommandGeneric(client *c, int ro) {

scriptRunCtx run_ctx;
if (scriptPrepareForRun(&run_ctx, fi->li->ei->c, c, fi->name, fi->f_flags, ro) != C_OK) return;
setupEngineModuleCtx(fi->li->ei, run_ctx.original_client);

engine->call(fi->li->ei->module_ctx,
engine->engine_ctx,
Expand All @@ -751,6 +772,8 @@ static void fcallCommandGeneric(client *c, int ro) {
numkeys,
c->argv + 3 + numkeys,
c->argc - 3 - numkeys);

teardownEngineModuleCtx(fi->li->ei);
scriptResetRun(&run_ctx);
}

Expand Down Expand Up @@ -1054,6 +1077,8 @@ static void freeCompiledFunctions(engineInfo *ei,
compiledFunction **compiled_functions,
size_t num_compiled_functions,
size_t free_function_from_idx) {
setupEngineModuleCtx(ei, NULL);

for (size_t i = 0; i < num_compiled_functions; i++) {
compiledFunction *func = compiled_functions[i];
decrRefCount(func->name);
Expand All @@ -1069,6 +1094,8 @@ static void freeCompiledFunctions(engineInfo *ei,
}

zfree(compiled_functions);

teardownEngineModuleCtx(ei);
}

/* Compile and save the given library, return the loaded library name on success
Expand Down Expand Up @@ -1111,13 +1138,15 @@ sds functionsCreateWithLibraryCtx(sds code, int replace, sds *err, functionsLibC
new_li = engineLibraryCreate(md.name, ei, code);
size_t num_compiled_functions = 0;
char *compile_error = NULL;
setupEngineModuleCtx(ei, NULL);
compiledFunction **compiled_functions =
engine->create(ei->module_ctx,
engine->engine_ctx,
md.code,
timeout,
&num_compiled_functions,
&compile_error);
teardownEngineModuleCtx(ei);
if (compiled_functions == NULL) {
serverAssert(num_compiled_functions == 0);
serverAssert(compile_error != NULL);
Expand Down Expand Up @@ -1235,8 +1264,10 @@ unsigned long functionsMemory(void) {
while ((entry = dictNext(iter))) {
engineInfo *ei = dictGetVal(entry);
engine *engine = ei->engine;
setupEngineModuleCtx(ei, NULL);
engines_memory += engine->get_used_memory(ei->module_ctx,
engine->engine_ctx);
teardownEngineModuleCtx(ei);
}
dictReleaseIterator(iter);

Expand Down
24 changes: 17 additions & 7 deletions src/module.c
Original file line number Diff line number Diff line change
Expand Up @@ -880,13 +880,13 @@ void moduleCallCommandUnblockedHandler(client *c) {
moduleReleaseTempClient(c);
}

/* Allocates the memory necessary to hold the ValkeyModuleCtx structure,
* initializes the context for use by the server to call scripting engines
* callback functions, and returns the pointer to the allocated memory. */
ValkeyModuleCtx *moduleAllocateScriptingEngineContext(ValkeyModule *module) {
ValkeyModuleCtx *ctx = (ValkeyModuleCtx *)zcalloc(sizeof(ValkeyModuleCtx));
moduleCreateContext(ctx, module, VALKEYMODULE_CTX_NEW_CLIENT);
return ctx;
/* Allocates the memory necessary to hold the ValkeyModuleCtx structure, and
* returns the pointer to the allocated memory.
*
* Used by the scripting engines implementation to cache the context structure.
*/
ValkeyModuleCtx *moduleAllocateContext(void) {
return (ValkeyModuleCtx *)zcalloc(sizeof(ValkeyModuleCtx));
}

/* Create a module ctx and keep track of the nesting level.
Expand Down Expand Up @@ -931,6 +931,16 @@ void moduleCreateContext(ValkeyModuleCtx *out_ctx, ValkeyModule *module, int ctx
}
}

/* Initialize a module context to be used by scripting engines callback
* functions.
*/
void moduleScriptingEngineInitContext(ValkeyModuleCtx *out_ctx,
ValkeyModule *module,
client *client) {
moduleCreateContext(out_ctx, module, VALKEYMODULE_CTX_NONE);
out_ctx->client = client;
}

/* This command binds the normal command invocation with commands
* exported by modules. */
void ValkeyModuleCommandDispatcher(client *c) {
Expand Down
5 changes: 4 additions & 1 deletion src/module.h
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,10 @@
typedef struct ValkeyModuleCtx ValkeyModuleCtx;
typedef struct ValkeyModule ValkeyModule;

ValkeyModuleCtx *moduleAllocateScriptingEngineContext(ValkeyModule *module);
ValkeyModuleCtx *moduleAllocateContext(void);
void moduleScriptingEngineInitContext(ValkeyModuleCtx *out_ctx,
ValkeyModule *module,
client *client);
void moduleFreeContext(ValkeyModuleCtx *ctx);

#endif /* _MODULE_H_ */

0 comments on commit 3cfea1c

Please sign in to comment.