fix locking of TDbDriverState (#45069)

2 files changed, 47 insertions, 32 deletions

diff --git a/ydb/public/sdk/cpp/src/client/impl/internal/db_driver_state/state.cpp b/ydb/public/sdk/cpp/src/client/impl/internal/db_driver_state/state.cpp
index 49e92a43bc3..46cdd3f8c9d 100644
--- a/ydb/public/sdk/cpp/src/client/impl/internal/db_driver_state/state.cpp
+++ b/ydb/public/sdk/cpp/src/client/impl/internal/db_driver_state/state.cpp
@@ -189,53 +189,67 @@ TDbDriverStatePtr TDbDriverStateTracker::GetDriverState(
         }
     }
     TDbDriverStatePtr strongState;
-    for (;;) {
+    {
         std::unique_lock lock(Lock_);
-        {
+        Notify_.wait(lock, [&]() {
             auto state = States_.find(key);
-            if (state != States_.end()) {
-                auto strong = state->second.lock();
-                if (strong) {
-                    return strong;
-                } else {
-                    // We could find state record, but couldn't promote weak to shared
-                    // this means weak ptr already expired but dtor hasn't been
-                    // called yet. Likely other thread now is waiting on mutex to
-                    // remove expired record from hashmap. So give him chance
-                    // to do it after that we will be able to create new state
-                    lock.unlock();
-                    std::this_thread::yield();
-                    continue;
-                }
+            if (state == States_.end()) {
+                return true;
             }
+            strongState = state->second.lock();
+            if (strongState) {
+                return true;
+            }
+            return false;
+        });
+        if (strongState) {
+            return strongState;
         }
         {
             auto deleter = [this, key](TDbDriverState* p) {
                 {
                     std::unique_lock lock(Lock_);
                     States_.erase(key);
+                    Notify_.notify_all();
                 }
                 delete p;
             };
-            strongState = std::shared_ptr<TDbDriverState>(
-                new TDbDriverState(
-                    quotedDatabase,
-                    discoveryEndpoint,
-                    discoveryMode,
-                    sslCredentials,
-                    DiscoveryClient_),
-                deleter);
 
-            strongState->SetCredentialsProvider(
-                credentialsProviderFactory
-                    ? credentialsProviderFactory->CreateProvider(strongState)
-                    : CreateInsecureCredentialsProviderFactory()->CreateProvider(strongState));
+            auto [it, inserted] = States_.try_emplace(key); // creates empty weak_ptr
+            auto& weakState = it->second;
+            lock.unlock(); // temporarily release lock
 
-            if (discoveryMode != EDiscoveryMode::Off) {
-                DiscoveryClient_->AddPeriodicTask(CreatePeriodicDiscoveryTask(strongState), DISCOVERY_RECHECK_PERIOD);
+            try {
+                Y_ABORT_UNLESS(inserted);
+                strongState = std::shared_ptr<TDbDriverState>(
+                    new TDbDriverState(
+                        quotedDatabase,
+                        discoveryEndpoint,
+                        discoveryMode,
+                        sslCredentials,
+                        DiscoveryClient_),
+                    deleter);
+
+                strongState->SetCredentialsProvider(
+                    credentialsProviderFactory
+                        ? credentialsProviderFactory->CreateProvider(strongState)
+                        : CreateInsecureCredentialsProviderFactory()->CreateProvider(strongState));
+
+                if (discoveryMode != EDiscoveryMode::Off) {
+                    DiscoveryClient_->AddPeriodicTask(CreatePeriodicDiscoveryTask(strongState), DISCOVERY_RECHECK_PERIOD);
+                }
+            } catch (...) {
+                lock.lock();
+                Y_ABORT_UNLESS(weakState.expired());
+                Y_ABORT_UNLESS(States_.erase(key));
+                Notify_.notify_all();
+                throw;
             }
-            Y_ABORT_UNLESS(States_.emplace(key, strongState).second);
-            break;
+
+            lock.lock(); // re-acquire lock
+            Y_ABORT_UNLESS(weakState.expired());
+            weakState = strongState; // reference remains valid
+            Notify_.notify_all();
         }
     }
 
diff --git a/ydb/public/sdk/cpp/src/client/impl/internal/db_driver_state/state.h b/ydb/public/sdk/cpp/src/client/impl/internal/db_driver_state/state.h
index 8b6e5961e00..7b56c4f2633 100644
--- a/ydb/public/sdk/cpp/src/client/impl/internal/db_driver_state/state.h
+++ b/ydb/public/sdk/cpp/src/client/impl/internal/db_driver_state/state.h
@@ -112,6 +112,7 @@ private:
     IInternalClient* DiscoveryClient_;
     std::unordered_map<TStateKey, std::weak_ptr<TDbDriverState>, TStateKeyHash> States_;
     std::shared_mutex Lock_;
+    std::condition_variable_any Notify_;
 };
 
 using TDbDriverStatePtr = TDbDriverState::TPtr;