YQ-3912 Add read group to row dispatcher (#12102)

kardymonds · web-flow · commit a15d8708109b · 2024-12-09T22:15:21.000+03:00
diff --git a/ydb/core/fq/libs/row_dispatcher/actors_factory.cpp b/ydb/core/fq/libs/row_dispatcher/actors_factory.cpp
@@ -9,6 +9,7 @@ struct TActorFactory : public IActorFactory {
     TActorFactory() {}
 
     NActors::TActorId RegisterTopicSession(
+        const TString& readGroup,
         const TString& topicPath,
         const TString& endpoint,
         const TString& database,
@@ -23,6 +24,7 @@ struct TActorFactory : public IActorFactory {
         ui64 maxBufferSize) const override {
 
         auto actorPtr = NFq::NewTopicSession(
+            readGroup,
             topicPath,
             endpoint,
             database,
diff --git a/ydb/core/fq/libs/row_dispatcher/actors_factory.h b/ydb/core/fq/libs/row_dispatcher/actors_factory.h
@@ -12,6 +12,7 @@ struct IActorFactory : public TThrRefBase {
     using TPtr = TIntrusivePtr<IActorFactory>;
 
     virtual NActors::TActorId RegisterTopicSession(
+        const TString& readGroup,
         const TString& topicPath,
         const TString& endpoint,
         const TString& database,
diff --git a/ydb/core/fq/libs/row_dispatcher/events/topic_session_stats.h b/ydb/core/fq/libs/row_dispatcher/events/topic_session_stats.h
@@ -51,6 +51,7 @@ struct TopicSessionCommonStatistic {
 };
 
 struct TopicSessionParams {
+    TString ReadGroup;
     TString Endpoint;
     TString Database;
     TString TopicPath;
diff --git a/ydb/core/fq/libs/row_dispatcher/row_dispatcher.cpp b/ydb/core/fq/libs/row_dispatcher/row_dispatcher.cpp
@@ -72,7 +72,25 @@ struct TEvPrivate {
     };
 };
 
-using TQueryStatKey = std::pair<TString, TString>;  // QueryId / Topic
+struct TQueryStatKey {
+    TString QueryId;
+    TString ReadGroup;
+
+    size_t Hash() const noexcept {
+        ui64 hash = std::hash<TString>()(QueryId);
+        hash = CombineHashes<ui64>(hash, std::hash<TString>()(ReadGroup));
+        return hash;
+    }
+    bool operator==(const TQueryStatKey& other) const {
+        return QueryId == other.QueryId && ReadGroup == other.ReadGroup;
+    }
+};
+
+struct TQueryStatKeyHash {
+    size_t operator()(const TQueryStatKey& k) const {
+        return k.Hash();
+    }
+};
 
 struct TAggQueryStat {
     NYql::TCounters::TEntry ReadBytes;
@@ -118,20 +136,22 @@ class TRowDispatcher : public TActorBootstrapped<TRowDispatcher> {
     };
 
     struct TopicSessionKey {
+        TString ReadGroup;
         TString Endpoint;
         TString Database;
         TString TopicPath;
         ui64 PartitionId;
 
         size_t Hash() const noexcept {
-            ui64 hash = std::hash<TString>()(Endpoint);
+            ui64 hash = std::hash<TString>()(ReadGroup);
+            hash = CombineHashes<ui64>(hash, std::hash<TString>()(Endpoint));
             hash = CombineHashes<ui64>(hash, std::hash<TString>()(Database));
             hash = CombineHashes<ui64>(hash, std::hash<TString>()(TopicPath));
             hash = CombineHashes<ui64>(hash, std::hash<ui64>()(PartitionId));
             return hash;
         }
         bool operator==(const TopicSessionKey& other) const {
-            return Endpoint == other.Endpoint && Database == other.Database
+            return ReadGroup == other.ReadGroup && Endpoint == other.Endpoint && Database == other.Database
                 && TopicPath == other.TopicPath && PartitionId == other.PartitionId;
         }
     };
@@ -243,7 +263,7 @@ class TRowDispatcher : public TActorBootstrapped<TRowDispatcher> {
 
     struct TAggregatedStats{
         NYql::TCounters::TEntry AllSessionsReadBytes;
-        TMap<TQueryStatKey, TMaybe<TAggQueryStat>> LastQueryStats;
+        THashMap<TQueryStatKey, TMaybe<TAggQueryStat>, TQueryStatKeyHash> LastQueryStats;
         TDuration LastUpdateMetricsPeriod;
     };
 
@@ -540,15 +560,14 @@ void TRowDispatcher::UpdateMetrics() {
     }
 
     for (auto& [key, sessionsInfo] : TopicSessions) {
-        const auto& topic = key.TopicPath;
         for (auto& [actorId, sessionInfo] : sessionsInfo.Sessions) {
             auto read = NYql::TCounters::TEntry(sessionInfo.Stat.ReadBytes);
             AggrStats.AllSessionsReadBytes.Add(read);
             sessionInfo.AggrReadBytes = read;
             sessionInfo.Stat.Clear();
 
             for (auto& [readActorId, consumer] : sessionInfo.Consumers) {
-                auto& stat = AggrStats.LastQueryStats[TQueryStatKey{consumer->QueryId, topic}];
+                auto& stat = AggrStats.LastQueryStats[TQueryStatKey{consumer->QueryId, key.ReadGroup}];
                 if (!stat) {
                     stat = TAggQueryStat();
                 }
@@ -557,22 +576,24 @@ void TRowDispatcher::UpdateMetrics() {
             }
         }
     }
-    for (auto it = AggrStats.LastQueryStats.begin(); it != AggrStats.LastQueryStats.end();) {
-        const auto& stats = it->second;
+    THashSet<TQueryStatKey, TQueryStatKeyHash> toDelete;
+    for (const auto& [key, stats] : AggrStats.LastQueryStats) {
         if (!stats) {
-            SetQueryMetrics(it->first, 0, 0, 0);
-            it = AggrStats.LastQueryStats.erase(it);
+            toDelete.insert(key);
             continue;
         }
-        SetQueryMetrics(it->first, stats->UnreadBytes.Max, stats->UnreadBytes.Avg, stats->ReadLagMessages.Max);
-        ++it;
+        SetQueryMetrics(key, stats->UnreadBytes.Max, stats->UnreadBytes.Avg, stats->ReadLagMessages.Max);
+    }
+    for (const auto& key : toDelete) {
+         SetQueryMetrics(key, 0, 0, 0);
+         AggrStats.LastQueryStats.erase(key);
     }
     PrintStateToLog();
 }
 
 void TRowDispatcher::SetQueryMetrics(const TQueryStatKey& queryKey, ui64 unreadBytesMax, ui64 unreadBytesAvg, i64 readLagMessagesMax) {
-    auto queryGroup = Metrics.Counters->GetSubgroup("queryId", queryKey.first);
-    auto topicGroup = queryGroup->GetSubgroup("topic", CleanupCounterValueString(queryKey.second));
+    auto queryGroup = Metrics.Counters->GetSubgroup("query_id", queryKey.QueryId);
+    auto topicGroup = queryGroup->GetSubgroup("read_group", CleanupCounterValueString(queryKey.ReadGroup));
     topicGroup->GetCounter("MaxUnreadBytes")->Set(unreadBytesMax);
     topicGroup->GetCounter("AvgUnreadBytes")->Set(unreadBytesAvg);
     topicGroup->GetCounter("MaxReadLag")->Set(readLagMessagesMax);
@@ -608,16 +629,15 @@ TString TRowDispatcher::GetInternalState() {
     printDataRate(AggrStats.AllSessionsReadBytes);
     str << "\n";
 
-    TMap<TQueryStatKey, TAggQueryStat> queryState;
-    TMap<TQueryStatKey, ui64> sessionCountByQuery;
+    THashMap<TQueryStatKey, TAggQueryStat, TQueryStatKeyHash> queryState;
+    THashMap<TQueryStatKey, ui64, TQueryStatKeyHash> sessionCountByQuery;
     ui64 unreadBytesSum = 0;
 
-    for (auto& [key, sessionsInfo] : TopicSessions) {
+    for (auto& [sessionKey, sessionsInfo] : TopicSessions) {
         for (auto& [actorId, sessionInfo] : sessionsInfo.Sessions) {
-            const auto& topic = key.TopicPath;
             unreadBytesSum += sessionInfo.Stat.UnreadBytes;
             for (auto& [readActorId, consumer] : sessionInfo.Consumers) {
-                auto key = TQueryStatKey{consumer->QueryId, topic};
+                auto key = TQueryStatKey{consumer->QueryId, sessionKey.ReadGroup};
                 ++sessionCountByQuery[key];
                 queryState[key].Add(consumer->Stat);
             }
@@ -630,11 +650,11 @@ TString TRowDispatcher::GetInternalState() {
 
     str << "Queries:\n";
     for (const auto& [queryStatKey, stat]: queryState) {
-        auto [queryId, topic] = queryStatKey;
+        auto [queryId, readGroup] = queryStatKey;
         const auto& aggStat = AggrStats.LastQueryStats[queryStatKey];
         auto sessionsBufferSumSize = sessionCountByQuery[queryStatKey] * MaxSessionBufferSizeBytes;
         auto used = sessionsBufferSumSize ? (stat.UnreadBytes.Sum * 100.0 / sessionsBufferSumSize) : 0.0;
-        str << "  " << queryId << " / " << topic << ": buffer used (all partitions) " << LeftPad(Prec(used, 4), 10) << "% (" << toHuman(stat.UnreadBytes.Sum) <<  ") unread max (one partition) " << toHuman(stat.UnreadBytes.Max) << " data rate";
+        str << "  " << queryId << " / " << readGroup << ": buffer used (all partitions) " << LeftPad(Prec(used, 4), 10) << "% (" << toHuman(stat.UnreadBytes.Sum) <<  ") unread max (one partition) " << toHuman(stat.UnreadBytes.Max) << " data rate";
         if (aggStat) {
             printDataRate(aggStat->ReadBytes);
         }
@@ -643,7 +663,7 @@ TString TRowDispatcher::GetInternalState() {
     }
     str << "TopicSessions:\n";
     for (auto& [key, sessionsInfo] : TopicSessions) {
-        str << "  " << key.TopicPath << " / " << key.PartitionId;
+        str << "  " << key.TopicPath << " / " << key.PartitionId  << " / " << key.ReadGroup;
         for (auto& [actorId, sessionInfo] : sessionsInfo.Sessions) {
             str << " / " << LeftPad(actorId, 32)
                 << " data rate " << toHumanDR(sessionInfo.AggrReadBytes.Sum) << " unread bytes " << toHuman(sessionInfo.Stat.UnreadBytes)
@@ -707,10 +727,10 @@ void TRowDispatcher::UpdateReadActorsInternalState() {
 }
 
 void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvStartSession::TPtr& ev) {
-    LOG_ROW_DISPATCHER_DEBUG("Received TEvStartSession from " << ev->Sender << ", topicPath " << ev->Get()->Record.GetSource().GetTopicPath() <<
+    LOG_ROW_DISPATCHER_DEBUG("Received TEvStartSession from " << ev->Sender << ", read group " << ev->Get()->Record.GetSource().GetReadGroup() << ", topicPath " << ev->Get()->Record.GetSource().GetTopicPath() <<
         " part id " << ev->Get()->Record.GetPartitionId() << " query id " << ev->Get()->Record.GetQueryId() << " cookie " << ev->Cookie);
-    auto queryGroup = Metrics.Counters->GetSubgroup("queryId", ev->Get()->Record.GetQueryId());
-    auto topicGroup = queryGroup->GetSubgroup("topic", CleanupCounterValueString(ev->Get()->Record.GetSource().GetTopicPath()));
+    auto queryGroup = Metrics.Counters->GetSubgroup("query_id", ev->Get()->Record.GetQueryId());
+    auto topicGroup = queryGroup->GetSubgroup("read_group", CleanupCounterValueString(ev->Get()->Record.GetSource().GetReadGroup()));
     topicGroup->GetCounter("StartSession", true)->Inc();
 
     NodesTracker.AddNode(ev->Sender.NodeId());
@@ -733,7 +753,7 @@ void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvStartSession::TPtr& ev) {
     }
     const auto& source = ev->Get()->Record.GetSource();
     TActorId sessionActorId;
-    TopicSessionKey topicKey{source.GetEndpoint(), source.GetDatabase(), source.GetTopicPath(), ev->Get()->Record.GetPartitionId()};
+    TopicSessionKey topicKey{source.GetReadGroup(), source.GetEndpoint(), source.GetDatabase(), source.GetTopicPath(), ev->Get()->Record.GetPartitionId()};
     TopicSessionInfo& topicSessionInfo = TopicSessions[topicKey];
     Y_ENSURE(topicSessionInfo.Sessions.size() <= 1);
 
@@ -745,8 +765,10 @@ void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvStartSession::TPtr& ev) {
     }
 
     if (topicSessionInfo.Sessions.empty()) {
-        LOG_ROW_DISPATCHER_DEBUG("Create new session, offset " << readOffset);
+        LOG_ROW_DISPATCHER_DEBUG("Create new session: read group " << source.GetReadGroup() << " topic " << source.GetTopicPath() 
+            << " part id " << ev->Get()->Record.GetPartitionId() << " offset " << readOffset);
         sessionActorId = ActorFactory->RegisterTopicSession(
+            source.GetReadGroup(),
             source.GetTopicPath(),
             source.GetEndpoint(),
             source.GetDatabase(),
@@ -850,6 +872,7 @@ void TRowDispatcher::DeleteConsumer(const ConsumerSessionKey& key) {
     Send(new IEventHandle(consumerIt->second->TopicSessionId, consumer->ReadActorId, event.release(), 0));
 
     TopicSessionKey topicKey{
+        consumer->SourceParams.GetReadGroup(),
         consumer->SourceParams.GetEndpoint(),
         consumer->SourceParams.GetDatabase(),
         consumer->SourceParams.GetTopicPath(),
@@ -879,7 +902,7 @@ void TRowDispatcher::Handle(const TEvPrivate::TEvTryConnect::TPtr& ev) {
 void TRowDispatcher::Handle(const NYql::NDq::TEvRetryQueuePrivate::TEvEvHeartbeat::TPtr& ev) {
     auto it = ConsumersByEventQueueId.find(ev->Get()->EventQueueId);
     if (it == ConsumersByEventQueueId.end()) {
-        LOG_ROW_DISPATCHER_WARN("No consumer with EventQueueId = " << ev->Get()->EventQueueId);
+        LOG_ROW_DISPATCHER_TRACE("No consumer with EventQueueId = " << ev->Get()->EventQueueId);
         return;
     }
     auto& sessionInfo = it->second;
@@ -984,7 +1007,7 @@ void TRowDispatcher::Handle(const NMon::TEvHttpInfo::TPtr& ev) {
 void TRowDispatcher::Handle(NFq::TEvRowDispatcher::TEvSessionStatistic::TPtr& ev) {
     LOG_ROW_DISPATCHER_TRACE("TEvSessionStatistic from " << ev->Sender);
     const auto& key = ev->Get()->Stat.SessionKey;
-    TopicSessionKey sessionKey{key.Endpoint, key.Database, key.TopicPath, key.PartitionId};
+    TopicSessionKey sessionKey{key.ReadGroup, key.Endpoint, key.Database, key.TopicPath, key.PartitionId};
 
     auto sessionsIt = TopicSessions.find(sessionKey);
     if (sessionsIt == TopicSessions.end()) {
diff --git a/ydb/core/fq/libs/row_dispatcher/topic_session.cpp b/ydb/core/fq/libs/row_dispatcher/topic_session.cpp
@@ -190,6 +190,7 @@ class TTopicSession : public TActorBootstrapped<TTopicSession> {
 
     bool InflightReconnect = false;
     TDuration ReconnectPeriod;
+    const TString ReadGroup;
     const TString TopicPath;
     const TString TopicPathPartition;
     const TString Endpoint;
@@ -225,6 +226,7 @@ class TTopicSession : public TActorBootstrapped<TTopicSession> {
 
 public:
     explicit TTopicSession(
+        const TString& readGroup,
         const TString& topicPath,
         const TString& endpoint,
         const TString& database,
@@ -322,6 +324,7 @@ class TTopicSession : public TActorBootstrapped<TTopicSession> {
 };
 
 TTopicSession::TTopicSession(
+    const TString& readGroup,
     const TString& topicPath,
     const TString& endpoint,
     const TString& database,
@@ -334,7 +337,8 @@ TTopicSession::TTopicSession(
     const ::NMonitoring::TDynamicCounterPtr& counters,
     const NYql::IPqGateway::TPtr& pqGateway,
     ui64 maxBufferSize)
-    : TopicPath(topicPath)
+    : ReadGroup(readGroup)
+    , TopicPath(topicPath)
     , TopicPathPartition(TStringBuilder() << topicPath << "/" << partitionId)
     , Endpoint(endpoint)
     , Database(database)
@@ -833,12 +837,16 @@ void TTopicSession::Handle(NFq::TEvRowDispatcher::TEvStartSession::TPtr& ev) {
     auto types = GetVector(ev->Get()->Record.GetSource().GetColumnTypes());
 
     try {
-        auto queryGroup = Counters->GetSubgroup("queryId", ev->Get()->Record.GetQueryId());
-        auto topicGroup = queryGroup->GetSubgroup("topic", CleanupCounterValueString(TopicPath));
+        if (Parser) {
+            // Parse remains data before adding new client
+            DoParsing(true);
+        }
+        auto queryGroup = Counters->GetSubgroup("query_id", ev->Get()->Record.GetQueryId());
+        auto readGroup = queryGroup->GetSubgroup("read_group", CleanupCounterValueString(ReadGroup));
         auto& clientInfo = Clients.emplace(
             std::piecewise_construct,
             std::forward_as_tuple(ev->Sender), 
-            std::forward_as_tuple(ev, topicGroup)).first->second;
+            std::forward_as_tuple(ev, readGroup)).first->second;
         UpdateFieldsIds(clientInfo);
 
         const auto& source = clientInfo.Settings.GetSource();
@@ -1044,7 +1052,7 @@ void TTopicSession::SendStatisticToRowDispatcher() {
     commonStatistic.LastReadedOffset = LastMessageOffset;
     SessionStats.Clear();
 
-    sessionStatistic.SessionKey = TopicSessionParams{Endpoint, Database, TopicPath, PartitionId};
+    sessionStatistic.SessionKey = TopicSessionParams{ReadGroup, Endpoint, Database, TopicPath, PartitionId};
     sessionStatistic.Clients.reserve(Clients.size());
     for (auto& [readActorId, info] : Clients) {
         TopicSessionClientStatistic clientStatistic;
@@ -1155,6 +1163,7 @@ TString TTopicSession::GetAnyQueryIdByFieldName(const TString& fieldName) {
 ////////////////////////////////////////////////////////////////////////////////
     
 std::unique_ptr<NActors::IActor> NewTopicSession(
+    const TString& readGroup,
     const TString& topicPath,
     const TString& endpoint,
     const TString& database,
@@ -1167,7 +1176,7 @@ std::unique_ptr<NActors::IActor> NewTopicSession(
     const ::NMonitoring::TDynamicCounterPtr& counters,
     const NYql::IPqGateway::TPtr& pqGateway,
     ui64 maxBufferSize) {
-    return std::unique_ptr<NActors::IActor>(new TTopicSession(topicPath, endpoint, database, config, rowDispatcherActorId, compileServiceActorId, partitionId, std::move(driver), credentialsProviderFactory, counters, pqGateway, maxBufferSize));
+    return std::unique_ptr<NActors::IActor>(new TTopicSession(readGroup, topicPath, endpoint, database, config, rowDispatcherActorId, compileServiceActorId, partitionId, std::move(driver), credentialsProviderFactory, counters, pqGateway, maxBufferSize));
 }
 
 } // namespace NFq
diff --git a/ydb/core/fq/libs/row_dispatcher/topic_session.h b/ydb/core/fq/libs/row_dispatcher/topic_session.h
@@ -16,6 +16,7 @@
 namespace NFq {
 
 std::unique_ptr<NActors::IActor> NewTopicSession(
+    const TString& readGroup,
     const TString& topicPath,
     const TString& endpoint,
     const TString& database,
diff --git a/ydb/core/fq/libs/row_dispatcher/ut/row_dispatcher_ut.cpp b/ydb/core/fq/libs/row_dispatcher/ut/row_dispatcher_ut.cpp
diff --git a/ydb/core/fq/libs/row_dispatcher/ut/topic_session_ut.cpp b/ydb/core/fq/libs/row_dispatcher/ut/topic_session_ut.cpp
diff --git a/ydb/tests/fq/yds/test_row_dispatcher.py b/ydb/tests/fq/yds/test_row_dispatcher.py
