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

Author: the-ancient-1

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 {

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);

ydb/core/util/concurrent_rw_hash.h

@@ -49,6 +49,20 @@ class TConcurrentRWHashMap {
         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);