From f73c7d0e426a09f28d0b048114d25115fc2aef84 Mon Sep 17 00:00:00 2001
From: Vladislav <vlad@dragonflydb.io>
Date: Wed, 24 Jul 2024 14:05:00 +0300
Subject: [PATCH] fix(transaction): Properly store block cancel status (#3371)

---
 src/server/server_family.cc | 2 +-
 src/server/transaction.cc   | 6 ++++--
 src/server/transaction.h    | 3 +++
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/server/server_family.cc b/src/server/server_family.cc
index d5d826c61..4c9dc9f2d 100644
--- a/src/server/server_family.cc
+++ b/src/server/server_family.cc
@@ -1568,7 +1568,7 @@ void ServerFamily::CancelBlockingOnThread(std::function<OpStatus(ArgSlice)> stat
   auto cb = [status_cb](unsigned thread_index, util::Connection* conn) {
     if (auto fcntx = static_cast<facade::Connection*>(conn)->cntx(); fcntx) {
       auto* cntx = static_cast<ConnectionContext*>(fcntx);
-      if (cntx->transaction && cntx->blocked) {
+      if (cntx->transaction) {
         cntx->transaction->CancelBlocking(status_cb);
       }
     }
diff --git a/src/server/transaction.cc b/src/server/transaction.cc
index 4d6d2f8a7..694a4ec46 100644
--- a/src/server/transaction.cc
+++ b/src/server/transaction.cc
@@ -1178,7 +1178,8 @@ OpStatus Transaction::WaitOnWatch(const time_point& tp, WaitKeysProvider wkeys_p
   if (status == cv_status::timeout) {
     result = OpStatus::TIMED_OUT;
   } else if (coordinator_state_ & COORD_CANCELLED) {
-    result = local_result_;
+    DCHECK_GT(block_cancel_result_, OpStatus::OK);
+    result = block_cancel_result_;
   }
 
   // If we don't follow up with an "action" hop, we must clean up manually on all shards.
@@ -1412,7 +1413,8 @@ void Transaction::CancelBlocking(std::function<OpStatus(ArgSlice)> status_cb) {
     return;
 
   coordinator_state_ |= COORD_CANCELLED;
-  local_result_ = status;
+  // don't use local_result_ because it can be overwirtten if we cancel ahead
+  block_cancel_result_ = status;
   blocking_barrier_.Close();
 }
 
diff --git a/src/server/transaction.h b/src/server/transaction.h
index da0836dc5..4fa526600 100644
--- a/src/server/transaction.h
+++ b/src/server/transaction.h
@@ -630,6 +630,9 @@ class Transaction {
 
   // Barrier for waking blocking transactions that ensures exclusivity of waking operation.
   BatonBarrier blocking_barrier_{};
+  // Stores status if COORD_CANCELLED was set. Apart from cancelled, it can be moved for cluster
+  // changes
+  OpStatus block_cancel_result_ = OpStatus::OK;
 
   // Transaction coordinator state, written and read by coordinator thread.
   uint8_t coordinator_state_ = 0;