Remove query from Compute Scheduler in all cases (#27598) (#27732)

Author: maximyurchuk

ydb/core/kqp/executer_actor/kqp_executer_impl.h

@@ -608,7 +608,7 @@ class TKqpExecuterBase : public TActor<TDerived> {
 
         LWTRACK(KqpBaseExecuterHandleReady, ResponseEv->Orbit, TxId);
 
-        if (!databaseId.empty() && (poolId != NResourcePool::DEFAULT_POOL_ID || AccountDefaultPoolInScheduler)) {
+        if (IsSchedulable()) {
             const auto schedulerServiceId = MakeKqpSchedulerServiceId(SelfId().NodeId());
 
             // TODO: deliberately create the database here - since database doesn't have any useful scheduling properties for now.
@@ -1193,6 +1193,12 @@ class TKqpExecuterBase : public TActor<TDerived> {
         Request.Transactions.crop(0);
         this->Send(Target, ResponseEv.release());
 
+        if (IsSchedulable()) {
+            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());
 
@@ -1269,6 +1275,12 @@ class TKqpExecuterBase : public TActor<TDerived> {
         return TasksGraph.ArenaSerializeTaskToProto(task, serializeAsyncIoSettings);
     }
 
+    inline bool IsSchedulable() const {
+        const auto& databaseId = GetUserRequestContext()->DatabaseId;
+        const auto& poolId = GetUserRequestContext()->PoolId.empty() ? NResourcePool::DEFAULT_POOL_ID : GetUserRequestContext()->PoolId;
+        return !databaseId.empty() && (poolId != NResourcePool::DEFAULT_POOL_ID || AccountDefaultPoolInScheduler);
+    }
+
 protected:
     IKqpGateway::TExecPhysicalRequest Request;
     NYql::NDq::IDqAsyncIoFactory::TPtr AsyncIoFactory;

ydb/core/kqp/executer_actor/ut/kqp_executer_ut.cpp renamed to ydb/core/kqp/executer_actor/kqp_executer_ut.cpp

@@ -17,6 +17,7 @@ Y_UNIT_TEST_SUITE(KqpExecuter) {
         - Start query execution and receive TEvTxRequest.
         - When sending TEvAddQuery from executer to scheduler, immediately receive TEvAbortExecution.
         - Imitate receiving TEvQueryResponse before receiving self TEvPoison by executer.
+        - Check that scheduler got TEvRemoveQuery.
         - Do not crash or get undefined behavior.
      */
     Y_UNIT_TEST(TestSuddenAbortAfterReady) {
@@ -27,27 +28,18 @@ 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) {
-            {
-                TStringStream ss;
-                ss << "Got " << ev->GetTypeName() << " " << ev->Recipient << " " << ev->Sender << Endl;
-                Cerr << ss.Str();
-            }
+            Cerr << (TStringBuilder() << "Got " << ev->GetTypeName() << " " << ev->Recipient << " " << ev->Sender << Endl);
 
             if (ev->GetTypeRewrite() == TEvKqpExecuter::TEvTxRequest::EventType) {
                 targetId = ActorIdFromProto(ev->Get<TEvKqpExecuter::TEvTxRequest>()->Record.GetTarget());
             }
 
             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 +52,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);
 
@@ -141,4 +138,4 @@ Y_UNIT_TEST_SUITE(KqpExecuter) {
     */
 }
 
-} // namespace NKikimr
+} // namespace NKikimr::NKqp

ydb/core/kqp/executer_actor/ut/ya.make

@@ -1,10 +1,4 @@
-UNITTEST_FOR(ydb/core/kqp)
-
-FORK_SUBTESTS()
-
-IF (SANITIZER_TYPE OR WITH_VALGRIND)
-    SIZE(MEDIUM)
-ENDIF()
+UNITTEST_FOR(ydb/core/kqp/executer_actor)
 
 SRCS(
     kqp_executer_ut.cpp

ydb/core/kqp/node_service/kqp_node_state.h

@@ -104,9 +104,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;
                         const auto& actorCtx = NActors::TActorContext::AsActorContext();
                         actorCtx.Send(MakeKqpSchedulerServiceId(actorCtx.SelfID.NodeId()), removeQueryEvent.Release());
                     }

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() {

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> {

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

ydb/library/services/services.proto

@@ -249,6 +249,7 @@ enum EServiceKikimr {
     KQP_SESSION = 545;
     KQP_COMPILE_COMPUTATION_PATTERN_SERVICE = 546;
     KQP_WORKLOAD_SERVICE = 547;
+    KQP_COMPUTE_SCHEDULER = 548;
 
     TABLET_RESOURCE_BROKER = 540;