Added resource waiting

Author: GrigoriyPA

ydb/tests/tools/kqprun/src/actors.cpp

@@ -1,6 +1,7 @@
 #include "actors.h"
 
 #include <ydb/core/kqp/common/simple/services.h>
+#include <ydb/core/kqp/rm_service/kqp_rm_service.h>
 
 
 namespace NKqpRun {
@@ -91,6 +92,54 @@ class TRunScriptActorMock : public NActors::TActorBootstrapped<TRunScriptActorMo
     TProgressCallback ProgressCallback_;
 };
 
+class TResourceWaiterActor : public NActors::TActorBootstrapped<TResourceWaiterActor> {
+public:
+    TResourceWaiterActor(NThreading::TPromise<void> promise, i32 expectedNodeCount)
+        : ExpectedNodeCount_(expectedNodeCount)
+        , Promise_(promise)
+    {}
+
+    void Bootstrap() {
+        GetResourceManager();
+        WaitResourcePublish();
+
+        Promise_.SetValue();
+        PassAway();
+    }
+
+private:
+    void GetResourceManager() {
+        while (true) {
+            ResourceManager_ = NKikimr::NKqp::TryGetKqpResourceManager(SelfId().NodeId());
+            if (ResourceManager_) {
+                break;
+            }
+
+            Sleep(TDuration::MilliSeconds(10));
+        }
+    }
+
+    void WaitResourcePublish() {
+        while (true) {
+            auto resourcesPromise = NThreading::NewPromise<i32>();
+            ResourceManager_->RequestClusterResourcesInfo([resourcesPromise](TVector<NKikimrKqp::TKqpNodeResources>&& resources) mutable {
+                resourcesPromise.SetValue(resources.size());
+            });
+
+            if (resourcesPromise.GetFuture().GetValueSync() == ExpectedNodeCount_) {
+                break;
+            }
+            Sleep(TDuration::MilliSeconds(10));
+        }
+    }
+
+private:
+    const i32 ExpectedNodeCount_;
+
+    NThreading::TPromise<void> Promise_;
+    std::shared_ptr<NKikimr::NKqp::NRm::IKqpResourceManager> ResourceManager_;
+};
+
 }  // anonymous namespace
 
 NActors::IActor* CreateRunScriptActorMock(THolder<NKikimr::NKqp::TEvKqp::TEvQueryRequest> request,
@@ -100,4 +149,8 @@ NActors::IActor* CreateRunScriptActorMock(THolder<NKikimr::NKqp::TEvKqp::TEvQuer
     return new TRunScriptActorMock(std::move(request), promise, resultRowsLimit, resultSizeLimit, resultSets, progressCallback);
 }
 
+NActors::IActor* CreateResourceWaiterActor(NThreading::TPromise<void> promise, i32 expectedNodeCount) {
+    return new TResourceWaiterActor(promise, expectedNodeCount);
+}
+
 }  // namespace NKqpRun

ydb/tests/tools/kqprun/src/actors.h

@@ -11,4 +11,6 @@ NActors::IActor* CreateRunScriptActorMock(THolder<NKikimr::NKqp::TEvKqp::TEvQuer
     ui64 resultRowsLimit, ui64 resultSizeLimit, std::vector<Ydb::ResultSet>& resultSets,
     TProgressCallback progressCallback);
 
+NActors::IActor* CreateResourceWaiterActor(NThreading::TPromise<void> promise, i32 expectedNodeCount);
+
 }  // namespace NKqpRun

ydb/tests/tools/kqprun/src/common.h

@@ -16,6 +16,7 @@ constexpr char YQL_TOKEN_VARIABLE[] = "YQL_TOKEN";
 struct TYdbSetupSettings {
     i32 NodeCount = 1;
     TString DomainName = "Root";
+    TDuration InitializationTimeout = TDuration::Seconds(10);
 
     bool TraceOptEnabled = false;
     TMaybe<TString> LogOutputFile;

ydb/tests/tools/kqprun/src/ydb_setup.cpp

@@ -163,12 +163,27 @@ class TYdbSetup::TImpl {
         NYql::NLog::InitLogger(NActors::CreateNullBackend());
     }
 
+    void WaitResourcePublishing() const {
+        auto promise = NThreading::NewPromise();
+        GetRuntime()->Register(CreateResourceWaiterActor(promise, Settings_.NodeCount));
+
+        try {
+            promise.GetFuture().GetValue(Settings_.InitializationTimeout);
+        } catch (...) {
+            ythrow yexception() << "Failed to initialize all resources: " << CurrentExceptionMessage();
+        }
+    }
+
 public:
     explicit TImpl(const TYdbSetupSettings& settings)
         : Settings_(settings)
     {
         InitializeYqlLogger();
         InitializeServer();
+
+        if (Settings_.NodeCount > 1) {
+            WaitResourcePublishing();
+        }
     }
 
     NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr SchemeQueryRequest(const TString& query, const TString& traceId) const {