YQ fixed hanging and TLI errors for distributed WM (#6263)

GrigoriyPA · web-flow · commit f60b60f9de61 · 2024-07-05T09:20:28.000+03:00
diff --git a/ydb/core/kqp/workload_service/actors/pool_handlers_acors.cpp b/ydb/core/kqp/workload_service/actors/pool_handlers_acors.cpp
@@ -1,5 +1,7 @@
 #include "actors.h"
 
+#include <ydb/core/base/path.h>
+
 #include <ydb/core/kqp/common/events/events.h>
 #include <ydb/core/kqp/common/simple/services.h>
 #include <ydb/core/kqp/workload_service/common/events.h>
@@ -46,7 +48,7 @@ class TPoolHandlerActorBase : public TActor<TDerived> {
     TPoolHandlerActorBase(void (TDerived::* requestFunc)(TAutoPtr<IEventHandle>& ev), const TString& database, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, NMonitoring::TDynamicCounterPtr counters)
         : TBase(requestFunc)
         , CountersRoot(counters)
-        , CountersSubgroup(counters->GetSubgroup("pool", TStringBuilder() << database << "/" << poolId))
+        , CountersSubgroup(counters->GetSubgroup("pool", CanonizePath(TStringBuilder() << database << "/" << poolId)))
         , Database(database)
         , PoolId(poolId)
         , QueueSizeLimit(GetMaxQueueSize(poolConfig))
@@ -475,6 +477,13 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBase<TFifoPoolHandlerActor
         }
     }
 
