From 7f9a13bb500d279e39cc29cd6b53a1116b090272 Mon Sep 17 00:00:00 2001
From: Kostas Kyrimis <kostas@dragonflydb.io>
Date: Tue, 11 Jun 2024 16:24:44 +0300
Subject: [PATCH] fix: lua and client tracking (#3163)

* add missing tracking_cb_ invkoe in RunSquashedMultiCb
* add test
---
 src/server/server_family_test.cc | 18 ++++++++++++++++++
 src/server/transaction.cc        |  5 ++---
 src/server/transaction.h         |  6 ++++++
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/src/server/server_family_test.cc b/src/server/server_family_test.cc
index cfea9ea78..d5afc01d5 100644
--- a/src/server/server_family_test.cc
+++ b/src/server/server_family_test.cc
@@ -6,6 +6,7 @@
 
 #include <absl/strings/match.h>
 
+#include "absl/strings/str_cat.h"
 #include "base/gtest.h"
 #include "base/logging.h"
 #include "facade/facade_test.h"
@@ -447,4 +448,21 @@ TEST_F(ServerFamilyTest, ClientTrackingNonTransactionalBug) {
 
   Run({"CLUSTER", "SLOTS"});
 }
+
+TEST_F(ServerFamilyTest, ClientTrackingLuaBug) {
+  Run({"HELLO", "3"});
+  // Check stickiness
+  Run({"CLIENT", "TRACKING", "ON"});
+  using namespace std::string_literals;
+  std::string eval = R"(redis.call('get', 'foo'); redis.call('set', 'foo', 'bar'); )";
+  Run({"EVAL", absl::StrCat(eval, "return 1"), "1", "foo"});
+  Run({"PING"});
+
+  EXPECT_EQ(InvalidationMessagesLen("IO0"), 1);
+  absl::StrAppend(&eval, R"(redis.call('get', 'oof'); redis.call('set', 'oof', 'bar'); return 1)");
+  Run({"EVAL", eval, "2", "foo", "oof"});
+  Run({"PING"});
+  EXPECT_EQ(InvalidationMessagesLen("IO0"), 3);
+}
+
 }  // namespace dfly
diff --git a/src/server/transaction.cc b/src/server/transaction.cc
index f0be75cb8..ffabc7436 100644
--- a/src/server/transaction.cc
+++ b/src/server/transaction.cc
@@ -665,9 +665,7 @@ void Transaction::RunCallback(EngineShard* shard) {
   // Log to journal only once the command finished running
   if ((coordinator_state_ & COORD_CONCLUDING) || (multi_ && multi_->concluding)) {
     LogAutoJournalOnShard(shard, result);
-    if (tracking_cb_) {
-      tracking_cb_(this);
-    }
+    MaybeInvokeTrackingCb();
   }
 }
 
@@ -1252,6 +1250,7 @@ OpStatus Transaction::RunSquashedMultiCb(RunnableType cb) {
   auto result = cb(this, shard);
   shard->db_slice().OnCbFinish();
   LogAutoJournalOnShard(shard, result);
+  MaybeInvokeTrackingCb();
 
   DCHECK_EQ(result.flags, 0);  // if it's sophisticated, we shouldn't squash it
   return result;
diff --git a/src/server/transaction.h b/src/server/transaction.h
index 66abf71cc..5a0427352 100644
--- a/src/server/transaction.h
+++ b/src/server/transaction.h
@@ -364,6 +364,12 @@ class Transaction {
     tracking_cb_ = std::move(f);
   }
 
+  void MaybeInvokeTrackingCb() {
+    if (tracking_cb_) {
+      tracking_cb_(this);
+    }
+  }
+
   // Remove once BZPOP is stabilized
   std::string DEBUGV18_BlockInfo() {
     return "claimed=" + std::to_string(blocking_barrier_.IsClaimed()) +