From 078db5caae314f6d667ada4a0839fba16638281e Mon Sep 17 00:00:00 2001
From: Vladislav <vlad@dragonflydb.io>
Date: Mon, 15 Jan 2024 13:51:30 +0300
Subject: [PATCH] fix(tx): guard parallel writes to local result (#2417)

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

diff --git a/src/server/transaction.cc b/src/server/transaction.cc
index 961b277a3..7dd771cf2 100644
--- a/src/server/transaction.cc
+++ b/src/server/transaction.cc
@@ -484,14 +484,16 @@ bool Transaction::RunInShard(EngineShard* shard, bool txq_ooo) {
       local_result_ = status;
     } else {
       if (status == OpStatus::OUT_OF_MEMORY) {
+        absl::base_internal::SpinLockHolder lk{&local_result_mu_};
+        CHECK(local_result_ == OpStatus::OK || local_result_ == OpStatus::OUT_OF_MEMORY);
         local_result_ = status;
       } else {
         CHECK_EQ(OpStatus::OK, status);
       }
     }
   } catch (std::bad_alloc&) {
-    // TODO: to log at most once per sec.
-    LOG_FIRST_N(ERROR, 16) << " out of memory";
+    LOG_FIRST_N(ERROR, 16) << " out of memory";  // TODO: to log at most once per sec.
+    absl::base_internal::SpinLockHolder lk{&local_result_mu_};
     local_result_ = OpStatus::OUT_OF_MEMORY;
   } catch (std::exception& e) {
     LOG(FATAL) << "Unexpected exception " << e.what();
diff --git a/src/server/transaction.h b/src/server/transaction.h
index 5a8bba12b..66d30f2be 100644
--- a/src/server/transaction.h
+++ b/src/server/transaction.h
@@ -4,6 +4,7 @@
 
 #pragma once
 
+#include <absl/base/internal/spinlock.h>
 #include <absl/container/flat_hash_map.h>
 #include <absl/container/flat_hash_set.h>
 #include <absl/container/inlined_vector.h>
@@ -569,7 +570,7 @@ class Transaction {
   DbIndex db_index_{0};
   uint64_t time_now_ms_{0};
 
-  std::atomic<uint32_t> wakeup_requested_{0};  // whether tx was woken up
+  std::atomic_uint32_t wakeup_requested_{0};  // whether tx was woken up
   std::atomic_uint32_t use_count_{0}, run_count_{0}, seqlock_{0};
 
   // unique_shard_cnt_ and unique_shard_id_ are accessed only by coordinator thread.
@@ -586,8 +587,9 @@ class Transaction {
   // If COORDINATOR_XXX has been set, it means we passed or crossed stage XXX.
   uint8_t coordinator_state_ = 0;
 
-  // Used for single-hop transactions with unique_shards_ == 1, hence no data-race.
+  // Result of callbacks. Usually written by single shard only, lock below for multishard oom error
   OpStatus local_result_ = OpStatus::OK;
+  absl::base_internal::SpinLock local_result_mu_;
 
  private:
   struct TLTmpSpace {