merge to stable-25-3 YQ-4312 streaming in ydb fixes and features (#25898)

Author: GrigoriyPA

ydb/core/fq/libs/row_dispatcher/format_handler/format_handler.cpp

@@ -437,7 +437,7 @@ class TTopicFormatHandler : public NActors::TActor<TTopicFormatHandler>, public
 
         if (const auto clientOffset = client->GetNextMessageOffset()) {
             if (Parser && CurrentOffset && *CurrentOffset > *clientOffset) {
-                LOG_ROW_DISPATCHER_DEBUG("Parser was flushed due to new historical offset " << *clientOffset << "(previous parser offset: " << *CurrentOffset << ")");
+                LOG_ROW_DISPATCHER_DEBUG("Parser was flushed due to new historical offset " << *clientOffset << " (previous parser offset: " << *CurrentOffset << ")");
                 Parser->Refresh(true);
             }
         }

ydb/core/fq/libs/row_dispatcher/leader_election.cpp

@@ -9,10 +9,13 @@
 #include <ydb/library/actors/core/actor_bootstrapped.h>
 #include <ydb/library/actors/core/hfunc.h>
 #include <ydb/library/actors/protos/actors.pb.h>
+#include <ydb/library/logger/actor.h>
 
 #include <ydb/core/base/path.h>
 #include <ydb/core/protos/config.pb.h>
 
+#include <memory>
+
 namespace NFq {
 
 using namespace NActors;
@@ -83,7 +86,11 @@ struct TLeaderElectionMetrics {
     ::NMonitoring::TDynamicCounters::TCounterPtr LeaderChanged;
 };
 
-class TLeaderElection: public TActorBootstrapped<TLeaderElection> {
+struct TActorSystemPtrMixin {
+    NKikimr::TDeferredActorLogBackend::TSharedAtomicActorSystemPtr ActorSystemPtr = std::make_shared<NKikimr::TDeferredActorLogBackend::TAtomicActorSystemPtr>(nullptr);
+};
+
+class TLeaderElection: public TActorBootstrapped<TLeaderElection>, public TActorSystemPtrMixin {
 
     enum class EState {
         Init,
@@ -93,8 +100,8 @@ class TLeaderElection: public TActorBootstrapped<TLeaderElection> {
         Started
     };
     NKikimrConfig::TSharedReadingConfig::TCoordinatorConfig Config;
-    const NKikimr::TYdbCredentialsProviderFactory& CredentialsProviderFactory;
-    NYdb::TDriver Driver;
+    NKikimr::TYdbCredentialsProviderFactory CredentialsProviderFactory;
+    std::unique_ptr<NYdb::TDriver> Driver;
     TYdbConnectionPtr YdbConnection;
     TString TablePathPrefix;
     const TString TenantId;
@@ -165,20 +172,19 @@ class TLeaderElection: public TActorBootstrapped<TLeaderElection> {
     void ProcessState();
     void ResetState();
     void SetTimeout();
+    NYdb::TDriverConfig GetYdbDriverConfig() const;
 };
 
 TLeaderElection::TLeaderElection(
     NActors::TActorId parentId,
     NActors::TActorId coordinatorId,
     const NKikimrConfig::TSharedReadingConfig::TCoordinatorConfig& config,
     const NKikimr::TYdbCredentialsProviderFactory& credentialsProviderFactory,
-    NYdb::TDriver driver,
+    NYdb::TDriver /*driver*/,
     const TString& tenant,
     const ::NMonitoring::TDynamicCounterPtr& counters)
     : Config(config)
     , CredentialsProviderFactory(credentialsProviderFactory)
-    , Driver(driver)
-    , YdbConnection(config.GetLocalMode() ? nullptr : NewYdbConnection(config.GetDatabase(), credentialsProviderFactory, Driver))
     , TablePathPrefix(JoinPath(config.GetDatabase().GetDatabase(), config.GetCoordinationNodePath()))
     , TenantId(JoinSeq(":", NKikimr::SplitPath(tenant)))
     , CoordinationNodePath(JoinPath(TablePathPrefix, TenantId))
@@ -218,13 +224,19 @@ TYdbSdkRetryPolicy::TPtr MakeSchemaRetryPolicy() {
 
 void TLeaderElection::Bootstrap() {
     Become(&TLeaderElection::StateFunc);
+    Y_ABORT_UNLESS(!ActorSystemPtr->load(std::memory_order_relaxed), "Double ActorSystemPtr init");
+    ActorSystemPtr->store(TActivationContext::ActorSystem(), std::memory_order_relaxed);
+
     LogPrefix = "TLeaderElection " + SelfId().ToString() + " ";
     LOG_ROW_DISPATCHER_DEBUG("Successfully bootstrapped, local coordinator id " << CoordinatorId.ToString()
          << ", tenant id " << TenantId << ", local mode " << Config.GetLocalMode() << ", coordination node path " << CoordinationNodePath);
     if (Config.GetLocalMode()) {
         TActivationContext::ActorSystem()->Send(ParentId, new NFq::TEvRowDispatcher::TEvCoordinatorChanged(CoordinatorId, 0));
         return;
     }
+
+    Driver = std::make_unique<NYdb::TDriver>(GetYdbDriverConfig());
+    YdbConnection = NewYdbConnection(Config.GetDatabase(), CredentialsProviderFactory, *Driver);
     ProcessState();
 }
 
@@ -469,6 +481,13 @@ void TLeaderElection::HandleException(const std::exception& e) {
     ResetState();
 }
 
+NYdb::TDriverConfig TLeaderElection::GetYdbDriverConfig() const {
+    NYdb::TDriverConfig cfg;
+    cfg.SetDiscoveryMode(NYdb::EDiscoveryMode::Async);
+    cfg.SetLog(std::make_unique<NKikimr::TDeferredActorLogBackend>(ActorSystemPtr, NKikimrServices::EServiceKikimr::YDB_SDK));
+    return cfg;
+}
+
 } // anonymous namespace
 
 ////////////////////////////////////////////////////////////////////////////////

ydb/core/fq/libs/row_dispatcher/topic_session.cpp

@@ -5,13 +5,11 @@
 #include <ydb/core/fq/libs/row_dispatcher/events/data_plane.h>
 #include <ydb/core/fq/libs/row_dispatcher/format_handler/format_handler.h>
 #include <ydb/core/protos/config.pb.h>
-
 #include <ydb/library/actors/core/actor_bootstrapped.h>
 #include <ydb/library/actors/core/hfunc.h>
 #include <ydb/library/yql/dq/actors/dq.h>
-
+#include <ydb/library/yql/providers/pq/common/pq_events_processor.h>
 #include <ydb/public/sdk/cpp/adapters/issue/issue.h>
-
 #include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/topic/client.h>
 
 #include <util/generic/queue.h>
@@ -62,6 +60,7 @@ struct TEvPrivate {
         EvCreateSession,
         EvSendStatistic,
         EvReconnectSession,
+        EvExecuteTopicEvent,
         EvEnd
     };
     static_assert(EvEnd < EventSpaceEnd(TEvents::ES_PRIVATE), "expect EvEnd < EventSpaceEnd(TEvents::ES_PRIVATE)");
@@ -71,13 +70,16 @@ struct TEvPrivate {
     struct TEvCreateSession : public TEventLocal<TEvCreateSession, EvCreateSession> {};
     struct TEvSendStatistic : public TEventLocal<TEvSendStatistic, EvSendStatistic> {};
     struct TEvReconnectSession : public TEventLocal<TEvReconnectSession, EvReconnectSession> {};
+    struct TEvExecuteTopicEvent : public NYql::TTopicEventBase<TEvExecuteTopicEvent, EvExecuteTopicEvent> {
+        using TTopicEventBase::TTopicEventBase;
+    };
 };
 
 constexpr ui64 SendStatisticPeriodSec = 2;
 constexpr ui64 MaxHandledEventsCount = 1000;
 constexpr ui64 MaxHandledEventsSize = 1000000;
 
-class TTopicSession : public TActorBootstrapped<TTopicSession> {
+class TTopicSession : public TActorBootstrapped<TTopicSession>, NYql::TTopicEventProcessor<TEvPrivate::TEvExecuteTopicEvent> {
 private:
     using TBase = TActorBootstrapped<TTopicSession>;
 
@@ -315,7 +317,7 @@ class TTopicSession : public TActorBootstrapped<TTopicSession> {
     [[maybe_unused]] static constexpr char ActorName[] = "FQ_ROW_DISPATCHER_SESSION";
 
 private:
-    NYdb::NTopic::TTopicClientSettings GetTopicClientSettings(bool useSsl) const;
+    NYdb::NTopic::TTopicClientSettings GetTopicClientSettings(bool useSsl);
     NYql::ITopicClient& GetTopicClient(bool useSsl);
     NYdb::NTopic::TReadSessionSettings GetReadSessionSettings(const TString& consumerName) const;
     void CreateTopicSession();
@@ -351,6 +353,7 @@ class TTopicSession : public TActorBootstrapped<TTopicSession> {
 private:
 
     STRICT_STFUNC_EXC(StateFunc,
+        hFunc(TEvPrivate::TEvExecuteTopicEvent, HandleTopicEvent);
         hFunc(NFq::TEvPrivate::TEvPqEventsReady, Handle);
         hFunc(NFq::TEvPrivate::TEvCreateSession, Handle);
         hFunc(NFq::TEvPrivate::TEvSendStatistic, Handle);
@@ -364,6 +367,7 @@ class TTopicSession : public TActorBootstrapped<TTopicSession> {
 
     STRICT_STFUNC_EXC(ErrorState,
         cFunc(TEvents::TEvPoisonPill::EventType, PassAway);
+        hFunc(TEvPrivate::TEvExecuteTopicEvent, HandleTopicEvent);
         IgnoreFunc(NFq::TEvPrivate::TEvPqEventsReady);
         IgnoreFunc(NFq::TEvPrivate::TEvCreateSession);
         IgnoreFunc(TEvRowDispatcher::TEvGetNextBatch);
@@ -427,6 +431,7 @@ void TTopicSession::PassAway() {
 
 void TTopicSession::SubscribeOnNextEvent() {
     if (!ReadSession || IsWaitingEvents) {
+        LOG_ROW_DISPATCHER_TRACE("Skip SubscribeOnNextEvent, has ReadSession: " << (ReadSession ? "true" : "false") << ", IsWaitingEvents: " << IsWaitingEvents);
         return;
     }
 
@@ -445,12 +450,16 @@ void TTopicSession::SubscribeOnNextEvent() {
     });
 }
 
-NYdb::NTopic::TTopicClientSettings TTopicSession::GetTopicClientSettings(bool useSsl) const {
-    return PqGateway->GetTopicClientSettings()
-        .Database(Database)
+NYdb::NTopic::TTopicClientSettings TTopicSession::GetTopicClientSettings(bool useSsl) {
+    auto opts = PqGateway->GetTopicClientSettings();
+    SetupTopicClientSettings(ActorContext().ActorSystem(), SelfId(), opts);
+
+    opts.Database(Database)
         .DiscoveryEndpoint(Endpoint)
         .SslCredentials(NYdb::TSslCredentials(useSsl))
         .CredentialsProviderFactory(CredentialsProviderFactory);
+
+    return opts;
 }
 
 NYql::ITopicClient& TTopicSession::GetTopicClient(bool useSsl) {
@@ -982,7 +991,7 @@ void TTopicSession::RefreshParsers() {
     }
 }
 
-}  // anonymous namespace
+} // anonymous namespace
 
 ////////////////////////////////////////////////////////////////////////////////
 
@@ -1005,4 +1014,4 @@ std::unique_ptr<IActor> NewTopicSession(
     return std::unique_ptr<IActor>(new TTopicSession(readGroup, topicPath, endpoint, database, config, functionRegistry, rowDispatcherActorId, compileServiceActorId, partitionId, std::move(driver), credentialsProviderFactory, counters, countersRoot, pqGateway, maxBufferSize));
 }
 
-}  // namespace NFq
+} // namespace NFq

ydb/core/fq/libs/row_dispatcher/ut/topic_session_ut.cpp

@@ -30,7 +30,7 @@ constexpr ui64 TimeoutBeforeStartSessionSec = 3;
 constexpr ui64 GrabTimeoutSec = 4 * TimeoutBeforeStartSessionSec;
 static_assert(GrabTimeoutSec <= WAIT_TIMEOUT.Seconds());
 
-template<bool MockTopicSession>
+template <bool MockTopicSession>
 class TFixture : public NTests::TBaseFixture {
 public:
     using TBase = NTests::TBaseFixture;
@@ -69,7 +69,7 @@ class TFixture : public NTests::TBaseFixture {
         CompileNotifier = Runtime.AllocateEdgeActor();
         const auto compileServiceActorId = Runtime.Register(CreatePurecalcCompileServiceMock(CompileNotifier));
 
-        if (MockTopicSession) {
+        if constexpr (MockTopicSession) {
             PqGatewayNotifier = Runtime.AllocateEdgeActor();
             MockPqGateway = CreateMockPqGateway({.Runtime = &Runtime, .Notifier = PqGatewayNotifier});
         }
@@ -115,9 +115,10 @@ class TFixture : public NTests::TBaseFixture {
             UNIT_ASSERT_C(ping, "Compilation is not performed for predicate: " << predicate);
         }
 
-        if (MockTopicSession) {
+        if constexpr (MockTopicSession) {
             Runtime.GrabEdgeEvent<TEvMockPqEvents::TEvCreateSession>(PqGatewayNotifier, TDuration::Seconds(GrabTimeoutSec));
-            MockPqGateway->AddEvent(TopicPath, NYdb::NTopic::TReadSessionEvent::TStartPartitionSessionEvent(nullptr, 0, 0), 0);
+            MockReadSession = MockPqGateway->ExtractReadSession(TopicPath);
+            MockReadSession->AddStartSessionEvent();
         }
     }
 
@@ -223,40 +224,18 @@ class TFixture : public NTests::TBaseFixture {
         return TRow().AddUint64(100 * index).AddString(TStringBuilder() << "value" << index);
     }
 
-    using TMessageInformation = NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessageInformation;
-    using TMessage = NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent::TMessage;
-
-    TMessageInformation MakeNextMessageInformation(size_t offset, size_t uncompressedSize) {
-        auto now = TInstant::Now();
-        TMessageInformation msgInfo(
-            offset,
-            "ProducerId",
-            0,
-            now,
-            now,
-            MakeIntrusive<NYdb::NTopic::TWriteSessionMeta>(),
-            MakeIntrusive<NYdb::NTopic::TMessageMeta>(),
-            uncompressedSize,
-            "messageGroupId"
-        );
-        return msgInfo;
-    }
-
     void PQWrite(
         const std::vector<TString>& sequence,
         ui64 firstMessageOffset = 0) {
-        if (!MockTopicSession) {
+        if constexpr (!MockTopicSession) {
             NYql::NDq::PQWrite(sequence, TopicPath, GetDefaultPqEndpoint());
         } else {
             ui64 offset = firstMessageOffset;
-            TVector<TMessage> msgs;
-            size_t size = 0;
+            TVector<IMockPqReadSession::TMessage> msgs;
             for (const auto& s : sequence) {
-                TMessage msg(s, nullptr, MakeNextMessageInformation(offset++, s.size()), CreatePartitionSession());
-                msgs.emplace_back(msg);
-                size += s.size();
+                msgs.push_back({.Offset = offset++, .Data = s});
             }
-            MockPqGateway->AddEvent(TopicPath, NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent(msgs, {}, CreatePartitionSession()), size);
+            MockReadSession->AddDataReceivedEvent(msgs);
         }
     }
 
@@ -280,6 +259,7 @@ class TFixture : public NTests::TBaseFixture {
     ui32 PartitionId = 0;
     NKikimrConfig::TSharedReadingConfig Config;
     TIntrusivePtr<IMockPqGateway> MockPqGateway;
+    IMockPqReadSession::TPtr MockReadSession;
 
     const TString Json1 = "{\"dt\":100,\"value\":\"value1\"}";
     const TString Json2 = "{\"dt\":200,\"value\":\"value2\"}";
@@ -628,7 +608,8 @@ Y_UNIT_TEST_SUITE(TopicSessionTests) {
 
         StopSession(ReadActorId2, source);
         Runtime.GrabEdgeEvent<TEvMockPqEvents::TEvCreateSession>(PqGatewayNotifier, TDuration::Seconds(GrabTimeoutSec));
-        MockPqGateway->AddEvent(TopicPath, NYdb::NTopic::TReadSessionEvent::TStartPartitionSessionEvent(nullptr, 0, 0), 0);
+        MockReadSession = MockPqGateway->ExtractReadSession(TopicPath);
+        MockReadSession->AddStartSessionEvent();
 
         std::vector<TString> data3 = { Json4 };
         PQWrite(data3, 4);

ydb/core/fq/libs/row_dispatcher/ya.make

@@ -29,6 +29,7 @@ PEERDIR(
     ydb/library/yql/dq/actors/common
     ydb/library/yql/dq/actors/compute
     ydb/library/yql/dq/proto
+    ydb/library/yql/providers/pq/common
     ydb/library/yql/providers/pq/provider
 
     ydb/public/sdk/cpp/adapters/issue

ydb/core/kqp/federated_query/kqp_federated_query_helpers.cpp

@@ -122,6 +122,9 @@ namespace {
         return allTypes.contains(type);
     }
 
+    void IKqpFederatedQuerySetupFactory::Cleanup() {
+    }
+
     // TKqpFederatedQuerySetupFactoryDefault contains network clients and service actors necessary
     // for federated queries. HTTP Gateway (required by S3 provider) is run by default even without
     // explicit configuration. Token Accessor and Connector Client are run only if config is provided.
@@ -234,6 +237,11 @@ namespace {
         return result;
     }
 
+    void TKqpFederatedQuerySetupFactoryDefault::Cleanup() {
+        HttpGateway.reset();
+        PqGateway.Reset();
+    }
+
     IKqpFederatedQuerySetupFactory::TPtr MakeKqpFederatedQuerySetupFactory(
         NActors::TActorSystemSetup* setup,
         const NKikimr::TAppData* appData,

ydb/core/kqp/federated_query/kqp_federated_query_helpers.h

@@ -62,6 +62,7 @@ namespace NKikimr::NKqp {
 
     struct IKqpFederatedQuerySetupFactory {
         using TPtr = std::shared_ptr<IKqpFederatedQuerySetupFactory>;
+        virtual void Cleanup();
         virtual std::optional<TKqpFederatedQuerySetup> Make(NActors::TActorSystem* actorSystem) = 0;
         virtual ~IKqpFederatedQuerySetupFactory() = default;
     };
@@ -82,6 +83,8 @@ namespace NKikimr::NKqp {
 
         std::optional<TKqpFederatedQuerySetup> Make(NActors::TActorSystem* actorSystem) override;
 
+        void Cleanup() override;
+
     private:
         NYql::THttpGatewayConfig HttpGatewayConfig;
         NYql::IHTTPGateway::TPtr HttpGateway;
@@ -153,6 +156,11 @@ namespace NKikimr::NKqp {
                 DqTaskTransformFactory, PqGatewayConfig, PqGateway, ActorSystemPtr, Driver};
         }
 
+        void Cleanup() override {
+            HttpGateway.reset();
+            PqGateway.Reset();
+        }
+
     private:
         NYql::IHTTPGateway::TPtr HttpGateway;
         NYql::NConnector::IClient::TPtr ConnectorClient;

ydb/core/kqp/ut/common/kqp_ut_common.cpp

@@ -165,11 +165,13 @@ TKikimrRunner::TKikimrRunner(const TKikimrSettings& settings) {
     appConfig.MutableColumnShardConfig()->SetDisabledOnSchemeShard(false);
     appConfig.MutableTableServiceConfig()->SetEnableRowsDuplicationCheck(true);
     if (settings.EnableStorageProxy) {
-        appConfig.MutableQueryServiceConfig()->MutableCheckpointsConfig()->SetEnabled(true);
-        appConfig.MutableQueryServiceConfig()->MutableCheckpointsConfig()->SetCheckpointingPeriodMillis(200);
-        appConfig.MutableQueryServiceConfig()->MutableCheckpointsConfig()->SetMaxInflight(1);
-        appConfig.MutableQueryServiceConfig()->MutableCheckpointsConfig()->MutableExternalStorage()->SetEndpoint(GetEnv("YDB_ENDPOINT"));
-        appConfig.MutableQueryServiceConfig()->MutableCheckpointsConfig()->MutableExternalStorage()->SetDatabase(GetEnv("YDB_DATABASE"));
+        auto& checkpoints = *appConfig.MutableQueryServiceConfig()->MutableCheckpointsConfig();
+        checkpoints.SetEnabled(true);
+        checkpoints.SetCheckpointingPeriodMillis(settings.CheckpointPeriod.MilliSeconds());
+        checkpoints.SetMaxInflight(1);
+        checkpoints.MutableExternalStorage()->SetEndpoint(GetEnv("YDB_ENDPOINT"));
+        checkpoints.MutableExternalStorage()->SetDatabase(GetEnv("YDB_DATABASE"));
+        checkpoints.MutableCheckpointGarbageConfig()->SetEnabled(true);
     }
     if (!appConfig.GetQueryServiceConfig().HasAllExternalDataSourcesAreAvailable()) {
         appConfig.MutableQueryServiceConfig()->SetAllExternalDataSourcesAreAvailable(true);