stable-25-3-1: EXT-1430 Fix race in resource manager (#25738)

Author: ydbot

ydb/core/kqp/rm_service/kqp_resource_info_exchanger.cpp

@@ -316,19 +316,18 @@ class TKqpResourceInfoExchangerActor : public TActorBootstrapped<TKqpResourceInf
             Schedule(at - now, new TEvPrivate::TEvUpdateSnapshotState);
             return;
         }
-        TVector<NKikimrKqp::TKqpNodeResources> resources;
-        resources.reserve(NodesState.size());
+        std::shared_ptr<TVector<NKikimrKqp::TKqpNodeResources>> resources = std::make_shared<TVector<NKikimrKqp::TKqpNodeResources>>();
+        resources->reserve(NodesState.size());
 
         for (const auto& [nodeId, state] : NodesState) {
             const auto& currentResources = state.NodeData.GetResources();
             if (currentResources.HasNodeId()) {
-                resources.push_back(std::move(currentResources));
+                resources->push_back(std::move(currentResources));
             }
         }
 
         with_lock (ResourceSnapshotState->Lock) {
-            ResourceSnapshotState->Snapshot =
-                std::make_shared<TVector<NKikimrKqp::TKqpNodeResources>>(std::move(resources));
+            ResourceSnapshotState->Snapshot = std::move(resources);
         }
 
         ResourceSnapshotRetryState.LastRetryAt = now;

ydb/core/kqp/rm_service/kqp_rm_service.cpp

@@ -171,12 +171,12 @@ class TKqpResourceManager : public IKqpResourceManager {
         SetConfigValues(config);
     }
 
-    void Bootstrap(NKikimrConfig::TTableServiceConfig::TResourceManager& config, TActorSystem* actorSystem, TActorId selfId) {
-        if (!Counters) {
-            Counters = MakeIntrusive<TKqpCounters>(AppData()->Counters);
-        }
+    void Registered(NKikimrConfig::TTableServiceConfig::TResourceManager& config, TActorSystem* actorSystem, TActorId selfId) {
         ActorSystem = actorSystem;
         SelfId = selfId;
+        if (!Counters) {
+            Counters = MakeIntrusive<TKqpCounters>(AppData(ActorSystem)->Counters);
+        }
         UpdatePatternCache(config.GetKqpPatternCacheCapacityBytes(),
             config.GetKqpPatternCacheCompiledCapacityBytes(),
             config.GetKqpPatternCachePatternAccessTimesBeforeTryToCompile());
@@ -582,20 +582,30 @@ class TKqpResourceManagerActor : public TActorBootstrapped<TKqpResourceManagerAc
     TKqpResourceManagerActor(const NKikimrConfig::TTableServiceConfig::TResourceManager& config,
         TIntrusivePtr<TKqpCounters> counters, const TActorId& resourceBrokerId,
         std::shared_ptr<TKqpProxySharedResources>&& kqpProxySharedResources, ui32 nodeId)
-        : Config(config)
+        : NodeId(nodeId)
+        , Config(config)
         , ResourceBrokerId(resourceBrokerId ? resourceBrokerId : MakeResourceBrokerID())
         , KqpProxySharedResources(std::move(kqpProxySharedResources))
     {
         ResourceManager = std::make_shared<TKqpResourceManager>(config, counters);
+    }
+
+    // Is called right after service registration
+    // and before any usual actor can try to get ResourceManager
+    void Registered(TActorSystem* sys, const TActorId& owner) override {
+        TActorBootstrapped::Registered(sys, owner);
+
+        ResourceManager->Registered(Config, sys, SelfId());
+
         with_lock (ResourceManagers.Lock) {
-            ResourceManagers.ByNodeId[nodeId] = ResourceManager;
-            ResourceManagers.Default = ResourceManager;
+            if (ResourceManagers.Default.expired()) { // There can be several managers in tests
+                ResourceManagers.Default = ResourceManager;
+            }
+            ResourceManagers.ByNodeId[NodeId] = ResourceManager;
         }
     }
 
     void Bootstrap() {
-        ResourceManager->Bootstrap(Config, TlsActivationContext->ActorSystem(), SelfId());
-
         LOG_D("Start KqpResourceManagerActor at " << SelfId() << " with ResourceBroker at " << ResourceBrokerId);
 
         // Subscribe for tenant changes
@@ -935,6 +945,7 @@ class TKqpResourceManagerActor : public TActorBootstrapped<TKqpResourceManagerAc
     }
 
 private:
+    const ui32 NodeId;
     NKikimrConfig::TTableServiceConfig::TResourceManager Config;
 
     const TActorId ResourceBrokerId;
@@ -984,7 +995,7 @@ std::shared_ptr<NRm::IKqpResourceManager> GetKqpResourceManager(TMaybe<ui32> _no
 
 std::shared_ptr<NRm::IKqpResourceManager> TryGetKqpResourceManager(TMaybe<ui32> _nodeId) {
     ui32 nodeId = _nodeId ? *_nodeId : TActivationContext::ActorSystem()->NodeId;
-    auto rm = NRm::ResourceManagers.Default.lock();
+    std::shared_ptr<NRm::TKqpResourceManager> rm = NRm::ResourceManagers.Default.lock();
     if (Y_LIKELY(rm && rm->GetNodeId() == nodeId)) {
         return rm;
     }

ydb/core/kqp/rm_service/kqp_rm_ut.cpp

@@ -150,10 +150,14 @@ class KqpRm : public TTestBase {
             const NKikimrConfig::TTableServiceConfig::TResourceManager& config, ui32 nodeInd = 0) {
         auto kqpCounters = MakeIntrusive<TKqpCounters>(Counters);
         auto resman = CreateKqpResourceManagerActor(config, kqpCounters, ResourceBrokers[nodeInd], nullptr, Runtime->GetNodeId(nodeInd));
+        // RM creates children during its registration, we need to enable schedule for them
+        auto prevObserver = Runtime->SetRegistrationObserverFunc([](TTestActorRuntimeBase& runtime, const TActorId& /*parentId*/, const TActorId& actorId) {
+            runtime.EnableScheduleForActor(actorId, true);
+        });
         ResourceManagers.push_back(Runtime->Register(resman, nodeInd));
         Runtime->RegisterService(MakeKqpResourceManagerServiceID(
             Runtime->GetNodeId(nodeInd)), ResourceManagers.back(), nodeInd);
-        Runtime->EnableScheduleForActor(ResourceManagers.back(), true);
+        Runtime->SetRegistrationObserverFunc(prevObserver);
     }
 
     void StartRms(const TVector<NKikimrConfig::TTableServiceConfig::TResourceManager>& configs = {}) {