feat(server): introduce oom_deny_commands flag (#3718)

* server: introduce oom_deny_commands flag Signed-off-by: adi_holden <[email protected]>
dragonflydb · Sep 22, 2024 · 5cf9178 · 5cf9178
1 parent 69db21d
commit 5cf9178
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 0 deletions.
diff --git a/src/facade/command_id.h b/src/facade/command_id.h
@@ -84,6 +84,10 @@ class CommandId {
     restricted_ = restricted;
   }
 
+  void SetFlag(uint32_t flag) {
+    opt_mask_ |= flag;
+  }
+
   static uint32_t OptCount(uint32_t mask);
 
   // PUBLISH/SUBSCRIBE/UNSUBSCRIBE variant

diff --git a/src/server/command_registry.cc b/src/server/command_registry.cc
@@ -25,6 +25,8 @@ ABSL_FLAG(vector<string>, rename_command, {},
 ABSL_FLAG(vector<string>, restricted_commands, {},
           "Commands restricted to connections on the admin port");
 
+ABSL_FLAG(vector<string>, oom_deny_commands, {},
+          "Additinal commands that will be marked as denyoom");
 namespace dfly {
 
 using namespace facade;
@@ -108,6 +110,10 @@ CommandRegistry::CommandRegistry() {
   for (string name : GetFlag(FLAGS_restricted_commands)) {
     restricted_cmds_.emplace(AsciiStrToUpper(name));
   }
+
+  for (string name : GetFlag(FLAGS_oom_deny_commands)) {
+    oomdeny_cmds_.emplace(AsciiStrToUpper(name));
+  }
 }
 
 void CommandRegistry::Init(unsigned int thread_count) {
@@ -133,6 +139,10 @@ CommandRegistry& CommandRegistry::operator<<(CommandId cmd) {
     cmd.SetRestricted(true);
   }
 
+  if (oomdeny_cmds_.find(k) != oomdeny_cmds_.end()) {
+    cmd.SetFlag(CO::DENYOOM);
+  }
+
   cmd.SetFamily(family_of_commands_.size() - 1);
   if (!is_sub_command || absl::StartsWith(cmd.name(), "ACL")) {
     cmd.SetBitIndex(1ULL << bit_index_);

diff --git a/src/server/command_registry.h b/src/server/command_registry.h
@@ -194,6 +194,7 @@ class CommandRegistry {
   absl::flat_hash_map<std::string, CommandId> cmd_map_;
   absl::flat_hash_map<std::string, std::string> cmd_rename_map_;
   absl::flat_hash_set<std::string> restricted_cmds_;
+  absl::flat_hash_set<std::string> oomdeny_cmds_;
 
   FamiliesVec family_of_commands_;
   size_t bit_index_;

diff --git a/tests/dragonfly/generic_test.py b/tests/dragonfly/generic_test.py
@@ -138,3 +138,29 @@ async def test_reply_guard_oom(df_factory, df_seeder_factory):
 
     info = await c_master.info("stats")
     assert info["evicted_keys"] > 0, "Weak testcase: policy based eviction was not triggered."
+
+
+@pytest.mark.asyncio
+async def test_denyoom_commands(df_factory):
+    df_server = df_factory.create(
+        proactor_threads=1, maxmemory="256mb", oom_deny_commands="get", oom_deny_ratio=0.7
+    )
+    df_server.start()
+    client = df_server.client()
+    await client.execute_command("DEBUG POPULATE 7000 size 44000")
+
+    min_deny = 250 * 1024 * 1024  # 250mb
+    info = await client.info("memory")
+    print(f'Used memory {info["used_memory"]}, rss {info["used_memory_rss"]}')
+    assert info["used_memory"] > min_deny, "Weak testcase: too little used memory"
+
+    # reject set due to oom
+    with pytest.raises(redis.exceptions.ResponseError):
+        await client.execute_command("set x y")
+
+    # reject get because it is set in oom_deny_commands
+    with pytest.raises(redis.exceptions.ResponseError):
+        await client.execute_command("get x")
+
+    # mget should not be rejected
+    await client.execute_command("mget x")