Reconnect session has been supported (#9862)

Author: dorooleg

ydb/library/yql/providers/common/proto/gateways_config.proto

@@ -326,7 +326,8 @@ message TPqClusterConfig {
     optional bool AddBearerToToken = 11; // whether to use prefix "Bearer " in token
     optional string DatabaseId = 12;
     repeated TAttr Settings = 100;
-    optional bool SharedReading = 101; 
+    optional bool SharedReading = 101;
+    optional string ReconnectPeriod = 102; // disabled by default, example of a parameter: 5m
 }
 
 message TPqGatewayConfig {

ydb/library/yql/providers/pq/async_io/dq_pq_read_actor.cpp

@@ -74,6 +74,7 @@ struct TEvPrivate {
         EvBegin = EventSpaceBegin(NActors::TEvents::ES_PRIVATE),
 
         EvSourceDataReady = EvBegin,
+        EvReconnectSession,
 
         EvEnd
     };
@@ -83,6 +84,7 @@ struct TEvPrivate {
     // Events
 
     struct TEvSourceDataReady : public TEventLocal<TEvSourceDataReady, EvSourceDataReady> {};
+    struct TEvReconnectSession : public TEventLocal<TEvReconnectSession, EvReconnectSession> {};
 };
 
 } // namespace
@@ -98,6 +100,7 @@ class TDqPqReadActor : public NActors::TActor<TDqPqReadActor>, public NYql::NDq:
             InFlyAsyncInputData = task->GetCounter("InFlyAsyncInputData");
             InFlySubscribe = task->GetCounter("InFlySubscribe");
             AsyncInputDataRate = task->GetCounter("AsyncInputDataRate", true);
+            ReconnectRate = task->GetCounter("ReconnectRate", true);
         }
 
         ~TMetrics() {
@@ -110,6 +113,7 @@ class TDqPqReadActor : public NActors::TActor<TDqPqReadActor>, public NYql::NDq:
         ::NMonitoring::TDynamicCounters::TCounterPtr InFlyAsyncInputData;
         ::NMonitoring::TDynamicCounters::TCounterPtr InFlySubscribe;
         ::NMonitoring::TDynamicCounters::TCounterPtr AsyncInputDataRate;
+        ::NMonitoring::TDynamicCounters::TCounterPtr ReconnectRate;
     };
 
 public:
