fix: Add selective non-atomic squash dispatch (#1641)

fix: Add selecting non-atomic squash dispatch

Signed-off-by: Vladislav Oleshko <[email protected]>

src/server/engine_shard_set.h

@@ -279,11 +279,17 @@ class EngineShardSet {
     RunBriefInParallel(std::forward<U>(func), [](auto i) { return true; });
   }
 
-  // Runs a brief function on selected shard thread. Waits for it to complete.
+  // Runs a brief function on selected shards. Waits for it to complete.
   // `func` must not preempt.
   template <typename U, typename P> void RunBriefInParallel(U&& func, P&& pred) const;
 
-  template <typename U> void RunBlockingInParallel(U&& func);
+  // Runs a possibly blocking function on all shards. Waits for it to complete.
+  template <typename U> void RunBlockingInParallel(U&& func) {
+    RunBlockingInParallel(std::forward<U>(func), [](auto i) { return true; });
+  }
+
+  // Runs a possibly blocking function on selected shards. Waits for it to complete.
+  template <typename U, typename P> void RunBlockingInParallel(U&& func, P&& pred);
 
   // Runs func on all shards via the same shard queue that's been used by transactions framework.
   // The functions running inside the shard queue run atomically (sequentially)
@@ -329,14 +335,18 @@ void EngineShardSet::RunBriefInParallel(U&& func, P&& pred) const {
   bc.Wait();
 }
 
-template <typename U> void EngineShardSet::RunBlockingInParallel(U&& func) {
-  BlockingCounter bc{size()};
+template <typename U, typename P> void EngineShardSet::RunBlockingInParallel(U&& func, P&& pred) {
+  BlockingCounter bc{0};
   static_assert(std::is_invocable_v<U, EngineShard*>,
                 "Argument must be invocable EngineShard* as argument.");
   static_assert(std::is_void_v<std::invoke_result_t<U, EngineShard*>>,
                 "Callable must not have a return value!");
 
   for (uint32_t i = 0; i < size(); ++i) {
+    if (!pred(i))
+      continue;
+
+    bc.Add(1);
     util::ProactorBase* dest = pp_->at(i);
 
     // the "Dispatch" call spawns a fiber underneath.

src/server/multi_command_squasher.cc

@@ -155,10 +155,8 @@ bool MultiCommandSquasher::ExecuteSquashed() {
     tx->PrepareSquashedMultiHop(base_cid_, cb);
     tx->ScheduleSingleHop([this](auto* tx, auto* es) { return SquashedHopCb(tx, es); });
   } else {
-    shard_set->RunBlockingInParallel([this, tx](auto* es) {
-      if (!sharded_[es->shard_id()].cmds.empty())
-        SquashedHopCb(tx, es);
-    });
+    shard_set->RunBlockingInParallel([this, tx](auto* es) { SquashedHopCb(tx, es); },
+                                     [this](auto sid) { return !sharded_[sid].cmds.empty(); });
   }
 
   bool aborted = false;