From 5cf917871c416d4e34049aa58fb2ed9d7f60bf91 Mon Sep 17 00:00:00 2001
From: adiholden <adi@dragonflydb.io>
Date: Sun, 22 Sep 2024 09:32:18 +0300
Subject: [PATCH] feat(server): introduce oom_deny_commands flag (#3718)

* server: introduce oom_deny_commands flag

Signed-off-by: adi_holden <adi@dragonflydb.io>
---
 src/facade/command_id.h         |  4 ++++
 src/server/command_registry.cc  | 10 ++++++++++
 src/server/command_registry.h   |  1 +
 tests/dragonfly/generic_test.py | 26 ++++++++++++++++++++++++++
 4 files changed, 41 insertions(+)

diff --git a/src/facade/command_id.h b/src/facade/command_id.h
index 085d88385ddf..9d4dc82563ed 100644
--- 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
index 19769824c875..aad459e08d80 100644
--- 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
index 408d5f7caf9b..82ce102a785d 100644
--- 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
index 2de1cab1e5f9..b6648d26df44 100644
--- 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")