From db67b35f8e276871e866efcc96445242b9732690 Mon Sep 17 00:00:00 2001
From: adiholden <adi@dragonflydb.io>
Date: Thu, 14 Nov 2024 20:57:05 +0200
Subject: [PATCH] fix server: fix stats of pipeline squashed commands (#4132)

Signed-off-by: adi_holden <adi@dragonflydb.io>
---
 src/server/main_service.cc           |  8 ++++----
 src/server/multi_command_squasher.cc |  3 ++-
 src/server/multi_command_squasher.h  | 12 ++++++------
 3 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/src/server/main_service.cc b/src/server/main_service.cc
index e4f08c25b38b..8787aa4c0e76 100644
--- a/src/server/main_service.cc
+++ b/src/server/main_service.cc
@@ -1407,13 +1407,13 @@ size_t Service::DispatchManyCommands(absl::Span<CmdArgList> args_list, SinkReply
     }
 
     dfly_cntx->transaction = dist_trans.get();
-    MultiCommandSquasher::Execute(absl::MakeSpan(stored_cmds),
-                                  static_cast<RedisReplyBuilder*>(builder), dfly_cntx, this, true,
-                                  false);
+    size_t squashed_num = MultiCommandSquasher::Execute(absl::MakeSpan(stored_cmds),
+                                                        static_cast<RedisReplyBuilder*>(builder),
+                                                        dfly_cntx, this, true, false);
     dfly_cntx->transaction = nullptr;
 
     dispatched += stored_cmds.size();
-    ss->stats.squashed_commands += stored_cmds.size();
+    ss->stats.squashed_commands += squashed_num;
     stored_cmds.clear();
   };
 
diff --git a/src/server/multi_command_squasher.cc b/src/server/multi_command_squasher.cc
index d36baf1dd220..9d40d59125ac 100644
--- a/src/server/multi_command_squasher.cc
+++ b/src/server/multi_command_squasher.cc
@@ -254,7 +254,7 @@ bool MultiCommandSquasher::ExecuteSquashed(facade::RedisReplyBuilder* rb) {
   return !aborted;
 }
 
-void MultiCommandSquasher::Run(RedisReplyBuilder* rb) {
+size_t MultiCommandSquasher::Run(RedisReplyBuilder* rb) {
   DVLOG(1) << "Trying to squash " << cmds_.size() << " commands for transaction "
            << cntx_->transaction->DebugId();
 
@@ -291,6 +291,7 @@ void MultiCommandSquasher::Run(RedisReplyBuilder* rb) {
 
   VLOG(1) << "Squashed " << num_squashed_ << " of " << cmds_.size()
           << " commands, max fanout: " << num_shards_ << ", atomic: " << atomic_;
+  return num_squashed_;
 }
 
 bool MultiCommandSquasher::IsAtomic() const {
diff --git a/src/server/multi_command_squasher.h b/src/server/multi_command_squasher.h
index 1a813c6133b8..889248818b35 100644
--- a/src/server/multi_command_squasher.h
+++ b/src/server/multi_command_squasher.h
@@ -22,10 +22,10 @@ namespace dfly {
 // contains a non-atomic multi transaction to execute squashed commands.
 class MultiCommandSquasher {
  public:
-  static void Execute(absl::Span<StoredCmd> cmds, facade::RedisReplyBuilder* rb,
-                      ConnectionContext* cntx, Service* service, bool verify_commands = false,
-                      bool error_abort = false) {
-    MultiCommandSquasher{cmds, cntx, service, verify_commands, error_abort}.Run(rb);
+  static size_t Execute(absl::Span<StoredCmd> cmds, facade::RedisReplyBuilder* rb,
+                        ConnectionContext* cntx, Service* service, bool verify_commands = false,
+                        bool error_abort = false) {
+    return MultiCommandSquasher{cmds, cntx, service, verify_commands, error_abort}.Run(rb);
   }
 
  private:
@@ -62,8 +62,8 @@ class MultiCommandSquasher {
   // Execute all currently squashed commands. Return false if aborting on error.
   bool ExecuteSquashed(facade::RedisReplyBuilder* rb);
 
-  // Run all commands until completion.
-  void Run(facade::RedisReplyBuilder* rb);
+  // Run all commands until completion. Returns number of squashed commands.
+  size_t Run(facade::RedisReplyBuilder* rb);
 
   bool IsAtomic() const;