Remove query from Compute Scheduler in all cases

abyss7 · abyss7 · commit 4ed0d32fa922 · 2025-10-25T17:11:59.000+03:00
diff --git a/ydb/core/kqp/executer_actor/kqp_executer_impl.h b/ydb/core/kqp/executer_actor/kqp_executer_impl.h
@@ -1197,6 +1197,12 @@ class TKqpExecuterBase : public TActor<TDerived> {
         Request.Transactions.crop(0);
         this->Send(Target, ResponseEv.release());
 
+        {
+            auto removeQueryEvent = MakeHolder<NScheduler::TEvRemoveQuery>();
+            removeQueryEvent->QueryId = TxId;
+            this->Send(MakeKqpSchedulerServiceId(SelfId().NodeId()), removeQueryEvent.Release());
+        }
+
         for (auto channelPair: ResultChannelProxies) {
             LOG_D("terminate result channel " << channelPair.first << " proxy at " << channelPair.second->SelfId());
 
diff --git a/ydb/core/kqp/executer_actor/ut/kqp_executer_ut.cpp b/ydb/core/kqp/executer_actor/ut/kqp_executer_ut.cpp
@@ -27,16 +27,11 @@ Y_UNIT_TEST_SUITE(KqpExecuter) {
         auto db = kikimr.RunCall([&] { return kikimr.GetTableClient(); } );
         auto session = kikimr.RunCall([&] { return db.CreateSession().GetValueSync().GetSession(); } );
 
-        auto prepareResult = kikimr.RunCall([&] { return session.PrepareDataQuery(Q_(R"(
-                SELECT COUNT(*) FROM `/Root/TwoShard`;
-            )")).GetValueSync();
-        });
-        UNIT_ASSERT_VALUES_EQUAL_C(prepareResult.GetStatus(), EStatus::SUCCESS, prepareResult.GetIssues().ToString());
-        auto dataQuery = prepareResult.GetQuery();
-
         TActorId executerId, targetId;
+        ui8 queries = 0;
         auto& runtime = *kikimr.GetTestServer().GetRuntime();
         runtime.SetObserverFunc([&](TAutoPtr<IEventHandle>& ev) {
+            // TODO: remove debug logging
             {
                 TStringStream ss;
                 ss << "Got " << ev->GetTypeName() << " " << ev->Recipient << " " << ev->Sender << Endl;
@@ -48,6 +43,7 @@ Y_UNIT_TEST_SUITE(KqpExecuter) {
             }
 
             if (ev->GetTypeRewrite() == NScheduler::TEvAddQuery::EventType) {
+                ++queries;
                 executerId = ev->Sender;
                 auto* abortExecution = new TEvKqp::TEvAbortExecution(NYql::NDqProto::StatusIds::UNSPECIFIED, NYql::TIssues());
                 runtime.Send(new IEventHandle(ev->Sender, targetId, abortExecution));
@@ -60,11 +56,16 @@ Y_UNIT_TEST_SUITE(KqpExecuter) {
             return TTestActorRuntime::EEventAction::PROCESS;
         });
 
-        auto future = kikimr.RunInThreadPool([&] { return dataQuery.Execute(TTxControl::BeginTx().CommitTx(), TExecDataQuerySettings()).GetValueSync(); });
+        auto future = kikimr.RunInThreadPool([&] {
+            return session.ExecuteDataQuery("SELECT COUNT(*) FROM `/Root/TwoShard`;", TTxControl::BeginTx().CommitTx()).ExtractValueSync();
+        });
 
         TDispatchOptions opts;
         opts.FinalEvents.emplace_back([&](IEventHandle& ev) {
-            return ev.GetTypeRewrite() == TEvKqpExecuter::TEvTxResponse::EventType;
+            if (ev.GetTypeRewrite() == NScheduler::TEvRemoveQuery::EventType) {
+                --queries;
+            }
+            return (ev.GetTypeRewrite() == TEvKqpExecuter::TEvTxResponse::EventType || ev.GetTypeRewrite() == NScheduler::TEvRemoveQuery::EventType) && !queries;
         });
         runtime.DispatchEvents(opts);
 
diff --git a/ydb/core/kqp/node_service/kqp_node_state.h b/ydb/core/kqp/node_service/kqp_node_state.h
@@ -103,9 +103,9 @@ class TState {
                         ExpiringRequests.erase(expireIt);
                     }
 
-                    if (auto query = requestIt->second.Query) {
+                    if (requestIt->second.Query) {
                         auto removeQueryEvent = MakeHolder<NScheduler::TEvRemoveQuery>();
-                        removeQueryEvent->Query = query;
+                        removeQueryEvent->QueryId = txId;
                         actorSystem->Send(MakeKqpSchedulerServiceId(actorSystem->NodeId), removeQueryEvent.Release());
                     }
 
diff --git a/ydb/core/kqp/runtime/scheduler/kqp_compute_scheduler_service.cpp b/ydb/core/kqp/runtime/scheduler/kqp_compute_scheduler_service.cpp
@@ -118,6 +118,7 @@ class TComputeSchedulerService : public NActors::TActorBootstrapped<TComputeSche
                 // TODO: Scheduler->UpdatePool(…);
             }
         } else {
+            LOG_ERROR_S(*NActors::TlsActivationContext, NKikimrServices::KQP_COMPUTE_SCHEDULER, "Trying to remove unknown pool: " << databaseId << "/" << poolId);
             // TODO: the removing message for unknown pool - should we check?
         }
     }
@@ -137,7 +138,9 @@ class TComputeSchedulerService : public NActors::TActorBootstrapped<TComputeSche
     }
 
     void Handle(TEvRemoveQuery::TPtr& ev) {
-        Scheduler->RemoveQuery(ev->Get()->Query);
+        if (!Scheduler->RemoveQuery(ev->Get()->QueryId)) {
+            LOG_ERROR_S(*NActors::TlsActivationContext, NKikimrServices::KQP_COMPUTE_SCHEDULER, "Trying to remove unknown query: " << ev->Get()->QueryId);
+        }
     }
 
     void Handle(NActors::TEvents::TEvWakeup::TPtr&) {
@@ -226,7 +229,7 @@ void TComputeScheduler::AddOrUpdatePool(const TString& databaseId, const TString
     }
 }
 
-TQueryPtr TComputeScheduler::AddOrUpdateQuery(const TString& databaseId, const TString& poolId, const NHdrf::TQueryId& queryId, const NHdrf::TStaticAttributes& attrs) {
+TQueryPtr TComputeScheduler::AddOrUpdateQuery(const NHdrf::TDatabaseId& databaseId, const NHdrf::TPoolId& poolId, const NHdrf::TQueryId& queryId, const NHdrf::TStaticAttributes& attrs) {
     Y_ENSURE(!poolId.empty());
 
     TWriteGuard lock(Mutex);
@@ -248,14 +251,16 @@ TQueryPtr TComputeScheduler::AddOrUpdateQuery(const TString& databaseId, const T
     return query;
 }
 
-void TComputeScheduler::RemoveQuery(const TQueryPtr& query) {
-    Y_ENSURE(query);
-
+bool TComputeScheduler::RemoveQuery(const NHdrf::TQueryId& queryId) {
     TWriteGuard lock(Mutex);
-    const auto& queryId = std::get<NHdrf::TQueryId>(query->GetId());
 
-    Y_ENSURE(Queries.erase(queryId));
-    query->GetParent()->RemoveQuery(queryId);
+    if (auto queryIt = Queries.find(queryId); queryIt != Queries.end()) {
+        queryIt->second->GetParent()->RemoveQuery(queryId);
+        Queries.erase(queryIt);
+        return true;
+    }
+
+    return false;
 }
 
 void TComputeScheduler::UpdateFairShare() {
diff --git a/ydb/core/kqp/runtime/scheduler/kqp_compute_scheduler_service.h b/ydb/core/kqp/runtime/scheduler/kqp_compute_scheduler_service.h
@@ -15,12 +15,12 @@ class TComputeScheduler : public std::enable_shared_from_this<TComputeScheduler>
     void SetTotalCpuLimit(ui64 cpu);
     ui64 GetTotalCpuLimit() const;
 
-    void AddOrUpdateDatabase(const TString& databaseId, const NHdrf::TStaticAttributes& attrs);
+    void AddOrUpdateDatabase(const NHdrf::TDatabaseId& databaseId, const NHdrf::TStaticAttributes& attrs);
 
-    void AddOrUpdatePool(const TString& databaseId, const TString& poolId, const NHdrf::TStaticAttributes& attrs);
+    void AddOrUpdatePool(const NHdrf::TDatabaseId& databaseId, const NHdrf::TPoolId& poolId, const NHdrf::TStaticAttributes& attrs);
 
-    NHdrf::NDynamic::TQueryPtr AddOrUpdateQuery(const TString& databaseId, const TString& poolId, const NHdrf::TQueryId& queryId, const NHdrf::TStaticAttributes& attrs);
-    void RemoveQuery(const NHdrf::NDynamic::TQueryPtr& query);
+    NHdrf::NDynamic::TQueryPtr AddOrUpdateQuery(const NHdrf::TDatabaseId& databaseId, const NHdrf::TPoolId& poolId, const NHdrf::TQueryId& queryId, const NHdrf::TStaticAttributes& attrs);
+    bool RemoveQuery(const NHdrf::TQueryId& queryId);
 
     void UpdateFairShare();
 
@@ -89,7 +89,7 @@ struct TEvAddQuery : public TEventLocal<TEvAddQuery, TEvents::EvAddQuery> {
 };
 
 struct TEvRemoveQuery : public TEventLocal<TEvRemoveQuery, TEvents::EvRemoveQuery> {
-    NHdrf::NDynamic::TQueryPtr Query;
+    NHdrf::TQueryId QueryId;
 };
 
 struct TEvQueryResponse : public TEventLocal<TEvQueryResponse, TEvents::EvQueryResponse> {
diff --git a/ydb/core/kqp/runtime/scheduler/kqp_schedulable_task.h b/ydb/core/kqp/runtime/scheduler/kqp_schedulable_task.h
@@ -28,7 +28,7 @@ struct TSchedulableTask : public std::enable_shared_from_this<TSchedulableTask>
     void IncreaseUsage();
     void DecreaseUsage(const TDuration& burstUsage, bool forcedResume);
 
-    // Returns parent pool's 'fair-share' - 'usage'
+    // Returns parent pool's 'fair-share' minus 'usage'
     size_t GetSpareUsage() const;
 
     // Account extra usage which doesn't affect scheduling
diff --git a/ydb/library/services/services.proto b/ydb/library/services/services.proto
@@ -256,6 +256,7 @@ enum EServiceKikimr {
     KQP_SESSION = 545;
     KQP_COMPILE_COMPUTATION_PATTERN_SERVICE = 546;
     KQP_WORKLOAD_SERVICE = 547;
+    KQP_COMPUTE_SCHEDULER = 548;
 
     TABLET_RESOURCE_BROKER = 540;
 

Original file line number	Diff line number	Diff line change
`@@ -103,9 +103,9 @@ class TState {`
`103`	`103`	`ExpiringRequests.erase(expireIt);`
`104`	`104`	`}`
`105`	`105`
`106`		`- if (auto query = requestIt->second.Query) {`
	`106`	`+ if (requestIt->second.Query) {`
`107`	`107`	`auto removeQueryEvent = MakeHolder<NScheduler::TEvRemoveQuery>();`
`108`		`- removeQueryEvent->Query = query;`
	`108`	`+ removeQueryEvent->QueryId = txId;`
`109`	`109`	`actorSystem->Send(MakeKqpSchedulerServiceId(actorSystem->NodeId), removeQueryEvent.Release());`
`110`	`110`	`}`
`111`	`111`
Original file line number	Diff line number	Diff line change
`@@ -118,6 +118,7 @@ class TComputeSchedulerService : public NActors::TActorBootstrapped<TComputeSche`
`118`	`118`	`// TODO: Scheduler->UpdatePool(…);`
`119`	`119`	`}`
`120`	`120`	`} else {`
	`121`	`+ LOG_ERROR_S(*NActors::TlsActivationContext, NKikimrServices::KQP_COMPUTE_SCHEDULER, "Trying to remove unknown pool: " << databaseId << "/" << poolId);`
`121`	`122`	`// TODO: the removing message for unknown pool - should we check?`
`122`	`123`	`}`
`123`	`124`	`}`
`@@ -137,7 +138,9 @@ class TComputeSchedulerService : public NActors::TActorBootstrapped<TComputeSche`
`137`	`138`	`}`
`138`	`139`
`139`	`140`	`void Handle(TEvRemoveQuery::TPtr& ev) {`
`140`		`- Scheduler->RemoveQuery(ev->Get()->Query);`
	`141`	`+ if (!Scheduler->RemoveQuery(ev->Get()->QueryId)) {`
	`142`	`+ LOG_ERROR_S(*NActors::TlsActivationContext, NKikimrServices::KQP_COMPUTE_SCHEDULER, "Trying to remove unknown query: " << ev->Get()->QueryId);`
	`143`	`+ }`
`141`	`144`	`}`
`142`	`145`
`143`	`146`	`void Handle(NActors::TEvents::TEvWakeup::TPtr&) {`
`@@ -226,7 +229,7 @@ void TComputeScheduler::AddOrUpdatePool(const TString& databaseId, const TString`
`226`	`229`	`}`
`227`	`230`	`}`
`228`	`231`
`229`		`-TQueryPtr TComputeScheduler::AddOrUpdateQuery(const TString& databaseId, const TString& poolId, const NHdrf::TQueryId& queryId, const NHdrf::TStaticAttributes& attrs) {`
	`232`	`+TQueryPtr TComputeScheduler::AddOrUpdateQuery(const NHdrf::TDatabaseId& databaseId, const NHdrf::TPoolId& poolId, const NHdrf::TQueryId& queryId, const NHdrf::TStaticAttributes& attrs) {`
`230`	`233`	`Y_ENSURE(!poolId.empty());`
`231`	`234`
`232`	`235`	`TWriteGuard lock(Mutex);`
`@@ -248,14 +251,16 @@ TQueryPtr TComputeScheduler::AddOrUpdateQuery(const TString& databaseId, const T`
`248`	`251`	`return query;`
`249`	`252`	`}`
`250`	`253`
`251`		`-void TComputeScheduler::RemoveQuery(const TQueryPtr& query) {`
`252`		`- Y_ENSURE(query);`
`253`		`-`
	`254`	`+bool TComputeScheduler::RemoveQuery(const NHdrf::TQueryId& queryId) {`
`254`	`255`	`TWriteGuard lock(Mutex);`
`255`		`- const auto& queryId = std::get<NHdrf::TQueryId>(query->GetId());`
`256`	`256`
`257`		`- Y_ENSURE(Queries.erase(queryId));`
`258`		`- query->GetParent()->RemoveQuery(queryId);`
	`257`	`+ if (auto queryIt = Queries.find(queryId); queryIt != Queries.end()) {`
	`258`	`+ queryIt->second->GetParent()->RemoveQuery(queryId);`
	`259`	`+ Queries.erase(queryIt);`
	`260`	`+ return true;`
	`261`	`+ }`
	`262`	`+`
	`263`	`+ return false;`
`259`	`264`	`}`
`260`	`265`
`261`	`266`	`void TComputeScheduler::UpdateFairShare() {`