Added loop option

ydb-platform · Jun 11, 2024 · e82ac97 · e82ac97
1 parent 4eb393a
commit e82ac97
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 6 deletions.
diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp
@@ -30,15 +30,23 @@ struct TExecutionOptions {
     std::vector<TString> ScriptQueries;
     TString SchemeQuery;
 
+    ui32 LoopCount = 1;
+    TDuration LoopDelay;
+
     bool ForgetExecution = false;
-    EExecutionCase ExecutionCase = EExecutionCase::GenericScript;
+    std::vector<EExecutionCase> ExecutionCases;
     NKikimrKqp::EQueryAction ScriptQueryAction = NKikimrKqp::QUERY_ACTION_EXECUTE;
 
     const TString TraceId = "kqprun_" + CreateGuidAsString();
 
     bool HasResults() const {
         return !ScriptQueries.empty() && ScriptQueryAction == NKikimrKqp::QUERY_ACTION_EXECUTE;
     }
+
+    EExecutionCase GetExecutionCase(size_t index) const {
+        Y_ABORT_UNLESS(!ExecutionCases.empty());
+        return ExecutionCases[std::min(index, ExecutionCases.size() - 1)];
+    }
 };
 
 
@@ -55,9 +63,24 @@ void RunScript(const TExecutionOptions& executionOptions, const NKqpRun::TRunner
         }
     }
 
-    for (size_t id = 0; id < executionOptions.ScriptQueries.size(); ++id) {
-        Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Executing script" << (executionOptions.ScriptQueries.size() > 1 ? TStringBuilder() << " " << id : TString()) << "..." << colors.Default() << Endl;
-        switch (executionOptions.ExecutionCase) {
+    const size_t numberQueries = executionOptions.ScriptQueries.size();
+    const size_t numberLoops = executionOptions.LoopCount;
+    for (size_t queryId = 0; queryId < numberQueries * numberLoops || numberLoops == 0; ++queryId) {
+        size_t id = queryId % numberQueries;
+        if (id == 0 && queryId > 0) {
+            Sleep(executionOptions.LoopDelay);
+        }
+
+        Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Executing script";
+        if (numberQueries > 1) {
+            Cout << " " << id;
+        }
+        if (numberLoops != 1) {
+            Cout << ", loop " << queryId / numberQueries;
+        }
+        Cout << "..." << colors.Default() << Endl;
+
+        switch (executionOptions.GetExecutionCase(id)) {
         case TExecutionOptions::EExecutionCase::GenericScript:
             if (!runner.ExecuteScript(executionOptions.ScriptQueries[id], executionOptions.ScriptQueryAction, executionOptions.TraceId)) {
                 ythrow yexception() << TInstant::Now().ToIsoStringLocal() << " Script execution failed";
@@ -334,7 +357,10 @@ class TMain : public TMainClassArgs {
             .RequiredArgument("query-type")
             .DefaultValue("script")
             .Choices(executionCase.GetChoices())
-            .StoreMappedResultT<TString>(&ExecutionOptions.ExecutionCase, executionCase);
+            .Handler1([this, executionCase](const NLastGetopt::TOptsParser* option) {
+                TString choice(option->CurValOrDef());
+                ExecutionOptions.ExecutionCases.emplace_back(executionCase(choice));
+            });
 
         TChoices<NKikimrKqp::EQueryAction> scriptAction({
             {"execute", NKikimrKqp::QUERY_ACTION_EXECUTE},
@@ -350,6 +376,15 @@ class TMain : public TMainClassArgs {
             .NoArgument()
             .SetFlag(&ExecutionOptions.ForgetExecution);
 
+        options.AddLongOption("loop-count", "Number of runs of the script query (use 0 to start infinite loop)")
+            .RequiredArgument("uint")
+            .DefaultValue(ExecutionOptions.LoopCount)
+            .StoreResult(&ExecutionOptions.LoopCount);
+        options.AddLongOption("loop-delay", "Delay in milliseconds between loop steps")
+            .RequiredArgument("uint")
+            .DefaultValue(1000)
+            .StoreMappedResultT<ui64>(&ExecutionOptions.LoopDelay, &TDuration::MilliSeconds<ui64>);
+
         // Cluster settings
 
         options.AddLongOption('N', "node-count", "Number of nodes to create")

diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.cpp b/ydb/tests/tools/kqprun/src/ydb_setup.cpp
@@ -198,7 +198,9 @@ class TYdbSetup::TImpl {
         WaitResourcesPublishing();
 
         if (Settings_.MonitoringEnabled) {
-            Cout << CoutColors_.Cyan() << "Monitoring port: " << CoutColors_.Default() << Server_->GetRuntime()->GetMonPort() << Endl;
+            for (ui32 nodeIndex = 0; nodeIndex < Settings_.NodeCount; ++nodeIndex) {
+                Cout << CoutColors_.Cyan() << "Monitoring port" << (Settings_.NodeCount > 1 ? TStringBuilder() << " for node " << nodeIndex + 1 : TString()) << ": " << CoutColors_.Default() << Server_->GetRuntime()->GetMonPort(nodeIndex) << Endl;
+            }
         }
     }