KqpRun added inprogress statistic (#4861)

Author: GrigoriyPA

ydb/tests/tools/kqprun/configuration/app_config.conf

@@ -73,6 +73,10 @@ QueryServiceConfig {
     MinDesiredDirectoriesOfFilesPerQuery: 1000
     RegexpCacheSize: 100
 
+    DefaultSettings {
+      Name: "AtomicUploadCommit"
+      Value: "true"
+    }
     DefaultSettings {
       Name: "UseBlocksSource"
       Value: "true"

ydb/tests/tools/kqprun/kqprun.cpp

@@ -32,7 +32,7 @@ struct TExecutionOptions {
     EClearExecutionCase ClearExecution = EClearExecutionCase::Disabled;
     NKikimrKqp::EQueryAction ScriptQueryAction = NKikimrKqp::QUERY_ACTION_EXECUTE;
 
-    TString TraceId = "kqprun";
+    TString TraceId = "kqprun_" + CreateGuidAsString();
 
     bool HasResults() const {
         return !ScriptQueries.empty() && ScriptQueryAction == NKikimrKqp::QUERY_ACTION_EXECUTE;
@@ -163,6 +163,7 @@ void RunMain(int argc, const char* argv[]) {
     TString schemeQueryAstFile;
     TString scriptQueryAstFile;
     TString scriptQueryPlanFile;
+    TString inProgressStatisticsFile;
     TString logFile = "-";
     TString appConfigFile = "./configuration/app_config.conf";
     std::vector<TString> tablesMappingList;
@@ -218,6 +219,10 @@ void RunMain(int argc, const char* argv[]) {
         .Optional()
         .RequiredArgument("FILE")
         .StoreResult(&scriptQueryPlanFile);
+    options.AddLongOption("in-progress-statistics", "File with script inprogress statistics")
+        .Optional()
+        .RequiredArgument("FILE")
+        .StoreResult(&inProgressStatisticsFile);
 
     options.AddLongOption('C', "clear-execution", "Execute script query without creating additional tables, one of { query | yql-script }")
         .Optional()
@@ -319,6 +324,10 @@ void RunMain(int argc, const char* argv[]) {
     THolder<TFileOutput> scriptQueryAstFileHolder = SetupDefaultFileOutput(scriptQueryAstFile, runnerOptions.ScriptQueryAstOutput);
     THolder<TFileOutput> scriptQueryPlanFileHolder = SetupDefaultFileOutput(scriptQueryPlanFile, runnerOptions.ScriptQueryPlanOutput);
 
+    if (inProgressStatisticsFile) {
+        runnerOptions.InProgressStatisticsOutputFile = inProgressStatisticsFile;
+    }
+
     runnerOptions.TraceOptType = GetCaseVariant<NKqpRun::TRunnerOptions::ETraceOptType>("trace-opt", traceOptType, {
         {"all", NKqpRun::TRunnerOptions::ETraceOptType::All},
         {"scheme", NKqpRun::TRunnerOptions::ETraceOptType::Scheme},
@@ -399,8 +408,20 @@ void KqprunTerminateHandler() {
 }
 
 
+void SegmentationFaultHandler(int) {
+    NColorizer::TColors colors = NColorizer::AutoColors(Cerr);
+
+    Cerr << colors.Red() << "======= segmentation fault call stack ========" << colors.Default() << Endl;
+    FormatBackTrace(&Cerr);
+    Cerr << colors.Red() << "==============================================" << colors.Default() << Endl;
+
+    abort();
+}
+
+
 int main(int argc, const char* argv[]) {
     std::set_terminate(KqprunTerminateHandler);
+    signal(SIGSEGV, &SegmentationFaultHandler);
 
     try {
         RunMain(argc, argv);

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 {
@@ -11,13 +12,14 @@ class TRunScriptActorMock : public NActors::TActorBootstrapped<TRunScriptActorMo
 public:
     TRunScriptActorMock(THolder<NKikimr::NKqp::TEvKqp::TEvQueryRequest> request,
         NThreading::TPromise<NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr> promise,
-        ui64 resultRowsLimit, ui64 resultSizeLimit, std::vector<Ydb::ResultSet>& resultSets, TString& queryPlan)
+        ui64 resultRowsLimit, ui64 resultSizeLimit, std::vector<Ydb::ResultSet>& resultSets,
+        TProgressCallback progressCallback)
         : Request_(std::move(request))
         , Promise_(promise)
         , ResultRowsLimit_(std::numeric_limits<ui64>::max())
         , ResultSizeLimit_(std::numeric_limits<i64>::max())
         , ResultSets_(resultSets)
-        , QueryPlan_(queryPlan)
+        , ProgressCallback_(progressCallback)
     {
         if (resultRowsLimit) {
             ResultRowsLimit_ = resultRowsLimit;
@@ -76,7 +78,9 @@ class TRunScriptActorMock : public NActors::TActorBootstrapped<TRunScriptActorMo
     }
 
     void Handle(NKikimr::NKqp::TEvKqpExecuter::TEvExecuterProgress::TPtr& ev) {
-        QueryPlan_ = ev->Get()->Record.GetQueryPlan();
+        if (ProgressCallback_) {
+            ProgressCallback_(ev->Get()->Record);
+        }
     }
 
 private:
@@ -85,15 +89,104 @@ class TRunScriptActorMock : public NActors::TActorBootstrapped<TRunScriptActorMo
     ui64 ResultRowsLimit_;
     ui64 ResultSizeLimit_;
     std::vector<Ydb::ResultSet>& ResultSets_;
-    TString& QueryPlan_;
+    TProgressCallback ProgressCallback_;
+};
+
+class TResourcesWaiterActor : public NActors::TActorBootstrapped<TResourcesWaiterActor> {
+    struct TEvPrivate {
+        enum EEv : ui32 {
+            EvResourcesInfo = EventSpaceBegin(NActors::TEvents::ES_PRIVATE),
+
+            EvEnd
+        };
+
+        static_assert(EvEnd < EventSpaceEnd(NActors::TEvents::ES_PRIVATE), "expect EvEnd < EventSpaceEnd(NActors::TEvents::ES_PRIVATE)");
+
+        struct TEvResourcesInfo : public NActors::TEventLocal<TEvResourcesInfo, EvResourcesInfo> {
+            explicit TEvResourcesInfo(i32 nodeCount)
+                : NodeCount(nodeCount)
+            {}
+
+            const i32 NodeCount;
+        };
+    };
+
+    static constexpr TDuration REFRESH_PERIOD = TDuration::MilliSeconds(10);
+
+public:
+    TResourcesWaiterActor(NThreading::TPromise<void> promise, i32 expectedNodeCount)
+        : ExpectedNodeCount_(expectedNodeCount)
+        , Promise_(promise)
+    {}
+
+    void Bootstrap() {
+        Become(&TResourcesWaiterActor::StateFunc);
+        CheckResourcesPublish();
+    }
+
+    void Handle(NActors::TEvents::TEvWakeup::TPtr&) {
+        CheckResourcesPublish();
+    }
+
+    void Handle(TEvPrivate::TEvResourcesInfo::TPtr& ev) {
+        if (ev->Get()->NodeCount == ExpectedNodeCount_) {
+            Promise_.SetValue();
+            PassAway();
+            return;
+        }
+
+        Schedule(REFRESH_PERIOD, new NActors::TEvents::TEvWakeup());
+    }
+
+    STRICT_STFUNC(StateFunc,
+        hFunc(NActors::TEvents::TEvWakeup, Handle);
+        hFunc(TEvPrivate::TEvResourcesInfo, Handle);
+    )
+
+private:
+    void CheckResourcesPublish() {
+        GetResourceManager();
+
+        if (!ResourceManager_) {
+            Schedule(REFRESH_PERIOD, new NActors::TEvents::TEvWakeup());
+            return;
+        }
+
+        UpdateResourcesInfo();
+    }
+
+    void GetResourceManager() {
+        if (ResourceManager_) {
+            return;
+        }
+        ResourceManager_ = NKikimr::NKqp::TryGetKqpResourceManager(SelfId().NodeId());
+    }
+
+    void UpdateResourcesInfo() const {
+        ResourceManager_->RequestClusterResourcesInfo(
+        [selfId = SelfId(), actorContext = ActorContext()](TVector<NKikimrKqp::TKqpNodeResources>&& resources) {
+            actorContext.Send(selfId, new TEvPrivate::TEvResourcesInfo(resources.size()));
+        });
+    }
+
+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,
     NThreading::TPromise<NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr> promise,
-    ui64 resultRowsLimit, ui64 resultSizeLimit, std::vector<Ydb::ResultSet>& resultSets, TString& queryPlan) {
-    return new TRunScriptActorMock(std::move(request), promise, resultRowsLimit, resultSizeLimit, resultSets, queryPlan);
+    ui64 resultRowsLimit, ui64 resultSizeLimit, std::vector<Ydb::ResultSet>& resultSets,
+    TProgressCallback progressCallback) {
+    return new TRunScriptActorMock(std::move(request), promise, resultRowsLimit, resultSizeLimit, resultSets, progressCallback);
+}
+
+NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise<void> promise, i32 expectedNodeCount) {
+    return new TResourcesWaiterActor(promise, expectedNodeCount);
 }
 
 }  // namespace NKqpRun

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

@@ -4,8 +4,13 @@
 
 namespace NKqpRun {
 
+using TProgressCallback = std::function<void(const NKikimrKqp::TEvExecuterProgress&)>;
+
 NActors::IActor* CreateRunScriptActorMock(THolder<NKikimr::NKqp::TEvKqp::TEvQueryRequest> request,
     NThreading::TPromise<NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr> promise,
-    ui64 resultRowsLimit, ui64 resultSizeLimit, std::vector<Ydb::ResultSet>& resultSets, TString& queryPlan);
+    ui64 resultRowsLimit, ui64 resultSizeLimit, std::vector<Ydb::ResultSet>& resultSets,
+    TProgressCallback progressCallback);
+
+NActors::IActor* CreateResourcesWaiterActor(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;
@@ -46,6 +47,7 @@ struct TRunnerOptions {
     IOutputStream* SchemeQueryAstOutput = nullptr;
     IOutputStream* ScriptQueryAstOutput = nullptr;
     IOutputStream* ScriptQueryPlanOutput = nullptr;
+    TMaybe<TString> InProgressStatisticsOutputFile;
 
     EResultOutputFormat ResultOutputFormat = EResultOutputFormat::RowsJson;
     NYdb::NConsoleClient::EOutputFormat PlanOutputFormat = NYdb::NConsoleClient::EOutputFormat::Default;