Merge 8be321f into 7cf9db7

GrigoriyPA · web-flow · commit 652e76befd13 · 2024-07-26T11:18:47.000Z
diff --git a/ydb/core/grpc_services/rpc_kqp_base.cpp b/ydb/core/grpc_services/rpc_kqp_base.cpp
@@ -67,6 +67,7 @@ void FillQueryStats(Ydb::TableStats::QueryStats& queryStats, const NKqpProto::TK
     queryStats.set_process_cpu_time_us(kqpStats.GetWorkerCpuTimeUs());
     queryStats.set_total_cpu_time_us(totalCpuTimeUs);
     queryStats.set_total_duration_us(kqpStats.GetDurationUs());
+    queryStats.set_queued_time_us(kqpStats.GetQueuedTimeUs());
 }
 
 void FillQueryStats(Ydb::TableStats::QueryStats& queryStats, const NKikimrKqp::TQueryResponse& kqpResponse) {
diff --git a/ydb/core/kqp/session_actor/kqp_query_state.h b/ydb/core/kqp/session_actor/kqp_query_state.h
@@ -114,6 +114,7 @@ class TKqpQueryState : public TNonCopyable {
     bool IsDocumentApiRestricted_ = false;
 
     TInstant StartTime;
+    TInstant ContinueTime;
     NYql::TKikimrQueryDeadlines QueryDeadlines;
     TKqpQueryStats QueryStats;
     bool KeepSession = false;
diff --git a/ydb/core/kqp/session_actor/kqp_query_stats.cpp b/ydb/core/kqp/session_actor/kqp_query_stats.cpp
@@ -210,6 +210,7 @@ ui64 CalcRequestUnit(const TKqpQueryStats& stats) {
 NKqpProto::TKqpStatsQuery TKqpQueryStats::ToProto() const {
     NKqpProto::TKqpStatsQuery result;
     result.SetDurationUs(DurationUs);
+    result.SetQueuedTimeUs(QueuedTimeUs);
 
     if (Compilation) {
         result.MutableCompilation()->SetFromCache(Compilation->FromCache);
diff --git a/ydb/core/kqp/session_actor/kqp_query_stats.h b/ydb/core/kqp/session_actor/kqp_query_stats.h
@@ -8,6 +8,7 @@ namespace NKikimr::NKqp {
 
 struct TKqpQueryStats {
     ui64 DurationUs = 0;
+    ui64 QueuedTimeUs = 0;
     std::optional<TKqpStatsCompile> Compilation;
 
     ui64 WorkerCpuTimeUs = 0;
diff --git a/ydb/core/kqp/session_actor/kqp_session_actor.cpp b/ydb/core/kqp/session_actor/kqp_session_actor.cpp
@@ -475,6 +475,7 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
 
     void Handle(NWorkload::TEvContinueRequest::TPtr& ev) {
         YQL_ENSURE(QueryState);
+        QueryState->ContinueTime = TInstant::Now();
 
         if (ev->Get()->Status == Ydb::StatusIds::UNSUPPORTED) {
             LOG_T("Failed to place request in resource pool, feature flag is disabled");
@@ -1552,6 +1553,9 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
 
         stats->DurationUs = ((TInstant::Now() - QueryState->StartTime).MicroSeconds());
         stats->WorkerCpuTimeUs = (QueryState->GetCpuTime().MicroSeconds());
+        if (const auto continueTime = QueryState->ContinueTime) {
+            stats->QueuedTimeUs = (continueTime - QueryState->StartTime).MicroSeconds();
+        }
         if (QueryState->CompileResult) {
             stats->Compilation.emplace();
             stats->Compilation->FromCache = (QueryState->CompileStats.FromCache);
diff --git a/ydb/core/kqp/ut/service/kqp_qs_queries_ut.cpp b/ydb/core/kqp/ut/service/kqp_qs_queries_ut.cpp
@@ -245,13 +245,15 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
         auto db = kikimr.GetQueryClient();
 
         TExecuteQuerySettings settings;
+        settings.StatsMode(EStatsMode::Full);
 
         {  // Existing pool
             settings.PoolId("default");
 
             const TString query = "SELECT Key, Value2 FROM TwoShard WHERE Value2 > 0 ORDER BY Key";
             auto result = db.ExecuteQuery(query, TTxControl::BeginTx().CommitTx(), settings).ExtractValueSync();
             CheckQueryResult(result);
+            UNIT_ASSERT_VALUES_UNEQUAL(result.GetStats()->GetQueuedTime(), TDuration::Zero());
         }
 
         {  // Not existing pool (check workload manager enabled)
diff --git a/ydb/core/kqp/ut/service/kqp_qs_scripts_ut.cpp b/ydb/core/kqp/ut/service/kqp_qs_scripts_ut.cpp
@@ -109,13 +109,16 @@ Y_UNIT_TEST_SUITE(KqpQueryServiceScripts) {
         auto db = kikimr.GetQueryClient();
 
         TExecuteScriptSettings settings;
+        settings.StatsMode(EStatsMode::Full);
 
         {  // Existing pool
             settings.PoolId("default");
 
             auto scripOp = db.ExecuteScript("SELECT 42", settings).ExtractValueSync();
             UNIT_ASSERT_VALUES_EQUAL_C(scripOp.Status().GetStatus(), EStatus::SUCCESS, scripOp.Status().GetIssues().ToString());
-            CheckScriptResults(scripOp, WaitScriptExecutionOperation(scripOp.Id(), kikimr.GetDriver()), db);
+            auto readyOp = WaitScriptExecutionOperation(scripOp.Id(), kikimr.GetDriver());
+            CheckScriptResults(scripOp, readyOp, db);
+            UNIT_ASSERT_VALUES_UNEQUAL(readyOp.Metadata().ExecStats.GetQueuedTime(), TDuration::Zero());
         }
 
         {  // Not existing pool (check workload manager enabled)
diff --git a/ydb/core/kqp/workload_service/actors/pool_handlers_acors.cpp b/ydb/core/kqp/workload_service/actors/pool_handlers_acors.cpp
@@ -54,6 +54,12 @@ class TPoolHandlerActorBase : public TActor<TDerived> {
             UpdateConfigCounters(poolConfig);
         }
 
+        void CollectRequestLatency(TInstant continueTime) {
+            if (continueTime) {
+                RequestsLatencyMs->Collect((TInstant::Now() - continueTime).MilliSeconds());
+            }
+        }
+
         void UpdateConfigCounters(const NResourcePool::TPoolSettings& poolConfig) {
             InFlightLimit->Set(std::max(poolConfig.ConcurrentQueryLimit, 0));
             QueueSizeLimit->Set(std::max(poolConfig.QueueSize, 0));
@@ -106,6 +112,7 @@ class TPoolHandlerActorBase : public TActor<TDerived> {
         const TActorId WorkerActorId;
         const TString SessionId;
         const TInstant StartTime = TInstant::Now();
+        TInstant ContinueTime;
 
         EState State = EState::Pending;
         bool Started = false;  // after TEvContinueRequest success
@@ -267,6 +274,7 @@ class TPoolHandlerActorBase : public TActor<TDerived> {
         if (status == Ydb::StatusIds::SUCCESS) {
             LocalInFlight++;
             request->Started = true;
+            request->ContinueTime = TInstant::Now();
             Counters.LocalInFly->Inc();
             Counters.ContinueOk->Inc();
             Counters.DelayedTimeMs->Collect((TInstant::Now() - request->StartTime).MilliSeconds());
@@ -387,7 +395,7 @@ class TPoolHandlerActorBase : public TActor<TDerived> {
 
         if (status == Ydb::StatusIds::SUCCESS) {
             Counters.CleanupOk->Inc();
-            Counters.RequestsLatencyMs->Collect((TInstant::Now() - request->StartTime).MilliSeconds());
+            Counters.CollectRequestLatency(request->ContinueTime);
             LOG_D("Reply cleanup success to " << request->WorkerActorId << ", session id: " << request->SessionId << ", local in flight: " << LocalInFlight);
         } else {
             Counters.CleanupError->Inc();
@@ -401,7 +409,7 @@ class TPoolHandlerActorBase : public TActor<TDerived> {
         this->Send(MakeKqpProxyID(this->SelfId().NodeId()), ev.release());
 
         Counters.Cancelled->Inc();
-        Counters.RequestsLatencyMs->Collect((TInstant::Now() - request->StartTime).MilliSeconds());
+        Counters.CollectRequestLatency(request->ContinueTime);
         LOG_I("Cancel request for worker " << request->WorkerActorId << ", session id: " << request->SessionId << ", local in flight: " << LocalInFlight);
     }
 
diff --git a/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.cpp b/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.cpp
@@ -487,6 +487,7 @@ class TWorkloadServiceYdbSetup : public IYdbSetup {
         request->SetQuery(query);
         request->SetType(NKikimrKqp::QUERY_TYPE_SQL_GENERIC_QUERY);
         request->SetAction(NKikimrKqp::QUERY_ACTION_EXECUTE);
+        request->SetCollectStats(Ydb::Table::QueryStatsCollection::STATS_COLLECTION_FULL);
         request->SetDatabase(Settings_.DomainName_);
         request->SetPoolId(settings.PoolId_);
 
@@ -560,6 +561,10 @@ const std::vector<Ydb::ResultSet>& TQueryRunnerResult::GetResultSets() const {
     return ResultSets;
 }
 
+const NKqpProto::TKqpStatsQuery& TQueryRunnerResult::GetQueryStats() const {
+    return Response.GetResponse().GetQueryStats();
+}
+
 //// TQueryRunnerResultAsync
 
 TQueryRunnerResult TQueryRunnerResultAsync::GetResult(TDuration timeout) const {
diff --git a/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.h b/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.h
@@ -44,6 +44,7 @@ struct TQueryRunnerResult {
 
     const Ydb::ResultSet& GetResultSet(size_t resultIndex) const;
     const std::vector<Ydb::ResultSet>& GetResultSets() const;
+    const NKqpProto::TKqpStatsQuery& GetQueryStats() const;
 };
 
 struct TQueryRunnerResultAsync {
diff --git a/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp b/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp
@@ -2,6 +2,8 @@
 
 #include <ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.h>
 
+#include <ydb/core/protos/kqp_stats.pb.h>
+
 
 namespace NKikimr::NKqp {
 
@@ -44,7 +46,9 @@ Y_UNIT_TEST_SUITE(KqpWorkloadService) {
             .EnableResourcePools(false)
             .Create();
 
-        TSampleQueries::TSelect42::CheckResult(ydb->ExecuteQuery(TSampleQueries::TSelect42::Query, TQueryRunnerSettings().PoolId("another_pool_id")));
+        auto result = ydb->ExecuteQuery(TSampleQueries::TSelect42::Query, TQueryRunnerSettings().PoolId("another_pool_id"));
+        TSampleQueries::TSelect42::CheckResult(result);
+        UNIT_ASSERT_VALUES_EQUAL(result.GetQueryStats().GetQueuedTimeUs(), 0);
     }
 
     TQueryRunnerResultAsync StartQueueSizeCheckRequests(TIntrusivePtr<IYdbSetup> ydb, const TQueryRunnerSettings& settings) {
diff --git a/ydb/core/protos/kqp_stats.proto b/ydb/core/protos/kqp_stats.proto
@@ -73,6 +73,7 @@ message TKqpExecutionExtraStats {
 message TKqpStatsQuery {
     // Basic stats
     uint64 DurationUs = 1;
+    uint64 QueuedTimeUs = 9;
     TKqpStatsCompile Compilation = 2;
 
     reserved 3; // repeated TKqpStatsExecution Executions = 3;
diff --git a/ydb/public/api/protos/ydb_query_stats.proto b/ydb/public/api/protos/ydb_query_stats.proto
@@ -43,4 +43,5 @@ message QueryStats {
     string query_ast = 5;
     uint64 total_duration_us = 6;
     uint64 total_cpu_time_us = 7;
+    uint64 queued_time_us = 8;
 }
diff --git a/ydb/public/sdk/cpp/client/ydb_query/stats.cpp b/ydb/public/sdk/cpp/client/ydb_query/stats.cpp
@@ -60,6 +60,10 @@ TDuration TExecStats::GetTotalDuration() const {
     return TDuration::MicroSeconds(Impl_->Proto.total_duration_us());
 }
 
+TDuration TExecStats::GetQueuedTime() const {
+    return TDuration::MicroSeconds(Impl_->Proto.queued_time_us());
+}
+
 TDuration TExecStats::GetTotalCpuTime() const {
     return TDuration::MicroSeconds(Impl_->Proto.total_cpu_time_us());
 }
diff --git a/ydb/public/sdk/cpp/client/ydb_query/stats.h b/ydb/public/sdk/cpp/client/ydb_query/stats.h
@@ -33,6 +33,7 @@ class TExecStats {
     TMaybe<TString> GetAst() const;
 
     TDuration GetTotalDuration() const;
+    TDuration GetQueuedTime() const;
     TDuration GetTotalCpuTime() const;
 
 private:
diff --git a/ydb/tests/tools/kqprun/src/kqp_runner.cpp b/ydb/tests/tools/kqprun/src/kqp_runner.cpp
@@ -390,6 +390,9 @@ class TKqpRunner::TImpl {
         } else {
             Cout << " Estimated duration: " << TInstant::Now() - StartTime_;
         }
+        if (meta.QueuedTime) {
+            Cout << ", Queued time: " << meta.QueuedTime;
+        }
         Cout << CoutColors_.Default() << Endl;
     }
 
diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.cpp b/ydb/tests/tools/kqprun/src/ydb_setup.cpp
@@ -66,6 +66,7 @@ void FillQueryMeta(TQueryMeta& meta, const NKikimrKqp::TQueryResponse& response)
         meta.Plan = plan;
     }
     meta.TotalDuration = TDuration::MicroSeconds(response.GetQueryStats().GetDurationUs());
+    meta.QueuedTime = TDuration::MicroSeconds(response.GetQueryStats().GetQueuedTimeUs());
 }
 
 }  // anonymous namespace
@@ -471,6 +472,7 @@ TRequestResult TYdbSetup::GetScriptExecutionOperationRequest(const TString& oper
             meta.Plan = deserializedMeta.exec_stats().query_plan();
         }
         meta.TotalDuration = TDuration::MicroSeconds(deserializedMeta.exec_stats().total_duration_us());
+        meta.QueuedTime = TDuration::MicroSeconds(deserializedMeta.exec_stats().queued_time_us());
     }
 
     return TRequestResult(scriptExecutionOperation->Get()->Status, scriptExecutionOperation->Get()->Issues);
diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.h b/ydb/tests/tools/kqprun/src/ydb_setup.h
@@ -17,6 +17,7 @@ struct TQueryMeta {
     TString Ast;
     TString Plan;
     TDuration TotalDuration;
+    TDuration QueuedTime;
 };
 
 

Original file line number	Diff line number	Diff line change
`@@ -67,6 +67,7 @@ void FillQueryStats(Ydb::TableStats::QueryStats& queryStats, const NKqpProto::TK`
`67`	`67`	`queryStats.set_process_cpu_time_us(kqpStats.GetWorkerCpuTimeUs());`
`68`	`68`	`queryStats.set_total_cpu_time_us(totalCpuTimeUs);`
`69`	`69`	`queryStats.set_total_duration_us(kqpStats.GetDurationUs());`
	`70`	`+ queryStats.set_queued_time_us(kqpStats.GetQueuedTimeUs());`
`70`	`71`	`}`
`71`	`72`
`72`	`73`	`void FillQueryStats(Ydb::TableStats::QueryStats& queryStats, const NKikimrKqp::TQueryResponse& kqpResponse) {`
Original file line number	Diff line number	Diff line change
`@@ -43,4 +43,5 @@ message QueryStats {`
`43`	`43`	`string query_ast = 5;`
`44`	`44`	`uint64 total_duration_us = 6;`
`45`	`45`	`uint64 total_cpu_time_us = 7;`
	`46`	`+ uint64 queued_time_us = 8;`
`46`	`47`	`}`
Original file line number	Diff line number	Diff line change
`@@ -60,6 +60,10 @@ TDuration TExecStats::GetTotalDuration() const {`
`60`	`60`	`return TDuration::MicroSeconds(Impl_->Proto.total_duration_us());`
`61`	`61`	`}`
`62`	`62`
	`63`	`+TDuration TExecStats::GetQueuedTime() const {`
	`64`	`+ return TDuration::MicroSeconds(Impl_->Proto.queued_time_us());`
	`65`	`+}`
	`66`	`+`
`63`	`67`	`TDuration TExecStats::GetTotalCpuTime() const {`
`64`	`68`	`return TDuration::MicroSeconds(Impl_->Proto.total_cpu_time_us());`
`65`	`69`	`}`
Original file line number	Diff line number	Diff line change
`@@ -390,6 +390,9 @@ class TKqpRunner::TImpl {`
`390`	`390`	`} else {`
`391`	`391`	`Cout << " Estimated duration: " << TInstant::Now() - StartTime_;`
`392`	`392`	`}`
	`393`	`+ if (meta.QueuedTime) {`
	`394`	`+ Cout << ", Queued time: " << meta.QueuedTime;`
	`395`	`+ }`
`393`	`396`	`Cout << CoutColors_.Default() << Endl;`
`394`	`397`	`}`
`395`	`398`
Original file line number	Diff line number	Diff line change
`@@ -66,6 +66,7 @@ void FillQueryMeta(TQueryMeta& meta, const NKikimrKqp::TQueryResponse& response)`
`66`	`66`	`meta.Plan = plan;`
`67`	`67`	`}`
`68`	`68`	`meta.TotalDuration = TDuration::MicroSeconds(response.GetQueryStats().GetDurationUs());`
	`69`	`+ meta.QueuedTime = TDuration::MicroSeconds(response.GetQueryStats().GetQueuedTimeUs());`
`69`	`70`	`}`
`70`	`71`
`71`	`72`	`} // anonymous namespace`
`@@ -471,6 +472,7 @@ TRequestResult TYdbSetup::GetScriptExecutionOperationRequest(const TString& oper`
`471`	`472`	`meta.Plan = deserializedMeta.exec_stats().query_plan();`
`472`	`473`	`}`
`473`	`474`	`meta.TotalDuration = TDuration::MicroSeconds(deserializedMeta.exec_stats().total_duration_us());`
	`475`	`+ meta.QueuedTime = TDuration::MicroSeconds(deserializedMeta.exec_stats().queued_time_us());`
`474`	`476`	`}`
`475`	`477`
`476`	`478`	`return TRequestResult(scriptExecutionOperation->Get()->Status, scriptExecutionOperation->Get()->Issues);`