@@ -139,6 +143,7 @@ class TDqPqReadActor : public NActors::TActor<TDqPqReadActor>, public NYql::NDq:
         , CredentialsProviderFactory(std::move(credentialsProviderFactory))
         , PqGateway(pqGateway)
     {
+        Y_UNUSED(TDuration::TryParse(SourceParams.GetReconnectPeriod(), ReconnectPeriod));
         MetadataFields.reserve(SourceParams.MetadataFieldsSize());
         TPqMetaExtractor fieldsExtractor;
         for (const auto& fieldName : SourceParams.GetMetadataFields()) {
@@ -209,6 +214,7 @@ class TDqPqReadActor : public NActors::TActor<TDqPqReadActor>, public NYql::NDq:
 private:
     STRICT_STFUNC(StateFunc,
         hFunc(TEvPrivate::TEvSourceDataReady, Handle);
+        hFunc(TEvPrivate::TEvReconnectSession, Handle);
     )
 
     void Handle(TEvPrivate::TEvSourceDataReady::TPtr& ev) {
@@ -222,6 +228,18 @@ class TDqPqReadActor : public NActors::TActor<TDqPqReadActor>, public NYql::NDq:
         Send(ComputeActorId, new TEvNewAsyncInputDataArrived(InputIndex));
     }
 
+    void Handle(TEvPrivate::TEvReconnectSession::TPtr&) {
+        SRC_LOG_D("SessionId: " << GetSessionId() << ", Reconnect epoch: " << Metrics.ReconnectRate->Val());
+        Metrics.ReconnectRate->Inc();
+        if (ReadSession) {
+            ReadSession->Close(TDuration::Zero());
+            ReadSession.reset();
+            ReadyBuffer = std::queue<TReadyBatch>{}; // clear read buffer
+        }
+
+        Schedule(ReconnectPeriod, new TEvPrivate::TEvReconnectSession());
+    }
+
     // IActor & IDqComputeActorAsyncInput
     void PassAway() override { // Is called from Compute Actor
         std::queue<TReadyBatch> empty;
@@ -259,6 +277,12 @@ class TDqPqReadActor : public NActors::TActor<TDqPqReadActor>, public NYql::NDq:
         const auto now = TInstant::Now();
         MaybeScheduleNextIdleCheck(now);
 
+        if (!InflightReconnect && ReconnectPeriod != TDuration::Zero()) {
+            Metrics.ReconnectRate->Inc();
+            Schedule(ReconnectPeriod, new TEvPrivate::TEvSourceDataReady());
+            InflightReconnect = true;
+        }
+
         i64 usedSpace = 0;
         if (MaybeReturnReadyBatch(buffer, watermark, usedSpace)) {
             return usedSpace;
@@ -565,6 +589,8 @@ class TDqPqReadActor : public NActors::TActor<TDqPqReadActor>, public NYql::NDq:
     };
 
 private:
+    bool InflightReconnect = false;
+    TDuration ReconnectPeriod;
     TMetrics Metrics;
     const i64 BufferSize;
     const THolderFactory& HolderFactory;

ydb/library/yql/providers/pq/common/yql_names.h

@@ -15,5 +15,6 @@ constexpr TStringBuf WatermarksEnableSetting = "WatermarksEnable";
 constexpr TStringBuf WatermarksGranularityUsSetting = "WatermarksGranularityUs";
 constexpr TStringBuf WatermarksLateArrivalDelayUsSetting = "WatermarksLateArrivalDelayUs";
 constexpr TStringBuf WatermarksIdlePartitionsSetting = "WatermarksIdlePartitions";
+constexpr TStringBuf ReconnectPeriod = "ReconnectPeriod";
 
 } // namespace NYql

ydb/library/yql/providers/pq/proto/dq_io.proto

@@ -37,6 +37,7 @@ message TDqPqTopicSource {
     repeated string ColumnTypes = 13;
     string Predicate = 14;
     bool SharedReading = 15;
+    string ReconnectPeriod = 16; // disabled by default, example of a parameter: 5m
 }
 
 message TDqPqTopicSink {

ydb/library/yql/providers/pq/provider/yql_pq_dq_integration.cpp

@@ -226,6 +226,8 @@ class TPqDqIntegration: public TDqIntegrationBase {
                         srcDesc.SetEndpoint(TString(Value(setting)));
                     } else if (name == SharedReading) {
                         sharedReading = FromString<bool>(Value(setting));
+                    } else if (name == ReconnectPeriod) {
+                        srcDesc.SetReconnectPeriod(TString(Value(setting)));
                     } else if (name == Format) {
                         format = TString(Value(setting));
                     } else if (name == UseSslSetting) {
@@ -348,6 +350,7 @@ class TPqDqIntegration: public TDqIntegrationBase {
 
         Add(props, EndpointSetting, clusterConfiguration->Endpoint, pos, ctx);
         Add(props, SharedReading, ToString(clusterConfiguration->SharedReading), pos, ctx);
+        Add(props, ReconnectPeriod, ToString(clusterConfiguration->ReconnectPeriod), pos, ctx);
         Add(props, Format, format, pos, ctx);
 
 

ydb/library/yql/providers/pq/provider/yql_pq_settings.cpp

@@ -43,6 +43,7 @@ void TPqConfiguration::Init(
         clusterSettings.UseSsl = cluster.GetUseSsl();
         clusterSettings.AddBearerToToken = cluster.GetAddBearerToToken();
         clusterSettings.SharedReading = cluster.GetSharedReading();
+        clusterSettings.ReconnectPeriod = cluster.GetReconnectPeriod();
 
         const TString authToken = typeCtx->Credentials->FindCredentialContent("cluster:default_" + clusterSettings.ClusterName, "default_pq", cluster.GetToken());
         clusterSettings.AuthToken = authToken;

ydb/library/yql/providers/pq/provider/yql_pq_settings.h

@@ -30,6 +30,7 @@ struct TPqClusterConfigurationSettings {
     TString AuthToken;
     bool AddBearerToToken = false;
     bool SharedReading = false;
+    TString ReconnectPeriod;
 };
 
 struct TPqConfiguration : public TPqSettings, public NCommon::TSettingDispatcher {