+    void PassAway() override {
+        GlobalInFly->Set(0);
+        GlobalDelayedRequests->Set(0);
+
+        TBase::PassAway();
+    }
+
 protected:
     bool ShouldResign() const override {
         return InFlightLimit == 0 || InFlightLimit == std::numeric_limits<ui64>::max();
@@ -519,6 +528,7 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBase<TFifoPoolHandlerActor
         RemoveFinishedRequests();
         RefreshRequired |= !PendingRequests.empty();
         RefreshRequired |= (GetLocalInFlight() || !DelayedRequests.empty()) && TInstant::Now() - LastRefreshTime > LEASE_DURATION / 4;
+        RefreshRequired |= GlobalState.AmountRequests() && TInstant::Now() - LastRefreshTime > LEASE_DURATION;
         if (RefreshRequired) {
             RefreshRequired = false;
             RunningOperation = true;
@@ -532,7 +542,7 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBase<TFifoPoolHandlerActor
         LOG_T("Try to start scheduled refresh");
 
         RefreshState();
-        if (GetLocalInFlight() + DelayedRequests.size() > 0) {
+        if (GetLocalSessionsCount() || GlobalState.AmountRequests()) {
             ScheduleRefresh();
         }
     }
@@ -564,6 +574,7 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBase<TFifoPoolHandlerActor
         if (ev->Get()->Status != Ydb::StatusIds::SUCCESS) {
             LOG_E("refresh pool state failed " << ev->Get()->Status << ", issues: " << ev->Get()->Issues.ToOneLineString());
             RefreshRequired = true;
+            ScheduleRefresh();
             return;
         }
 
@@ -573,6 +584,9 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBase<TFifoPoolHandlerActor
         LastRefreshTime = TInstant::Now();
 
         GlobalState = ev->Get()->PoolState;
+        if (GlobalState.AmountRequests()) {
+            ScheduleRefresh();
+        }
         GlobalInFly->Set(GlobalState.RunningRequests);
         GlobalDelayedRequests->Set(GlobalState.DelayedRequests);
         LOG_T("succefully refreshed pool state, in flight: " << GlobalState.RunningRequests << ", delayed: " << GlobalState.DelayedRequests);
@@ -598,10 +612,12 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBase<TFifoPoolHandlerActor
 
         DoDelayRequest();
         DoStartDelayedRequest();
+        RefreshState();
     };
 
     void Handle(TEvPrivate::TEvDelayRequestResponse::TPtr& ev) {
         RunningOperation = false;
+        ScheduleRefresh();
 
         if (ev->Get()->Status != Ydb::StatusIds::SUCCESS) {
             LOG_E("failed to delay request " << ev->Get()->Status << ", session id: " << ev->Get()->SessionId << ", issues: " << ev->Get()->Issues.ToOneLineString());
@@ -617,7 +633,6 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBase<TFifoPoolHandlerActor
         GlobalDelayedRequests->Inc();
         LOG_D("succefully delayed request, session id: " << ev->Get()->SessionId);
 
-        ScheduleRefresh();
         DoStartDelayedRequest();
         RefreshState();
     };
@@ -637,12 +652,13 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBase<TFifoPoolHandlerActor
             return;
         }
 
-        if (!sessionId) {
-            LOG_D("first request in queue is remote, send notification to node " << ev->Get()->NodeId);
+        const ui32 nodeId = ev->Get()->NodeId;
+        if (SelfId().NodeId() != nodeId) {
+            LOG_D("first request in queue is remote, send notification to node " << nodeId);
             auto event = std::make_unique<TEvPrivate::TEvRefreshPoolState>();
             event->Record.SetPoolId(PoolId);
             event->Record.SetDatabase(Database);
-            this->Send(MakeKqpWorkloadServiceId(ev->Get()->NodeId), std::move(event));
+            this->Send(MakeKqpWorkloadServiceId(nodeId), std::move(event));
             RefreshState();
             return;
         }
@@ -680,7 +696,7 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBase<TFifoPoolHandlerActor
         }
 
         for (const TString& sessionId : ev->Get()->SesssionIds) {
-            LOG_T("succefully cleanuped request, session id: " << sessionId);
+            LOG_T("cleanuped request, session id: " << sessionId);
             if (TRequest* request = GetRequestSafe(sessionId)) {
                 FinalReply(request, ev->Get()->Status, ev->Get()->Issues);
             }
@@ -696,7 +712,7 @@ class TFifoPoolHandlerActor : public TPoolHandlerActorBase<TFifoPoolHandlerActor
             return;
         }
 
-        if (!PendingRequests.empty() && DelayedRequests.empty() && GlobalState.DelayedRequests == 0 && GlobalState.RunningRequests < InFlightLimit) {
+        if (!PendingRequests.empty() && QueueSizeLimit == 0 && GlobalState.RunningRequests < InFlightLimit) {
             RunningOperation = true;
             const TString& sessionId = PopPendingRequest();
             this->Register(CreateStartRequestActor(this->SelfId(), Database, PoolId, sessionId, LEASE_DURATION, CountersSubgroup));
diff --git a/ydb/core/kqp/workload_service/tables/table_queries.cpp b/ydb/core/kqp/workload_service/tables/table_queries.cpp
@@ -92,7 +92,7 @@ class TTablesCreator : public NTableCreator::TMultiTableCreator {
                 Col("wait_deadline", NScheme::NTypeIds::Timestamp),
                 Col("lease_deadline", NScheme::NTypeIds::Timestamp),
             },
-            { "database", "pool_id", "start_time", "session_id" },
+            { "database", "pool_id", "node_id", "start_time", "session_id" },
             NKikimrServices::KQP_WORKLOAD_SERVICE,
             TtlCol("lease_deadline", DEADLINE_OFFSET, BRO_RUN_INTERVAL)
         );
@@ -383,7 +383,7 @@ class TRefreshPoolStateQuery : public TQueryBase {
                 .Utf8(PoolId)
                 .Build();
 
-        RunDataQuery(sql, &params);
+        RunDataQuery(sql, &params, TTxControl::BeginAndCommitTx(true));
         SetQueryResultHandler(&TRefreshPoolStateQuery::OnQueryResult, "Describe pool");
     }
 
@@ -515,6 +515,8 @@ class TStartFirstDelayedRequestQuery : public TQueryBase {
     {}
 
     void OnRunQuery() override {
+        RequestNodeId = SelfId().NodeId();
+
         TString sql = TStringBuilder() << R"(
             -- TStartFirstDelayedRequestQuery::OnRunQuery
             DECLARE $database AS Text;
@@ -539,7 +541,7 @@ class TStartFirstDelayedRequestQuery : public TQueryBase {
                 .Utf8(PoolId)
                 .Build();
 
-        RunDataQuery(sql, &params);
+        RunDataQuery(sql, &params, TTxControl::BeginAndCommitTx(true));
         SetQueryResultHandler(&TStartFirstDelayedRequestQuery::OnGetPoolInfo, "Describe pool");
     }
 
@@ -599,6 +601,7 @@ class TStartFirstDelayedRequestQuery : public TQueryBase {
             DELETE FROM `)" << TTablesCreator::GetDelayedRequestsPath() << R"(`
             WHERE database = $database
               AND pool_id = $pool_id
+              AND node_id = $node_id
               AND start_time = $start_time
               AND session_id = $session_id;
 
@@ -780,6 +783,7 @@ class TCleanupRequestsQuery : public TQueryBase {
             DELETE FROM `)" << TTablesCreator::GetDelayedRequestsPath() << R"(`
             WHERE database = $database
               AND pool_id = $pool_id
+              AND node_id = $node_id
               AND session_id IN $session_ids;
 
             DELETE FROM `)" << TTablesCreator::GetRunningRequestsPath() << R"(`
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
@@ -404,9 +404,7 @@ class TWorkloadServiceYdbSetup : public IYdbSetup {
     }
 
     void WaitPoolHandlersCount(i64 finalCount, std::optional<i64> initialCount = std::nullopt, TDuration timeout = FUTURE_WAIT_TIMEOUT) const override {
-        auto counter = GetServiceCounters(GetRuntime()->GetAppData().Counters, "kqp")
-            ->GetSubgroup("subsystem", "workload_manager")
-            ->GetCounter("ActivePoolHandlers");
+        auto counter = GetWorkloadManagerCounters(0)->GetCounter("ActivePoolHandlers");
 
         if (initialCount) {
             UNIT_ASSERT_VALUES_EQUAL_C(counter->Val(), *initialCount, "Unexpected pool handlers count");
@@ -429,6 +427,18 @@ class TWorkloadServiceYdbSetup : public IYdbSetup {
         Sleep(TDuration::Seconds(1));
     }
 
+    void ValidateWorkloadServiceCounters(bool checkTableCounters = true, const TString& poolId = "") const override {
+        for (ui32 nodeIndex = 0; nodeIndex < Settings_.NodeCount_; ++nodeIndex) {
+            auto subgroup = GetWorkloadManagerCounters(nodeIndex)
+                ->GetSubgroup("pool", CanonizePath(TStringBuilder() << Settings_.DomainName_ << "/" << (poolId ? poolId : Settings_.PoolId_)));
+
+            CheckCommonCounters(subgroup);
+            if (checkTableCounters) {
+                CheckTableCounters(subgroup);
+            }
+        }
+    }
+
     TTestActorRuntime* GetRuntime() const override {
         return Server_->GetRuntime();
     }
@@ -460,6 +470,45 @@ class TWorkloadServiceYdbSetup : public IYdbSetup {
         return event;
     }
 
+    NMonitoring::TDynamicCounterPtr GetWorkloadManagerCounters(ui32 nodeIndex) const {
+        return GetServiceCounters(GetRuntime()->GetAppData(nodeIndex).Counters, "kqp")
+            ->GetSubgroup("subsystem", "workload_manager");
+    }
+
+    static void CheckCommonCounters(NMonitoring::TDynamicCounterPtr subgroup) {
+        UNIT_ASSERT_VALUES_EQUAL(subgroup->GetCounter("LocalInFly", false)->Val(), 0);
+        UNIT_ASSERT_VALUES_EQUAL(subgroup->GetCounter("LocalDelayedRequests", false)->Val(), 0);
+        UNIT_ASSERT_VALUES_EQUAL(subgroup->GetCounter("ContinueOverloaded", true)->Val(), 0);
+        UNIT_ASSERT_VALUES_EQUAL(subgroup->GetCounter("ContinueError", true)->Val(), 0);
+        UNIT_ASSERT_VALUES_EQUAL(subgroup->GetCounter("CleanupError", true)->Val(), 0);
+        UNIT_ASSERT_VALUES_EQUAL(subgroup->GetCounter("Cancelled", true)->Val(), 0);
+
+        UNIT_ASSERT_GE(subgroup->GetCounter("ContinueOk", true)->Val(), 1);
+        UNIT_ASSERT_VALUES_EQUAL(subgroup->GetCounter("ContinueOk", true)->Val(), subgroup->GetCounter("CleanupOk", true)->Val());
+    }
+
+    static void CheckTableCounters(NMonitoring::TDynamicCounterPtr subgroup) {
+        UNIT_ASSERT_VALUES_EQUAL(subgroup->GetCounter("PendingRequestsCount", false)->Val(), 0);
+        UNIT_ASSERT_VALUES_EQUAL(subgroup->GetCounter("FinishingRequestsCount", false)->Val(), 0);
+
+        const std::vector<std::pair<TString, bool>> tableQueries = {
+            {"TCleanupTablesQuery", false},
+            {"TRefreshPoolStateQuery", true},
+            {"TDelayRequestQuery", true},
+            {"TStartFirstDelayedRequestQuery", true},
+            {"TStartRequestQuery", false},
+            {"TCleanupRequestsQuery", true},
+        };
+        for (const auto& [operation, runExpected] : tableQueries) {
+            auto operationSubgroup = subgroup->GetSubgroup("operation", operation);
+
+            UNIT_ASSERT_VALUES_EQUAL_C(operationSubgroup->GetCounter("FinishError", true)->Val(), 0, TStringBuilder() << "Unexpected vaule for operation " << operation);
+            if (runExpected) {
+                UNIT_ASSERT_GE_C(operationSubgroup->GetCounter("FinishOk", true)->Val(), 1, TStringBuilder() << "Unexpected vaule for operation " << operation);
+            }
+        }
+    }
+
 private:
     const TYdbSetupSettings Settings_;
 
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
@@ -100,6 +100,7 @@ class IYdbSetup : public TThrRefBase {
     virtual void WaitPoolState(const TPoolStateDescription& state, const TString& poolId = "") const = 0;
     virtual void WaitPoolHandlersCount(i64 finalCount, std::optional<i64> initialCount = std::nullopt, TDuration timeout = FUTURE_WAIT_TIMEOUT) const = 0;
     virtual void StopWorkloadService(ui64 nodeIndex = 0) const = 0;
+    virtual void ValidateWorkloadServiceCounters(bool checkTableCounters = true, const TString& poolId = "") const = 0;
 
     virtual TTestActorRuntime* GetRuntime() const = 0;
     virtual const TYdbSetupSettings& GetSettings() const = 0;
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
@@ -5,10 +5,39 @@
 
 namespace NKikimr::NKqp {
 
+namespace {
+
 using namespace NWorkload;
 using namespace NYdb;
 
 
+void StartTestConcurrentQueryLimit(const ui64 activeCountLimit, const ui64 queueSize, const ui64 nodeCount = 1) {
+    auto ydb = TYdbSetupSettings()
+        .NodeCount(nodeCount)
+        .ConcurrentQueryLimit(activeCountLimit)
+        .QueueSize(queueSize)
+        .QueryCancelAfter(FUTURE_WAIT_TIMEOUT * queueSize)
+        .Create();
+
+    auto settings = TQueryRunnerSettings()
+        .InFlightCoordinatorActorId(ydb->CreateInFlightCoordinator(queueSize, activeCountLimit))
+        .HangUpDuringExecution(true);
+
+    // Initialize queue
+    std::vector<TQueryRunnerResultAsync> asyncResults;
+    for (size_t i = 0; i < queueSize; ++i) {
+        asyncResults.emplace_back(ydb->ExecuteQueryAsync(TSampleQueries::TSelect42::Query, settings.NodeIndex(i % nodeCount)));
+    }
+
+    for (const auto& asyncResult : asyncResults) {
+        TSampleQueries::TSelect42::CheckResult(asyncResult.GetResult());
+    }
+
+    ydb->ValidateWorkloadServiceCounters();
+}
+
+}  // anonymous namespace
+
 Y_UNIT_TEST_SUITE(KqpWorkloadService) {
     Y_UNIT_TEST(WorkloadServiceDisabledByFeatureFlag) {
         auto ydb = TYdbSetupSettings()
@@ -123,10 +152,7 @@ Y_UNIT_TEST_SUITE(KqpWorkloadService) {
         UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::CANCELLED, result.GetIssues().ToString());
 
         // Check that queue is free
-        auto firstRequest = ydb->ExecuteQueryAsync(TSampleQueries::TSelect42::Query);
-        auto secondRequest = ydb->ExecuteQueryAsync(TSampleQueries::TSelect42::Query);
-        TSampleQueries::TSelect42::CheckResult(firstRequest.GetResult());
-        TSampleQueries::TSelect42::CheckResult(secondRequest.GetResult());
+        UNIT_ASSERT_VALUES_EQUAL(ydb->GetPoolDescription().AmountRequests(), 0);
     }
 
     Y_UNIT_TEST(TestStartQueryAfterCancel) {
@@ -145,28 +171,12 @@ Y_UNIT_TEST_SUITE(KqpWorkloadService) {
         TSampleQueries::CheckCancelled(hangingRequest.GetResult());
     }
 
-    Y_UNIT_TEST(TestConcurrentQueryLimit) {
-        const ui64 activeCountLimit = 5;
-        const ui64 queueSize = 50;
-        auto ydb = TYdbSetupSettings()
-            .ConcurrentQueryLimit(activeCountLimit)
-            .QueueSize(queueSize)
-            .QueryCancelAfter(FUTURE_WAIT_TIMEOUT * queueSize)
-            .Create();
-
-        auto settings = TQueryRunnerSettings()
-            .InFlightCoordinatorActorId(ydb->CreateInFlightCoordinator(queueSize, activeCountLimit))
-            .HangUpDuringExecution(true);
-
-        // Initialize queue
-        std::vector<TQueryRunnerResultAsync> asyncResults;
-        for (size_t i = 0; i < queueSize; ++i) {
-            asyncResults.emplace_back(ydb->ExecuteQueryAsync(TSampleQueries::TSelect42::Query, settings));
-        }
+    Y_UNIT_TEST(TestLargeConcurrentQueryLimit) {
+        StartTestConcurrentQueryLimit(5, 100);
+    }
 
-        for (const auto& asyncResult : asyncResults) {
-            TSampleQueries::TSelect42::CheckResult(asyncResult.GetResult());
-        }
+    Y_UNIT_TEST(TestLessConcurrentQueryLimit) {
+        StartTestConcurrentQueryLimit(1, 100);
     }
 
     Y_UNIT_TEST(TestZeroConcurrentQueryLimit) {
@@ -255,30 +265,12 @@ Y_UNIT_TEST_SUITE(KqpWorkloadServiceDistributed) {
         TSampleQueries::TSelect42::CheckResult(request.GetResult(TDuration::Seconds(50)));
     }
 
-    Y_UNIT_TEST(TestDistributedConcurrentQueryLimit) {
-        const ui64 nodeCount = 3;
-        const ui64 activeCountLimit = 5;
-        const ui64 queueSize = 50;
-        auto ydb = TYdbSetupSettings()
-            .NodeCount(nodeCount)
-            .ConcurrentQueryLimit(activeCountLimit)
-            .QueueSize(queueSize)
-            .QueryCancelAfter(FUTURE_WAIT_TIMEOUT * queueSize)
-            .Create();
-
-        auto settings = TQueryRunnerSettings()
-            .InFlightCoordinatorActorId(ydb->CreateInFlightCoordinator(queueSize, activeCountLimit))
-            .HangUpDuringExecution(true);
-
-        // Initialize queue
-        std::vector<TQueryRunnerResultAsync> asyncResults;
-        for (size_t i = 0; i < queueSize; ++i) {
-            asyncResults.emplace_back(ydb->ExecuteQueryAsync(TSampleQueries::TSelect42::Query, settings.NodeIndex(i % nodeCount)));
-        }
+    Y_UNIT_TEST(TestDistributedLargeConcurrentQueryLimit) {
+        StartTestConcurrentQueryLimit(5, 100, 3);
+    }
 
-        for (const auto& asyncResult : asyncResults) {
-            TSampleQueries::TSelect42::CheckResult(asyncResult.GetResult());
-        }
+    Y_UNIT_TEST(TestDistributedLessConcurrentQueryLimit) {
+        StartTestConcurrentQueryLimit(1, 100, 5);
     }
 }
 
diff --git a/ydb/library/query_actor/query_actor.cpp b/ydb/library/query_actor/query_actor.cpp
diff --git a/ydb/library/query_actor/query_actor.h b/ydb/library/query_actor/query_actor.h
