OltpSink: reattach shard (#14863)

Author: nikvas0

ydb/core/kqp/common/kqp_tx_manager.cpp

@@ -202,6 +202,29 @@ class TKqpTransactionManager : public IKqpTransactionManager {
         return locks;
     }
 
+    void Reattached(ui64 shardId) override {
+        auto& shardInfo = ShardsInfo.at(shardId);
+        shardInfo.Reattaching = false;
+    }
+
+    void SetRestarting(ui64 shardId) override {
+        auto& shardInfo = ShardsInfo.at(shardId);
+        shardInfo.Restarting = true;
+    }
+
+    bool ShouldReattach(ui64 shardId, TInstant now) override {
+        auto& shardInfo = ShardsInfo.at(shardId);
+        if (!std::exchange(shardInfo.Restarting, false) && !shardInfo.Reattaching) {
+            return false;
+        }
+        return ::NKikimr::NKqp::ShouldReattach(now, shardInfo.ReattachState.ReattachInfo);;
+    }
+
+    TReattachState& GetReattachState(ui64 shardId) override {
+        auto& shardInfo = ShardsInfo.at(shardId);
+        return shardInfo.ReattachState;
+    }
+
     bool IsTxPrepared() const override {
         for (const auto& [_, shardInfo] : ShardsInfo) {
             if (shardInfo.State != EShardState::PREPARED) {
@@ -454,6 +477,10 @@ class TKqpTransactionManager : public IKqpTransactionManager {
 
         bool IsOlap = false;
         THashSet<TStringBuf> Pathes;
+
+        bool Restarting = false;
+        bool Reattaching = false;
+        TReattachState ReattachState;
     };
 
     void MakeLocksIssue(const TShardInfo& shardInfo) {

ydb/core/kqp/common/kqp_tx_manager.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <ydb/core/kqp/common/kqp_yql.h>
+#include <ydb/core/kqp/common/simple/reattach.h>
 #include <ydb/core/kqp/gateway/kqp_gateway.h>
 #include <ydb/core/kqp/provider/yql_kikimr_provider.h>
 #include <ydb/core/util/ulid.h>
@@ -43,11 +44,21 @@ class IKqpTransactionManager {
     virtual bool AddLock(ui64 shardId, const NKikimrDataEvents::TLock& lock) = 0;
 
     virtual void BreakLock(ui64 shardId) = 0;
+    virtual TVector<NKikimrDataEvents::TLock> GetLocks() const = 0;
+    virtual TVector<NKikimrDataEvents::TLock> GetLocks(ui64 shardId) const = 0;
 
     virtual TTableInfo GetShardTableInfo(ui64 shardId) const = 0;
 
-    virtual TVector<NKikimrDataEvents::TLock> GetLocks() const = 0;
-    virtual TVector<NKikimrDataEvents::TLock> GetLocks(ui64 shardId) const = 0;
+    virtual bool ShouldReattach(ui64 shardId, TInstant now) = 0;
+    virtual void Reattached(ui64 shardId) = 0;
+    virtual void SetRestarting(ui64 shardId) = 0;
+
+    struct TReattachState {
+        TReattachInfo ReattachInfo;
+        ui64 Cookie = 0;
+    };
+
+    virtual TReattachState& GetReattachState(ui64 shardId) = 0;
 
     virtual EShardState GetState(ui64 shardId) const = 0;
     virtual void SetError(ui64 shardId) = 0;
@@ -57,11 +68,7 @@ class IKqpTransactionManager {
 
     virtual void SetTopicOperations(NTopic::TTopicOperations&& topicOperations) = 0;
     virtual const NTopic::TTopicOperations& GetTopicOperations() const = 0;
-
-    virtual void SetAllowVolatile(bool allowVolatile) = 0;
-
     virtual void BuildTopicTxs(NTopic::TTopicOperationTransactions& txs) = 0;
-
     virtual bool HasTopics() const = 0;
 
     virtual bool IsTxPrepared() const = 0;
@@ -74,6 +81,7 @@ class IKqpTransactionManager {
     virtual bool IsEmpty() const = 0;
     virtual bool HasLocks() const = 0;
 
+    virtual void SetAllowVolatile(bool allowVolatile) = 0;
     virtual bool IsVolatile() const = 0;
 
     virtual bool HasSnapshot() const = 0;

ydb/core/kqp/common/simple/reattach.cpp

@@ -0,0 +1,44 @@
+#include "reattach.h"
+
+#include <ydb/core/base/appdata_fwd.h>
+#include <library/cpp/random_provider/random_provider.h>
+
+namespace NKikimr::NKqp {
+
+namespace {
+    static constexpr TDuration MinReattachDelay = TDuration::MilliSeconds(10);
+    static constexpr TDuration MaxReattachDelay = TDuration::MilliSeconds(100);
+    static constexpr TDuration MaxReattachDuration = TDuration::Seconds(4);
+}
+
+bool ShouldReattach(TInstant now, TReattachInfo& reattachInfo) {
+    if (!reattachInfo.Reattaching) {
+        reattachInfo.Deadline = now + MaxReattachDuration;
+        reattachInfo.Delay = TDuration::Zero();
+        reattachInfo.Reattaching = true;
+        return true;
+    }
+
+    TDuration left = reattachInfo.Deadline - now;
+    if (!left) {
+        reattachInfo.Reattaching = false;
+        return false;
+    }
+
+    reattachInfo.Delay *= 2.0;
+    if (reattachInfo.Delay < MinReattachDelay) {
+        reattachInfo.Delay = MinReattachDelay;
+    } else if (reattachInfo.Delay > MaxReattachDelay) {
+        reattachInfo.Delay = MaxReattachDelay;
+    }
+
+    // Add ±10% jitter
+    reattachInfo.Delay *= 0.9 + 0.2 * TAppData::RandomProvider->GenRandReal4();
+    if (reattachInfo.Delay > left) {
+        reattachInfo.Delay = left;
+    }
+
+    return true;
+}
+
+}

ydb/core/kqp/common/simple/reattach.h

@@ -0,0 +1,15 @@
+#pragma once
+
+#include <util/datetime/base.h>
+
+namespace NKikimr::NKqp {
+
+struct TReattachInfo {
+    TDuration Delay;
+    TInstant Deadline;
+    bool Reattaching = false;
+};
+
+bool ShouldReattach(TInstant now, TReattachInfo& reattachInfo);
+
+}

ydb/core/kqp/common/simple/ya.make

@@ -2,10 +2,11 @@ LIBRARY()
 
 SRCS(
     helpers.cpp
+    kqp_event_ids.cpp
     query_id.cpp
-    settings.cpp
+    reattach.cpp
     services.cpp
-    kqp_event_ids.cpp
+    settings.cpp
     temp_tables.cpp
 )
 

ydb/core/kqp/executer_actor/kqp_data_executer.cpp

@@ -11,8 +11,9 @@
 #include <ydb/core/client/minikql_compile/db_key_resolver.h>
 #include <ydb/core/kqp/common/buffer/events.h>
 #include <ydb/core/kqp/common/kqp_data_integrity_trails.h>
-#include <ydb/core/kqp/common/kqp_yql.h>
 #include <ydb/core/kqp/common/kqp_tx_manager.h>
+#include <ydb/core/kqp/common/kqp_yql.h>
+#include <ydb/core/kqp/common/simple/reattach.h>
 #include <ydb/library/wilson_ids/wilson.h>
 #include <ydb/core/kqp/compute_actor/kqp_compute_actor.h>
 #include <ydb/core/kqp/common/kqp_tx.h>
@@ -41,55 +42,24 @@ using namespace NLongTxService;
 
 namespace {
 
-static constexpr TDuration MinReattachDelay = TDuration::MilliSeconds(10);
-static constexpr TDuration MaxReattachDelay = TDuration::MilliSeconds(100);
-static constexpr TDuration MaxReattachDuration = TDuration::Seconds(4);
 static constexpr ui32 ReplySizeLimit = 48 * 1024 * 1024; // 48 MB
 
 class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Data> {
     using TBase = TKqpExecuterBase<TKqpDataExecuter, EExecType::Data>;
     using TKqpSnapshot = IKqpGateway::TKqpSnapshot;
 
     struct TReattachState {
-        TDuration Delay;
-        TInstant Deadline;
+        TReattachInfo ReattachInfo;
         ui64 Cookie = 0;
-        bool Reattaching = false;
 
         bool ShouldReattach(TInstant now) {
             ++Cookie; // invalidate any previous cookie
 
-            if (!Reattaching) {
-                Deadline = now + MaxReattachDuration;
-                Delay = TDuration::Zero();
-                Reattaching = true;
-                return true;
-            }
-
-            TDuration left = Deadline - now;
-            if (!left) {
-                Reattaching = false;
-                return false;
-            }
-
-            Delay *= 2.0;
-            if (Delay < MinReattachDelay) {
-                Delay = MinReattachDelay;
-            } else if (Delay > MaxReattachDelay) {
-                Delay = MaxReattachDelay;
-            }
-
-            // Add ±10% jitter
-            Delay *= 0.9 + 0.2 * TAppData::RandomProvider->GenRandReal4();
-            if (Delay > left) {
-                Delay = left;
-            }
-
-            return true;
+            return ::NKikimr::NKqp::ShouldReattach(now, ReattachInfo);
         }
 
         void Reattached() {
-            Reattaching = false;
+            ReattachInfo.Reattaching = false;
         }
     };
 
@@ -729,13 +699,13 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
             case TShardState::EState::Prepared: {
                 // Disconnected while waiting for other shards to prepare
 
-                if ((wasRestarting || shardState->ReattachState.Reattaching) &&
+                if ((wasRestarting || shardState->ReattachState.ReattachInfo.Reattaching) &&
                     shardState->ReattachState.ShouldReattach(TlsActivationContext->Now()))
                 {
                     LOG_N("Shard " << msg->TabletId << " delivery problem (already prepared, reattaching in "
-                        << shardState->ReattachState.Delay << ")");
+                        << shardState->ReattachState.ReattachInfo.Delay << ")");
 
-                    Schedule(shardState->ReattachState.Delay, new TEvPrivate::TEvReattachToShard(msg->TabletId));
+                    Schedule(shardState->ReattachState.ReattachInfo.Delay, new TEvPrivate::TEvReattachToShard(msg->TabletId));
                     ++shardState->RestartCount;
                     return;
                 }
@@ -1542,13 +1512,13 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
                 // Proceed with query processing
                 [[fallthrough]];
             case TShardState::EState::Executing: {
-                if ((wasRestarting || shardState->ReattachState.Reattaching) &&
+                if ((wasRestarting || shardState->ReattachState.ReattachInfo.Reattaching) &&
                      shardState->ReattachState.ShouldReattach(TlsActivationContext->Now()))
                 {
                     LOG_N("Shard " << msg->TabletId << " lost pipe while waiting for reply (reattaching in "
-                        << shardState->ReattachState.Delay << ")");
+                        << shardState->ReattachState.ReattachInfo.Delay << ")");
 
-                    Schedule(shardState->ReattachState.Delay, new TEvPrivate::TEvReattachToShard(msg->TabletId));
+                    Schedule(shardState->ReattachState.ReattachInfo.Delay, new TEvPrivate::TEvReattachToShard(msg->TabletId));
                     ++shardState->RestartCount;
                     return;
                 }