Fix lock in Immediate Control Board #5290 (#5289)

3 files changed, 34 insertions, 16 deletions

diff --git a/ydb/core/control/immediate_control_board_impl.cpp b/ydb/core/control/immediate_control_board_impl.cpp
index 06f6dab1a83..40336f4df9e 100644
--- a/ydb/core/control/immediate_control_board_impl.cpp
+++ b/ydb/core/control/immediate_control_board_impl.cpp
@@ -7,16 +7,18 @@
 namespace NKikimr {
 
 bool TControlBoard::RegisterLocalControl(TControlWrapper control, TString name) {
-    bool result = true;
-    if (Board.Has(name)) {
-        result = false;
-    }
-    Board.Insert(name, control.Control);
-    return result;
+    TIntrusivePtr<TControl> ptr;
+    bool result = Board.Swap(name, control.Control, ptr);
+    return !result;
 }
 
 bool TControlBoard::RegisterSharedControl(TControlWrapper& control, TString name) {
-    auto& ptr = Board.InsertIfAbsent(name, control.Control);
+    TIntrusivePtr<TControl> ptr;
+    if (Board.Get(name, ptr)) {
+        control.Control = ptr;
+        return false;
+    }
+    ptr = Board.InsertIfAbsent(name, control.Control);
     if (control.Control == ptr) {
         return true;
     } else {
diff --git a/ydb/core/control/immediate_control_board_ut.cpp b/ydb/core/control/immediate_control_board_ut.cpp
index dba6280ab5b..35ca7524dee 100644
--- a/ydb/core/control/immediate_control_board_ut.cpp
+++ b/ydb/core/control/immediate_control_board_ut.cpp
@@ -111,15 +111,17 @@ Y_UNIT_TEST_SUITE(ControlImplementationTests) {
 
     Y_UNIT_TEST(TestParallelRegisterSharedControl) {
         void* (*parallelJob)(void*) = [](void *controlBoard) -> void *{
-            TControlBoard *Icb = reinterpret_cast<TControlBoard *>(controlBoard);
-            TControlWrapper control1(1, 1, 1);
-            Icb->RegisterSharedControl(control1, "sharedControl");
-            // Useless because running this test with --sanitize=thread cannot reveal
-            // race condition in Icb->RegisterLocalControl(...) without mutex
-            TControlWrapper control2(2, 2, 2);
-            TControlWrapper control2_origin(control2);
-            Icb->RegisterLocalControl(control2, "localControl");
-            UNIT_ASSERT_EQUAL(control2, control2_origin);
+            for (ui64 i = 0; i < 10000; ++i) {
+                TControlBoard *Icb = reinterpret_cast<TControlBoard *>(controlBoard);
+                TControlWrapper control1(1, 1, 1);
+                Icb->RegisterSharedControl(control1, "sharedControl");
+                // Useless because running this test with --sanitize=thread cannot reveal
+                // race condition in Icb->RegisterLocalControl(...) without mutex
+                TControlWrapper control2(2, 2, 2);
+                TControlWrapper control2_origin(control2);
+                Icb->RegisterLocalControl(control2, "localControl");
+                UNIT_ASSERT_EQUAL(control2, control2_origin);
+            }
             return nullptr;
         };
         TIntrusivePtr<TControlBoard> Icb(new TControlBoard);
diff --git a/ydb/core/util/concurrent_rw_hash.h b/ydb/core/util/concurrent_rw_hash.h
index 2e787d1022e..bb25b7d8c32 100644
--- a/ydb/core/util/concurrent_rw_hash.h
+++ b/ydb/core/util/concurrent_rw_hash.h
@@ -49,6 +49,20 @@ public:
         bucket.Map[key] = value;
     }
 
+    bool Swap(const K& key, const V& value, V& out_prev_value) {
+        TBucket& bucket = GetBucketForKey(key);
+        TWriteGuard guard(bucket.RWLock);
+
+        typename TActualMap::iterator it = bucket.Map.find(key);
+        if (it != bucket.Map.end()) {
+            out_prev_value = it->second;
+            it->second = value;
+            return true;
+        }
+        bucket.Map.insert(std::make_pair(key, value));
+        return false;
+    }
+
     V& InsertIfAbsent(const K& key, const V& value) {
         TBucket& bucket = GetBucketForKey(key);
         TWriteGuard guard(bucket.RWLock);