diff --git a/src/module.c b/src/module.c index f03743681..9dbaa3ca2 100644 --- a/src/module.c +++ b/src/module.c @@ -5654,6 +5654,7 @@ void RM_SetContextUser(RedisModuleCtx *ctx, const RedisModuleUser *user) { * "3" -> REDISMODULE_ARGV_RESP_3 * "0" -> REDISMODULE_ARGV_RESP_AUTO * "C" -> REDISMODULE_ARGV_RUN_AS_USER + * "M" -> REDISMODULE_ARGV_RESPECT_DENY_OOM * * On error (format specifier error) NULL is returned and nothing is * allocated. On success the argument vector is returned. */ @@ -5934,6 +5935,9 @@ RedisModuleCallReply *RM_Call(RedisModuleCtx *ctx, const char *cmdname, const ch goto cleanup; } } + } else { + /* if we aren't OOM checking in RM_Call, we want further executions from this client to also not fail on OOM */ + c->flags |= CLIENT_ALLOW_OOM; } if (flags & REDISMODULE_ARGV_NO_WRITES) { diff --git a/src/script.c b/src/script.c index 3d9785312..57fa7eeba 100644 --- a/src/script.c +++ b/src/script.c @@ -130,6 +130,7 @@ uint64_t scriptFlagsToCmdFlags(uint64_t cmd_flags, uint64_t script_flags) { /* Prepare the given run ctx for execution */ int scriptPrepareForRun(scriptRunCtx *run_ctx, client *engine_client, client *caller, const char *funcname, uint64_t script_flags, int ro) { serverAssert(!curr_run_ctx); + bool client_allow_oom = caller->flags & CLIENT_ALLOW_OOM; int running_stale = server.masterhost && server.repl_state != REPL_STATE_CONNECTED && @@ -189,7 +190,7 @@ int scriptPrepareForRun(scriptRunCtx *run_ctx, client *engine_client, client *ca /* Check OOM state. the no-writes flag imply allow-oom. we tested it * after the no-write error, so no need to mention it in the error reply. */ - if (server.pre_command_oom_state && server.maxmemory && + if (!client_allow_oom && server.pre_command_oom_state && server.maxmemory && !(script_flags & (SCRIPT_FLAG_ALLOW_OOM|SCRIPT_FLAG_NO_WRITES))) { addReplyError(caller, "-OOM allow-oom flag is not set on the script, " @@ -233,7 +234,7 @@ int scriptPrepareForRun(scriptRunCtx *run_ctx, client *engine_client, client *ca * flag, we will not allow write commands. */ run_ctx->flags |= SCRIPT_READ_ONLY; } - if (!(script_flags & SCRIPT_FLAG_EVAL_COMPAT_MODE) && (script_flags & SCRIPT_FLAG_ALLOW_OOM)) { + if (client_allow_oom || (!(script_flags & SCRIPT_FLAG_EVAL_COMPAT_MODE) && (script_flags & SCRIPT_FLAG_ALLOW_OOM))) { /* Note: we don't need to test the no-writes flag here and set this run_ctx flag, * since only write commands can are deny-oom. */ run_ctx->flags |= SCRIPT_ALLOW_OOM; diff --git a/src/server.h b/src/server.h index 092ee5758..80e07fb9d 100644 --- a/src/server.h +++ b/src/server.h @@ -369,6 +369,8 @@ extern int configOOMScoreAdjValuesDefaults[CONFIG_OOM_COUNT]; RDB without replication buffer. */ #define CLIENT_NO_EVICT (1ULL<<43) /* This client is protected against client memory eviction. */ +#define CLIENT_ALLOW_OOM (1ULL<<44) /* Client used by RM_Call is allowed to fully execute + scripts even when in OOM */ /* Client block type (btype field in client structure) * if CLIENT_BLOCKED flag is set. */ diff --git a/tests/unit/moduleapi/misc.tcl b/tests/unit/moduleapi/misc.tcl index 15e4f4c6c..0f076c6b5 100644 --- a/tests/unit/moduleapi/misc.tcl +++ b/tests/unit/moduleapi/misc.tcl @@ -245,29 +245,40 @@ start_server {tags {"modules"}} { } } - test {rm_call EVAL - OOM} { + # Note: each script is unique, to check that flags are extracted correctly + test {rm_call EVAL - OOM - with M flag} { r config set maxmemory 1 - assert_error {OOM command not allowed when used memory > 'maxmemory'. script*} { - r test.rm_call eval { + # script without shebang, but uses SET, so fails + assert_error {*OOM command not allowed when used memory > 'maxmemory'*} { + r test.rm_call_flags M eval { redis.call('set','x',1) return 1 } 1 x } - r test.rm_call eval {#!lua flags=no-writes + # script with an allow-oom flag, succeeds despite using SET + r test.rm_call_flags M eval {#!lua flags=allow-oom + redis.call('set','x', 1) + return 2 + } 1 x + + # script with no-writes flag, implies allow-oom, succeeds + r test.rm_call_flags M eval {#!lua flags=no-writes redis.call('get','x') return 2 } 1 x - assert_error {OOM allow-oom flag is not set on the script,*} { - r test.rm_call eval {#!lua + # script with shebang using default flags, so fails regardless of using only GET + assert_error {*OOM command not allowed when used memory > 'maxmemory'*} { + r test.rm_call_flags M eval {#!lua redis.call('get','x') return 3 } 1 x } - r test.rm_call eval { + # script without shebang, but uses GET, so succeeds + r test.rm_call_flags M eval { redis.call('get','x') return 4 } 1 x @@ -275,6 +286,31 @@ start_server {tags {"modules"}} { r config set maxmemory 0 } {OK} {needs:config-maxmemory} + # All RM_Call for script succeeds in OOM state without using the M flag + test {rm_call EVAL - OOM - without M flag} { + r config set maxmemory 1 + + # no shebang at all + r test.rm_call eval { + redis.call('set','x',1) + return 6 + } 1 x + + # Shebang without flags + r test.rm_call eval {#!lua + redis.call('set','x', 1) + return 7 + } 1 x + + # with allow-oom flag + r test.rm_call eval {#!lua flags=allow-oom + redis.call('set','x', 1) + return 8 + } 1 x + + r config set maxmemory 0 + } {OK} {needs:config-maxmemory} + test "not enough good replicas" { r set x "some value" r config set min-replicas-to-write 1