Added AsyncQueryRunnerActor

ydb/tests/tools/kqprun/kqprun.cpp

@@ -81,7 +81,8 @@ void RunScript(const TExecutionOptions& executionOptions, const NKqpRun::TRunner
             Sleep(executionOptions.LoopDelay);
         }
 
-        if (executionOptions.GetExecutionCase(id) != TExecutionOptions::EExecutionCase::AsyncQuery) {
+        const auto executionCase = executionOptions.GetExecutionCase(id);
+        if (executionCase != TExecutionOptions::EExecutionCase::AsyncQuery) {
             Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Executing script";
             if (numberQueries > 1) {
                 Cout << " " << id;
@@ -92,7 +93,7 @@ void RunScript(const TExecutionOptions& executionOptions, const NKqpRun::TRunner
             Cout << "..." << colors.Default() << Endl;
         }
 
-        switch (executionOptions.GetExecutionCase(id)) {
+        switch (executionCase) {
         case TExecutionOptions::EExecutionCase::GenericScript:
             if (!runner.ExecuteScript(executionOptions.ScriptQueries[id], executionOptions.ScriptQueryAction, executionOptions.TraceId)) {
                 ythrow yexception() << TInstant::Now().ToIsoStringLocal() << " Script execution failed";
@@ -381,8 +382,8 @@ class TMain : public TMainClassArgs {
             });
         options.AddLongOption("inflight-limit", "In flight limit for async queries (use 0 for unlimited)")
             .RequiredArgument("uint")
-            .DefaultValue(RunnerOptions.InFlightLimit)
-            .StoreResult(&RunnerOptions.InFlightLimit);
+            .DefaultValue(RunnerOptions.YdbSettings.InFlightLimit)
+            .StoreResult(&RunnerOptions.YdbSettings.InFlightLimit);
 
         TChoices<NKikimrKqp::EQueryAction> scriptAction({
             {"execute", NKikimrKqp::QUERY_ACTION_EXECUTE},

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

@@ -1,5 +1,7 @@
 #include "actors.h"
 
+#include <library/cpp/colorizer/colors.h>
+
 #include <ydb/core/kqp/common/simple/services.h>
 #include <ydb/core/kqp/rm_service/kqp_rm_service.h>
 
@@ -10,26 +12,25 @@ namespace {
 
 class TRunScriptActorMock : public NActors::TActorBootstrapped<TRunScriptActorMock> {
 public:
-    TRunScriptActorMock(THolder<NKikimr::NKqp::TEvKqp::TEvQueryRequest> request,
-        NThreading::TPromise<TQueryResponse> promise, ui64 resultRowsLimit, ui64 resultSizeLimit,
-        TProgressCallback progressCallback)
-        : Request_(std::move(request))
+    TRunScriptActorMock(TQueryRequest request, NThreading::TPromise<TQueryResponse> promise, TProgressCallback progressCallback)
+        : TargetNode(request.TargetNode)
+        , Request_(std::move(request.Event))
         , Promise_(promise)
         , ResultRowsLimit_(std::numeric_limits<ui64>::max())
         , ResultSizeLimit_(std::numeric_limits<i64>::max())
         , ProgressCallback_(progressCallback)
     {
-        if (resultRowsLimit) {
-            ResultRowsLimit_ = resultRowsLimit;
+        if (request.ResultRowsLimit) {
+            ResultRowsLimit_ = request.ResultRowsLimit;
         }
-        if (resultSizeLimit) {
-            ResultSizeLimit_ = resultSizeLimit;
+        if (request.ResultSizeLimit) {
+            ResultSizeLimit_ = request.ResultSizeLimit;
         }
     }
 
     void Bootstrap() {
         NActors::ActorIdToProto(SelfId(), Request_->Record.MutableRequestActorId());
-        Send(NKikimr::NKqp::MakeKqpProxyID(SelfId().NodeId()), std::move(Request_));
+        Send(NKikimr::NKqp::MakeKqpProxyID(TargetNode), std::move(Request_));
 
         Become(&TRunScriptActorMock::StateFunc);
     }
@@ -88,7 +89,8 @@ class TRunScriptActorMock : public NActors::TActorBootstrapped<TRunScriptActorMo
     }
 
 private:
-    THolder<NKikimr::NKqp::TEvKqp::TEvQueryRequest> Request_;
+    ui32 TargetNode = 0;
+    std::unique_ptr<NKikimr::NKqp::TEvKqp::TEvQueryRequest> Request_;
     NThreading::TPromise<TQueryResponse> Promise_;
     ui64 ResultRowsLimit_;
     ui64 ResultSizeLimit_;
@@ -97,25 +99,104 @@ class TRunScriptActorMock : public NActors::TActorBootstrapped<TRunScriptActorMo
     std::vector<ui64> ResultSetSizes_;
 };
 
-class TResourcesWaiterActor : public NActors::TActorBootstrapped<TResourcesWaiterActor> {
-    struct TEvPrivate {
-        enum EEv : ui32 {
-            EvResourcesInfo = EventSpaceBegin(NActors::TEvents::ES_PRIVATE),
+class TAsyncQueryRunnerActor : public NActors::TActor<TAsyncQueryRunnerActor> {
+    using TBase = NActors::TActor<TAsyncQueryRunnerActor>;
+
+public:
+    TAsyncQueryRunnerActor(ui64 inFlightLimit)
+        : TBase(&TAsyncQueryRunnerActor::StateFunc)
+        , InFlightLimit_(inFlightLimit)
+    {}
+
+    STRICT_STFUNC(StateFunc,
+        hFunc(TEvPrivate::TEvStartAsyncQuery, Handle);
+        hFunc(TEvPrivate::TEvAsyncQueryFinished, Handle);
+        hFunc(TEvPrivate::TEvFinalizeAsyncQueryRunner, Handle);
+    )
+
+    void Handle(TEvPrivate::TEvStartAsyncQuery::TPtr& ev) {
+        DelayedRequests_.emplace(std::move(ev));
+        StartDelayedRequests();
+    }
+
+    void Handle(TEvPrivate::TEvAsyncQueryFinished::TPtr& ev) {
+        const ui64 requestId = ev->Get()->RequestId;
+        RunningRequests_.erase(requestId);
+
+        const auto& response = ev->Get()->Result.Response->Get()->Record.GetRef();
+        const auto status = response.GetYdbStatus();
+
+        if (status == Ydb::StatusIds::SUCCESS) {
+            Completed_++;
+            Cout << CoutColors_.Green() << TInstant::Now().ToIsoStringLocal() << " Request #" << requestId << " completed. " << CoutColors_.Yellow() << GetInfoString() << CoutColors_.Default() << Endl;
+        } else {
+            Failed_++;
+            NYql::TIssues issues;
+            NYql::IssuesFromMessage(response.GetResponse().GetQueryIssues(), issues);
+            Cout << CoutColors_.Red() << TInstant::Now().ToIsoStringLocal() << " Request #" << requestId << " failed " << status << ". " << CoutColors_.Yellow() << GetInfoString() << "\n" << CoutColors_.Red() << "Issues:\n" << issues.ToString() << CoutColors_.Default();
+        }
+
+        StartDelayedRequests();
+        TryFinalize();
+    }
+
+    void Handle(TEvPrivate::TEvFinalizeAsyncQueryRunner::TPtr& ev) {
+        FinalizePromise_ = ev->Get()->FinalizePromise;
+        if (!TryFinalize()) {
+            Cout << CoutColors_.Yellow() << TInstant::Now().ToIsoStringLocal() << " Waiting for " << DelayedRequests_.size() + RunningRequests_.size() << " async queries..." << CoutColors_.Default() << Endl;
+        }
+    }
 
-            EvEnd
-        };
+private:
+    void StartDelayedRequests() {
+        while (!DelayedRequests_.empty() && (!InFlightLimit_ || RunningRequests_.size() < InFlightLimit_)) {
+            auto request = std::move(DelayedRequests_.front());
+            DelayedRequests_.pop();
+
+            auto promise = NThreading::NewPromise<TQueryResponse>();
+            Register(CreateRunScriptActorMock(std::move(request->Get()->Request), promise, nullptr));
+            RunningRequests_[RequestId_] = promise.GetFuture().Subscribe([id = RequestId_, this](const NThreading::TFuture<TQueryResponse>& f) {
+                Send(SelfId(), new TEvPrivate::TEvAsyncQueryFinished(id, std::move(f.GetValue())));
+            });
+
+            MaxInFlight_ = std::max(MaxInFlight_, RunningRequests_.size());
+            Cout << TStringBuilder() << CoutColors_.Cyan() << TInstant::Now().ToIsoStringLocal() << " Request #" << RequestId_ << " started. " << CoutColors_.Yellow() << GetInfoString() << CoutColors_.Default() << "\n";
+
+            RequestId_++;
+            request->Get()->StartPromise.SetValue();
+        }
+    }
 
-        static_assert(EvEnd < EventSpaceEnd(NActors::TEvents::ES_PRIVATE), "expect EvEnd < EventSpaceEnd(NActors::TEvents::ES_PRIVATE)");
+    bool TryFinalize() {
+        if (!FinalizePromise_ || !RunningRequests_.empty()) {
+            return false;
+        }
 
-        struct TEvResourcesInfo : public NActors::TEventLocal<TEvResourcesInfo, EvResourcesInfo> {
-            explicit TEvResourcesInfo(i32 nodeCount)
-                : NodeCount(nodeCount)
-            {}
+        FinalizePromise_->SetValue();
+        PassAway();
+        return true;
+    }
 
-            const i32 NodeCount;
-        };
-    };
+    TString GetInfoString() const {
+        return TStringBuilder() << "completed: " << Completed_ << ", failed: " << Failed_ << ", in flight: " << RunningRequests_.size() << ", max in flight: " << MaxInFlight_ << ", spend time: " << TInstant::Now() - StartTime_;
+    }
+
+private:
+    const ui64 InFlightLimit_;
+    const TInstant StartTime_ = TInstant::Now();
+    const NColorizer::TColors CoutColors_ = NColorizer::AutoColors(Cout);
+
+    std::optional<NThreading::TPromise<void>> FinalizePromise_;
+    std::queue<TEvPrivate::TEvStartAsyncQuery::TPtr> DelayedRequests_;
+    std::unordered_map<ui64, NThreading::TFuture<TQueryResponse>> RunningRequests_;
+
+    ui64 RequestId_ = 1;
+    ui64 MaxInFlight_ = 0;
+    ui64 Completed_ = 0;
+    ui64 Failed_ = 0;
+};
 
+class TResourcesWaiterActor : public NActors::TActorBootstrapped<TResourcesWaiterActor> {
     static constexpr TDuration REFRESH_PERIOD = TDuration::MilliSeconds(10);
 
 public:
@@ -183,10 +264,12 @@ class TResourcesWaiterActor : public NActors::TActorBootstrapped<TResourcesWaite
 
 }  // anonymous namespace
 
-NActors::IActor* CreateRunScriptActorMock(THolder<NKikimr::NKqp::TEvKqp::TEvQueryRequest> request,
-    NThreading::TPromise<TQueryResponse> promise, ui64 resultRowsLimit, ui64 resultSizeLimit,
-    TProgressCallback progressCallback) {
-    return new TRunScriptActorMock(std::move(request), promise, resultRowsLimit, resultSizeLimit, progressCallback);
+NActors::IActor* CreateRunScriptActorMock(TQueryRequest request, NThreading::TPromise<TQueryResponse> promise, TProgressCallback progressCallback) {
+    return new TRunScriptActorMock(std::move(request), promise, progressCallback);
+}
+
+NActors::IActor* CreateAsyncQueryRunnerActor(ui64 inFlightLimit) {
+    return new TAsyncQueryRunnerActor(inFlightLimit);
 }
 
 NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise<void> promise, i32 expectedNodeCount) {

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

@@ -9,11 +9,68 @@ struct TQueryResponse {
     std::vector<Ydb::ResultSet> ResultSets;
 };
 
+struct TQueryRequest {
+    std::unique_ptr<NKikimr::NKqp::TEvKqp::TEvQueryRequest> Event;
+    ui32 TargetNode;
+    ui64 ResultRowsLimit;
+    ui64 ResultSizeLimit;
+};
+
+struct TEvPrivate {
+    enum EEv : ui32 {
+        EvStartAsyncQuery = EventSpaceBegin(NActors::TEvents::ES_PRIVATE),
+        EvAsyncQueryFinished,
+        EvFinalizeAsyncQueryRunner,
+
+        EvResourcesInfo,
+
+        EvEnd
+    };
+
+    static_assert(EvEnd < EventSpaceEnd(NActors::TEvents::ES_PRIVATE), "expect EvEnd < EventSpaceEnd(NActors::TEvents::ES_PRIVATE)");
+
+    struct TEvStartAsyncQuery : public NActors::TEventLocal<TEvStartAsyncQuery, EvStartAsyncQuery> {
+        TEvStartAsyncQuery(TQueryRequest request, NThreading::TPromise<void> startPromise)
+            : Request(std::move(request))
+            , StartPromise(startPromise)
+        {}
+
+        TQueryRequest Request;
+        NThreading::TPromise<void> StartPromise;
+    };
+
+    struct TEvAsyncQueryFinished : public NActors::TEventLocal<TEvAsyncQueryFinished, EvAsyncQueryFinished> {
+        TEvAsyncQueryFinished(ui64 requestId, TQueryResponse result)
+            : RequestId(requestId)
+            , Result(result)
+        {}
+
+        const ui64 RequestId;
+        const TQueryResponse Result;
+    };
+
+    struct TEvFinalizeAsyncQueryRunner : public NActors::TEventLocal<TEvFinalizeAsyncQueryRunner, EvFinalizeAsyncQueryRunner> {
+        explicit TEvFinalizeAsyncQueryRunner(NThreading::TPromise<void> finalizePromise)
+            : FinalizePromise(finalizePromise)
+        {}
+
+        NThreading::TPromise<void> FinalizePromise;
+    };
+
+    struct TEvResourcesInfo : public NActors::TEventLocal<TEvResourcesInfo, EvResourcesInfo> {
+        explicit TEvResourcesInfo(i32 nodeCount)
+            : NodeCount(nodeCount)
+        {}
+
+        const i32 NodeCount;
+    };
+};
+
 using TProgressCallback = std::function<void(const NKikimrKqp::TEvExecuterProgress&)>;
 
-NActors::IActor* CreateRunScriptActorMock(THolder<NKikimr::NKqp::TEvKqp::TEvQueryRequest> request,
-    NThreading::TPromise<TQueryResponse> promise, ui64 resultRowsLimit, ui64 resultSizeLimit,
-    TProgressCallback progressCallback);
+NActors::IActor* CreateRunScriptActorMock(TQueryRequest request, NThreading::TPromise<TQueryResponse> promise, TProgressCallback progressCallback);
+
+NActors::IActor* CreateAsyncQueryRunnerActor(ui64 inFlightLimit);
 
 NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise<void> promise, i32 expectedNodeCount);
 

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

@@ -28,6 +28,8 @@ struct TYdbSetupSettings {
     NKikimr::NMiniKQL::TComputationNodeFactory ComputationFactory;
     TIntrusivePtr<NYql::IYtGateway> YtGateway;
     NKikimrConfig::TAppConfig AppConfig;
+
+    ui64 InFlightLimit = 0;
 };
 
 
@@ -55,8 +57,6 @@ struct TRunnerOptions {
     NYdb::NConsoleClient::EOutputFormat PlanOutputFormat = NYdb::NConsoleClient::EOutputFormat::Default;
     ETraceOptType TraceOptType = ETraceOptType::Disabled;
 
-    ui64 InFlightLimit = 0;
-
     TYdbSetupSettings YdbSettings;
 };