Merge 96c7e84 into 202d1ff

Author: GrigoriyPA

ydb/core/kqp/common/events/workload_service.h

@@ -51,19 +51,21 @@ struct TEvContinueRequest : public NActors::TEventLocal<TEvContinueRequest, TKqp
 };
 
 struct TEvCleanupRequest : public NActors::TEventLocal<TEvCleanupRequest, TKqpWorkloadServiceEvents::EvCleanupRequest> {
-    TEvCleanupRequest(const TString& database, const TString& sessionId, const TString& poolId, TDuration duration, TDuration cpuConsumed)
+    TEvCleanupRequest(const TString& database, const TString& sessionId, const TString& poolId, TDuration duration, TDuration cpuConsumed, bool sendResponseOnNotfound = true)
         : Database(database)
         , SessionId(sessionId)
         , PoolId(poolId)
         , Duration(duration)
         , CpuConsumed(cpuConsumed)
+        , SendResponseOnNotfound(sendResponseOnNotfound)
     {}
 
     const TString Database;
     const TString SessionId;
     const TString PoolId;
     const TDuration Duration;
     const TDuration CpuConsumed;
+    const bool SendResponseOnNotfound;
 };
 
 struct TEvCleanupResponse : public NActors::TEventLocal<TEvCleanupResponse, TKqpWorkloadServiceEvents::EvCleanupResponse> {

ydb/core/kqp/session_actor/kqp_session_actor.cpp

@@ -256,6 +256,7 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
             QueryState->UserToken
         ), IEventHandle::FlagTrackDelivery);
 
+        QueryState->PoolHandlerActor = MakeKqpWorkloadServiceId(SelfId().NodeId());
         Become(&TKqpSessionActor::ExecuteState);
     }
 
@@ -2421,6 +2422,7 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
                 hFunc(TEvTxProxySchemeCache::TEvNavigateKeySetResult, HandleNoop);
                 hFunc(TEvents::TEvUndelivered, HandleNoop);
                 hFunc(TEvTxUserProxy::TEvAllocateTxIdResult, HandleNoop);
+                hFunc(TEvKqpExecuter::TEvStreamData, HandleNoop);
                 hFunc(NWorkload::TEvContinueRequest, HandleNoop);
 
                 // always come from WorkerActor

ydb/core/kqp/workload_service/actors/pool_handlers_acors.cpp

@@ -161,6 +161,7 @@ class TPoolHandlerActorBase : public TActor<TDerived> {
         }
 
         SendPoolInfoUpdate(std::nullopt, std::nullopt, Subscribers);
+        this->Send(MakeKqpWorkloadServiceId(this->SelfId().NodeId()), new TEvPrivate::TEvStopPoolHandlerResponse(Database, PoolId));
 
         Counters.OnCleanup(ResetCountersOnStrop);
 
