add event count throttle

Author: kostasrim

src/server/main_service.cc

@@ -1327,7 +1327,6 @@ bool Service::InvokeCmd(const CommandId* cid, CmdArgList tail_args, SinkReplyBui
 
   ServerState::tlocal()->RecordCmd(cntx->has_main_or_memcache_listener);
   Transaction* tx = cntx->transaction;
-
   auto& info = cntx->conn_state.tracking_info_;
   const bool is_read_only = cid->opt_mask() & CO::READONLY;
   if (tx) {

src/server/multi_command_squasher.cc

@@ -69,6 +69,8 @@ atomic_uint64_t MultiCommandSquasher::current_reply_size_ = 0;
 thread_local size_t MultiCommandSquasher::throttle_size_limit_ =
     absl::GetFlag(FLAGS_throttle_squashed);
 
+thread_local util::fb2::EventCount MultiCommandSquasher::ec_;
+
 MultiCommandSquasher::MultiCommandSquasher(absl::Span<StoredCmd> cmds, ConnectionContext* cntx,
                                            Service* service, bool verify_commands, bool error_abort)
     : cmds_{cmds},
@@ -207,12 +209,7 @@ OpStatus MultiCommandSquasher::SquashedHopCb(EngineShard* es, RespVersion resp_v
 
     local_tx->InitByArgs(cntx_->ns, local_cntx.conn_state.db_index, args);
 
-    if (MultiCommandSquasher::IsMultiCommandSquasherOverLimit()) {
-      auto* rb = (facade::RedisReplyBuilder*)&crb;
-      rb->SendError("High Rss");
-    } else {
-      service_->InvokeCmd(cmd->Cid(), args, &crb, &local_cntx);
-    }
+    service_->InvokeCmd(cmd->Cid(), args, &crb, &local_cntx);
 
     sinfo.replies.emplace_back(crb.Take());
     current_reply_size_.fetch_add(Size(sinfo.replies.back()), std::memory_order_relaxed);
@@ -233,6 +230,9 @@ bool MultiCommandSquasher::ExecuteSquashed(facade::RedisReplyBuilder* rb) {
   if (order_.empty())
     return true;
 
+  MultiCommandSquasher::ec_.await(
+      []() { return !MultiCommandSquasher::IsMultiCommandSquasherOverLimit(); });
+
   unsigned num_shards = 0;
   for (auto& sd : sharded_) {
     sd.replies.reserve(sd.cmds.size());
@@ -287,6 +287,7 @@ bool MultiCommandSquasher::ExecuteSquashed(facade::RedisReplyBuilder* rb) {
       break;
   }
   current_reply_size_.fetch_sub(size, std::memory_order_relaxed);
+  MultiCommandSquasher::ec_.notifyAll();
 
   uint64_t after_reply = proactor->GetMonotonicTimeNs();
   ServerState::SafeTLocal()->stats.multi_squash_exec_hop_usec += (after_hop - start) / 1000;

src/server/multi_command_squasher.h

@@ -7,6 +7,7 @@
 #include "facade/reply_capture.h"
 #include "server/conn_context.h"
 #include "server/main_service.h"
+#include "util/fibers/synchronization.h"
 
 namespace dfly {
 
@@ -33,9 +34,12 @@ class MultiCommandSquasher {
   }
 
   static bool IsMultiCommandSquasherOverLimit() {
-    LOG(INFO) << throttle_size_limit_ << " current reply suze " << current_reply_size_;
-    return throttle_size_limit_ > 0 &&
-           current_reply_size_.load(std::memory_order_relaxed) > throttle_size_limit_;
+    const bool over_limit =
+        throttle_size_limit_ > 0 &&
+        current_reply_size_.load(std::memory_order_relaxed) > throttle_size_limit_;
+    VLOG_IF(2, over_limit) << "MultiCommandSquasher overlimit: " << throttle_size_limit_
+                           << " current reply size " << current_reply_size_;
+    return over_limit;
   }
 
  private:
@@ -99,6 +103,8 @@ class MultiCommandSquasher {
   // we increase size in one thread and decrease in another
   static atomic_uint64_t current_reply_size_;
   static thread_local size_t throttle_size_limit_;
+  // Used to throttle when memory is tight
+  static thread_local util::fb2::EventCount ec_;
 };
 
 }  // namespace dfly

tests/dragonfly/memory_test.py

@@ -1,4 +1,5 @@
 import pytest
+import asyncio
 from redis import asyncio as aioredis
 from .utility import *
 import logging
@@ -226,18 +227,36 @@ async def test_cache_eviction_with_rss_deny_oom(
 
 @pytest.mark.asyncio
 async def test_throttle_on_commands_squashing_replies_bytes(df_factory: DflyInstanceFactory):
-    df = df_factory.create(proactor_threads=2, throttle_squashed=1000000000)
+    df = df_factory.create(
+        proactor_threads=2, throttle_squashed=1_000_000_000, vmodule="multi_command_squasher=2"
+    )
     df.start()
 
     client = df.client()
     # 1gb
     await client.execute_command("debug populate 1 test 10000 rand type hash elements 100000")
 
-    await client.execute_command("multi")
-    await client.execute_command("hgetall test:0")
-    await client.execute_command("hgetall test:0")
-    await client.execute_command("hgetall test:0")
-    await client.execute_command("hgetall test:0")
-    res = await client.execute_command("exec")
-    assert type(res[1]) is redis.exceptions.ResponseError
-    assert type(res[2]) is redis.exceptions.ResponseError
+    async def poll():
+        # At any point we should not cross this limit
+        cl = df.client()
+        await cl.execute_command("multi")
+        await cl.execute_command("hgetall test:0")
+        await cl.execute_command("exec")
+
+    # With the current approach this will overshoot
+    #    await client.execute_command("multi")
+    #    await client.execute_command("hgetall test:0")
+    #    await client.execute_command("hgetall test:0")
+    #    await client.execute_command("hgetall test:0")
+    #    await client.execute_command("hgetall test:0")
+    #    res = await client.execute_command("exec")
+    tasks = []
+    for i in range(50):
+        tasks.append(asyncio.create_task(poll()))
+
+    for task in tasks:
+        await task
+
+    df.stop()
+    found = df.find_in_logs("MultiCommandSquasher overlimit: ")
+    assert len(found) > 0