From 0e4e971ad94570b11f8362103616e80aa00ce388 Mon Sep 17 00:00:00 2001
From: Vladislav <vlad@dragonflydb.io>
Date: Sat, 24 Aug 2024 13:32:21 +0300
Subject: [PATCH] chore(server): Fix watch (#3557)

---
 src/server/main_service.cc | 12 ++++++++----
 src/server/multi_test.cc   | 11 +++++++++--
 src/server/transaction.cc  |  4 ++--
 3 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/src/server/main_service.cc b/src/server/main_service.cc
index 0665a540b..b3b3c6fe6 100644
--- a/src/server/main_service.cc
+++ b/src/server/main_service.cc
@@ -2070,8 +2070,8 @@ void Service::Discard(CmdArgList args, ConnectionContext* cntx) {
 }
 
 // Return true if non of the connections watched keys expired.
-bool CheckWatchedKeyExpiry(ConnectionContext* cntx, const CommandRegistry& registry) {
-  static char EXISTS[] = "EXISTS";
+bool CheckWatchedKeyExpiry(ConnectionContext* cntx, const CommandId* exists_cid,
+                           const CommandId* exec_cid) {
   auto& exec_info = cntx->conn_state.exec_info;
 
   CmdArgVec str_list(exec_info.watched_keys.size());
@@ -2089,11 +2089,14 @@ bool CheckWatchedKeyExpiry(ConnectionContext* cntx, const CommandRegistry& regis
     return OpStatus::OK;
   };
 
-  cntx->transaction->MultiSwitchCmd(registry.Find(EXISTS));
+  cntx->transaction->MultiSwitchCmd(exists_cid);
   cntx->transaction->InitByArgs(cntx->ns, cntx->conn_state.db_index, CmdArgList{str_list});
   OpStatus status = cntx->transaction->ScheduleSingleHop(std::move(cb));
   CHECK_EQ(OpStatus::OK, status);
 
+  // Reset cid to EXEC as it was before
+  cntx->transaction->MultiSwitchCmd(exec_cid);
+
   // The comparison can still be true even if a key expired due to another one being created.
   // So we have to check the watched_dirty flag, which is set if a key expired.
   return watch_exist_count.load() == exec_info.watched_existed &&
@@ -2206,7 +2209,8 @@ void Service::Exec(CmdArgList args, ConnectionContext* cntx) {
   }
 
   // EXEC should not run if any of the watched keys expired.
-  if (!exec_info.watched_keys.empty() && !CheckWatchedKeyExpiry(cntx, registry_)) {
+  if (!exec_info.watched_keys.empty() &&
+      !CheckWatchedKeyExpiry(cntx, registry_.Find("EXISTS"), exec_cid_)) {
     cntx->transaction->UnlockMulti();
     return rb->SendNull();
   }
diff --git a/src/server/multi_test.cc b/src/server/multi_test.cc
index 5d6d66df8..6719b6fcf 100644
--- a/src/server/multi_test.cc
+++ b/src/server/multi_test.cc
@@ -484,6 +484,14 @@ TEST_F(MultiTest, Watch) {
   Run({"multi"});
   ASSERT_THAT(Run({"exec"}), kExecFail);
 
+  // Check watch with nonempty exec body
+  EXPECT_EQ(Run({"watch", "a"}), "OK");
+  Run({"multi"});
+  Run({"get", "a"});
+  Run({"get", "b"});
+  Run({"get", "c"});
+  ASSERT_THAT(Run({"exec"}), kExecSuccess);
+
   // Check watch data cleared after EXEC.
   Run({"set", "a", "1"});
   Run({"multi"});
@@ -499,10 +507,9 @@ TEST_F(MultiTest, Watch) {
   // Check EXEC doesn't miss watched key expiration.
   Run({"watch", "a"});
   Run({"expire", "a", "1"});
-
   AdvanceTime(1000);
-
   Run({"multi"});
+  Run({"get", "a"});
   ASSERT_THAT(Run({"exec"}), kExecFail);
 
   // Check unwatch.
diff --git a/src/server/transaction.cc b/src/server/transaction.cc
index b4531755a..4f8dd98a8 100644
--- a/src/server/transaction.cc
+++ b/src/server/transaction.cc
@@ -506,8 +506,8 @@ void Transaction::MultiUpdateWithParent(const Transaction* parent) {
 
 void Transaction::MultiBecomeSquasher() {
   DCHECK(multi_->mode == GLOBAL || multi_->mode == LOCK_AHEAD);
-  DCHECK_GT(GetUniqueShardCnt(), 0u);    // initialized and determined active shards
-  DCHECK(cid_->IsMultiTransactional());  // proper base command set
+  DCHECK_GT(GetUniqueShardCnt(), 0u);                    // initialized and determined active shards
+  DCHECK(cid_->IsMultiTransactional()) << cid_->name();  // proper base command set
   multi_->role = SQUASHER;
 }