@@ -184,16 +185,16 @@ class TPoolHandlerActorBase : public TActor<TDerived> {
     }
 
     void Handle(TEvPrivate::TEvResolvePoolResponse::TPtr& ev) {
-        this->Send(MakeKqpWorkloadServiceId(this->SelfId().NodeId()), new TEvPrivate::TEvPlaceRequestIntoPoolResponse(Database, PoolId));
-
         auto event = std::move(ev->Get()->Event);
+        const TString& sessionId = event->Get()->SessionId;
+        this->Send(MakeKqpWorkloadServiceId(this->SelfId().NodeId()), new TEvPrivate::TEvPlaceRequestIntoPoolResponse(Database, PoolId, sessionId));
+
         const TActorId& workerActorId = event->Sender;
         if (!InFlightLimit) {
             this->Send(workerActorId, new TEvContinueRequest(Ydb::StatusIds::PRECONDITION_FAILED, PoolId, PoolConfig, {NYql::TIssue(TStringBuilder() << "Resource pool " << PoolId << " was disabled due to zero concurrent query limit")}));
             return;
         }
 
-        const TString& sessionId = event->Get()->SessionId;
         if (LocalSessions.contains(sessionId)) {
             this->Send(workerActorId, new TEvContinueRequest(Ydb::StatusIds::INTERNAL_ERROR, PoolId, PoolConfig, {NYql::TIssue(TStringBuilder() << "Got duplicate session id " << sessionId << " for pool " << PoolId)}));
             return;
@@ -217,7 +218,13 @@ class TPoolHandlerActorBase : public TActor<TDerived> {
         const TActorId& workerActorId = ev->Sender;
 
         TRequest* request = GetRequestSafe(sessionId);
-        if (!request || request->State == TRequest::EState::Canceling) {
+        if (!request) {
+            if (ev->Get()->SendResponseOnNotfound) {
+                this->Send(workerActorId, new TEvCleanupResponse(Ydb::StatusIds::SUCCESS));
+            }
+            return;
+        }
+        if (request->State == TRequest::EState::Canceling) {
             this->Send(workerActorId, new TEvCleanupResponse(Ydb::StatusIds::SUCCESS));
             return;
         }

ydb/core/kqp/workload_service/common/events.h

@@ -29,6 +29,7 @@ struct TEvPrivate {
         EvFinishRequestInPool,
         EvResignPoolHandler,
         EvStopPoolHandler,
+        EvStopPoolHandlerResponse,
         EvCancelRequest,
         EvUpdatePoolSubscription,
 
@@ -128,13 +129,15 @@ struct TEvPrivate {
     };
 
     struct TEvPlaceRequestIntoPoolResponse : public NActors::TEventLocal<TEvPlaceRequestIntoPoolResponse, EvPlaceRequestIntoPoolResponse> {
-        TEvPlaceRequestIntoPoolResponse(const TString& database, const TString& poolId)
+        TEvPlaceRequestIntoPoolResponse(const TString& database, const TString& poolId, const TString& sessionId)
             : Database(database)
             , PoolId(poolId)
+            , SessionId(sessionId)
         {}
 
         const TString Database;
         const TString PoolId;
+        const TString SessionId;
     };
 
     struct TEvFinishRequestInPool : public NActors::TEventLocal<TEvFinishRequestInPool, EvFinishRequestInPool> {
@@ -173,6 +176,16 @@ struct TEvPrivate {
         const bool ResetCounters;
     };
 
+    struct TEvStopPoolHandlerResponse : public NActors::TEventLocal<TEvStopPoolHandlerResponse, EvStopPoolHandlerResponse> {
+        TEvStopPoolHandlerResponse(const TString& database, const TString& poolId)
+            : Database(database)
+            , PoolId(poolId)
+        {}
+
+        const TString Database;
+        const TString PoolId;
+    };
+
     struct TEvCancelRequest : public NActors::TEventLocal<TEvCancelRequest, EvCancelRequest> {
         explicit TEvCancelRequest(const TString& sessionId)
             : SessionId(sessionId)

ydb/core/kqp/workload_service/kqp_workload_service.cpp

@@ -168,6 +168,12 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
 
     void Handle(TEvCleanupRequest::TPtr& ev) {
         const TString& database = ev->Get()->Database;
+        const TString& sessionId = ev->Get()->SessionId;
+        if (GetOrCreateDatabaseState(database)->PendingSessionIds.contains(sessionId)) {
+            PendingCancelRequests[sessionId].emplace_back(std::move(ev));
+            return;
+        }
+
         const TString& poolId = ev->Get()->PoolId;
         auto poolState = GetPoolState(database, poolId);
         if (!poolState) {
@@ -176,7 +182,7 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
         }
 
         LOG_D("Finished request with worker actor " << ev->Sender << ", Database: " << database << ", PoolId: " << poolId << ", SessionId: " << ev->Get()->SessionId);
-        Send(ev->Forward(poolState->PoolHandler));
+        poolState->DoCleanupRequest(std::move(ev));
     }
 
     void Handle(TEvents::TEvWakeup::TPtr& ev) {
@@ -220,6 +226,7 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
         hFunc(TEvPrivate::TEvTablesCreationFinished, Handle);
         hFunc(TEvPrivate::TEvCpuLoadResponse, Handle);
         hFunc(TEvPrivate::TEvResignPoolHandler, Handle);
+        hFunc(TEvPrivate::TEvStopPoolHandlerResponse, Handle);
     )
 
 private:
@@ -265,9 +272,22 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
     void Handle(TEvPrivate::TEvPlaceRequestIntoPoolResponse::TPtr& ev) {
         const TString& database = ev->Get()->Database;
         const TString& poolId = ev->Get()->PoolId;
-        LOG_T("Request placed into pool, Database: " << database << ", PoolId: " << poolId);
+        const TString& sessionId = ev->Get()->SessionId;
+        LOG_T("Request placed into pool, Database: " << database << ", PoolId: " << poolId << ", SessionId: " << sessionId);
 
-        if (auto poolState = GetPoolState(database, poolId)) {
+        GetOrCreateDatabaseState(database)->PendingSessionIds.erase(sessionId);
+
+        auto poolState = GetPoolState(database, poolId);
+        for (auto& event : PendingCancelRequests[sessionId]) {
+            if (poolState) {
+                poolState->DoCleanupRequest(std::move(event));
+            } else {
+                ReplyCleanupError(ev->Sender, Ydb::StatusIds::NOT_FOUND, TStringBuilder() << "Pool " << poolId << " not found");
+            }
+        }
+        PendingCancelRequests.erase(sessionId);
+
+        if (poolState) {
             poolState->PlaceRequestRunning = false;
             poolState->UpdateHandler();
             poolState->StartPlaceRequest();
@@ -388,6 +408,17 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
         }
     }
 
+    void Handle(TEvPrivate::TEvStopPoolHandlerResponse::TPtr& ev) {
+        const TString& database = ev->Get()->Database;
+        const TString& poolId = ev->Get()->PoolId;
+        LOG_T("Got stop pool handler response, Database: " << database << ", PoolId: " << poolId);
+
+        Counters.ActivePools->Dec();
+        if (auto poolState = GetPoolState(database, poolId)) {
+            poolState->PreviousPoolHandlers.erase(ev->Sender);
+        }
+    }
+
 private:
     void InitializeWorkloadService() {
         if (ServiceInitialized) {
@@ -449,7 +480,6 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
         }
         for (const auto& poolKey : poolsToDelete) {
             PoolIdToState.erase(poolKey);
-            Counters.ActivePools->Dec();
         }
 
         if (!PoolIdToState.empty()) {
@@ -567,6 +597,7 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
     bool IdleChecksStarted = false;
     ETablesCreationStatus TablesCreationStatus = ETablesCreationStatus::Cleanup;
     std::unordered_set<TString> PendingHandlers;
+    std::unordered_map<TString, std::vector<TEvCleanupRequest::TPtr>> PendingCancelRequests;
 
     std::unordered_map<TString, TDatabaseState> DatabaseToState;
     std::unordered_map<TString, TPoolState> PoolIdToState;

ydb/core/kqp/workload_service/kqp_workload_service_impl.h

@@ -19,6 +19,7 @@ struct TDatabaseState {
     bool& EnabledResourcePoolsOnServerless;
 
     std::vector<TEvPlaceRequestIntoPool::TPtr> PendingRequersts = {};
+    std::unordered_set<TString> PendingSessionIds = {};
     std::unordered_map<TString, std::unordered_set<TActorId>> PendingSubscriptions = {};
     bool HasDefaultPool = false;
     bool Serverless = false;
@@ -38,6 +39,7 @@ struct TDatabaseState {
 
     void DoPlaceRequest(TEvPlaceRequestIntoPool::TPtr ev) {
         TString database = ev->Get()->Database;
+        PendingSessionIds.emplace(ev->Get()->SessionId);
         PendingRequersts.emplace_back(std::move(ev));
 
         if (!EnabledResourcePoolsOnServerless && (TInstant::Now() - LastUpdateTime) > IDLE_DURATION) {
@@ -112,6 +114,7 @@ struct TPoolState {
     bool WaitingInitialization = false;
     bool PlaceRequestRunning = false;
     std::optional<TActorId> NewPoolHandler = std::nullopt;
+    std::unordered_set<TActorId> PreviousPoolHandlers = {};
 
     ui64 InFlightRequests = 0;
     TInstant LastUpdateTime = TInstant::Now();
@@ -122,6 +125,7 @@ struct TPoolState {
         }
 
         ActorContext.Send(PoolHandler, new TEvPrivate::TEvStopPoolHandler(false));
+        PreviousPoolHandlers.insert(PoolHandler);
         PoolHandler = *NewPoolHandler;
         NewPoolHandler = std::nullopt;
         InFlightRequests = 0;
@@ -143,6 +147,16 @@ struct TPoolState {
         InFlightRequests--;
         LastUpdateTime = TInstant::Now();
     }
+
+    void DoCleanupRequest(TEvCleanupRequest::TPtr event) {
+        for (const auto& poolHandler : PreviousPoolHandlers) {
+            ActorContext.Send(poolHandler, new TEvCleanupRequest(
+                event->Get()->Database, event->Get()->SessionId, event->Get()->PoolId,
+                event->Get()->Duration, event->Get()->CpuConsumed, false
+            ));
+        }
+        ActorContext.Send(event->Forward(PoolHandler));
+    }
 };
 
 struct TCpuQuotaManagerState {