Kafka server balancing (#26299)

Author: nshestakov

ydb/core/kafka_proxy/actors/actors.h

@@ -1,14 +1,13 @@
 #pragma once
 
 #include <ydb/core/base/path.h>
+#include <ydb/core/kafka_proxy/kafka_messages.h>
 #include <ydb/core/persqueue/public/pq_rl_helpers.h>
 #include <ydb/core/protos/config.pb.h>
 #include <ydb/library/aclib/aclib.h>
 #include <ydb/public/api/protos/persqueue_error_codes_v1.pb.h>
 #include <ydb/public/api/protos/draft/persqueue_error_codes.pb.h> // strange
 
-#include <ydb/core/kafka_proxy/kafka_messages.h>
-
 namespace NKafka {
 
 static constexpr int ProxyNodeId = 1;
@@ -23,13 +22,29 @@ enum EAuthSteps {
     FAILED
 };
 
+enum class EBalancingMode {
+    Server,
+    Native,
+};
+
+struct TReadSession {
+    EBalancingMode BalancingMode = EBalancingMode::Server;
+    std::optional<EBalancingMode> PendingBalancingMode;
+    TActorId ProxyActorId;
+};
+
 struct TContext {
     using TPtr = std::shared_ptr<TContext>;
 
     TContext(const NKikimrConfig::TKafkaProxyConfig& config)
         : Config(config) {
     }
 
+    TContext(const TContext& other)
+        : Config(other.Config)
+    {
+    }
+
     const NKikimrConfig::TKafkaProxyConfig& Config;
 
     TActorId ConnectionId;
@@ -49,13 +64,13 @@ struct TContext {
     TString ClientDC;
     bool IsServerless = false;
     bool RequireAuthentication = false;
+    TReadSession ReadSession;
 
     NKikimr::NPQ::TRlContext RlContext;
 
-    bool Authenticated() {
 
+    bool Authenticated() {
         return !RequireAuthentication || AuthenticationStep == SUCCESS;
-
     }
 };
 
@@ -77,6 +92,10 @@ class TMessagePtr {
         return Ptr;
     }
 
+    T& operator*() const {
+        return *Ptr;
+    } 
+
     operator bool() const {
         return nullptr != Ptr;
     }
@@ -178,6 +197,7 @@ NActors::IActor* CreateKafkaMetadataActor(const TContext::TPtr context, const ui
                                           const TMessagePtr<TMetadataRequestData>& message,
                                           const TActorId& discoveryCacheActor);
 NActors::IActor* CreateKafkaProduceActor(const TContext::TPtr context);
+NActors::IActor* CreateKafkaReadSessionProxyActor(const TContext::TPtr context, ui64 cookie);
 NActors::IActor* CreateKafkaReadSessionActor(const TContext::TPtr context, ui64 cookie);
 NActors::IActor* CreateKafkaBalancerActor(const TContext::TPtr context, ui64 cookie);
 NActors::IActor* CreateKafkaSaslHandshakeActor(const TContext::TPtr context, const ui64 correlationId, const TMessagePtr<TSaslHandshakeRequestData>& message);

ydb/core/kafka_proxy/actors/kafka_balancer_actor.cpp

@@ -351,7 +351,7 @@ void TKafkaBalancerActor::JoinStepCheckGroupState(NKqp::TEvKqp::TEvQueryResponse
     if (SessionTimeoutMs > MAX_SESSION_TIMEOUT_MS || SessionTimeoutMs < MIN_SESSION_TIMEOUT_MS) {
         SendJoinGroupResponseFail(ctx, CorrelationId,
             EKafkaErrors::INVALID_SESSION_TIMEOUT,
-            "Invalid session timeout");
+            TStringBuilder() << "Invalid session timeout " << SessionTimeoutMs);
         return;
     }
 

ydb/core/kafka_proxy/actors/kafka_fetch_actor.cpp

@@ -30,10 +30,10 @@ void TKafkaFetchActor::Bootstrap(const NActors::TActorContext& ctx) {
 
 void TKafkaFetchActor::SendFetchRequests(const TActorContext& ctx) {
     Response->Responses.resize(FetchRequestData->Topics.size());
-    KAFKA_LOG_D(TStringBuilder() << "Fetch actor: New request. DatabasePath: " << Context->DatabasePath << " MaxWaitMs: " << FetchRequestData->MaxWaitMs << " MaxBytes: " << FetchRequestData->MaxBytes);
+    KAFKA_LOG_D(TStringBuilder() << "Fetch actor: New request. DatabasePath: " << Context->DatabasePath 
+        << " MaxWaitMs: " << FetchRequestData->MaxWaitMs << " MaxBytes: " << FetchRequestData->MaxBytes);
     for (size_t topicIndex = 0; topicIndex <  Response->Responses.size(); topicIndex++) {
-        TVector<NKikimr::NPQ::TPartitionFetchRequest> partPQRequests;
-        PrepareFetchRequestData(topicIndex, partPQRequests);
+        auto partPQRequests = PrepareFetchRequestData(topicIndex);
         auto ruPerRequest = topicIndex == 0 && Context->Config.GetMeteringV2Enabled();
         auto consumer = Context->GroupId.empty() ? NKikimr::NPQ::CLIENTID_WITHOUT_CONSUMER : Context->GroupId;
         NKikimr::NPQ::TFetchRequestSettings request {
@@ -55,11 +55,13 @@ void TKafkaFetchActor::SendFetchRequests(const TActorContext& ctx) {
     }
 }
 
-void TKafkaFetchActor::PrepareFetchRequestData(const size_t topicIndex, TVector<NKikimr::NPQ::TPartitionFetchRequest>& partPQRequests) {
+TVector<NKikimr::NPQ::TPartitionFetchRequest> TKafkaFetchActor::PrepareFetchRequestData(const size_t topicIndex) {
     auto& topicKafkaRequest = FetchRequestData->Topics[topicIndex];
     TFetchResponseData::TFetchableTopicResponse& topicKafkaResponse = Response->Responses[topicIndex];
+
     topicKafkaResponse.Topic = topicKafkaRequest.Topic;
 
+    TVector<NKikimr::NPQ::TPartitionFetchRequest> partPQRequests;
     partPQRequests.resize(topicKafkaRequest.Partitions.size());
     topicKafkaResponse.Partitions.resize(topicKafkaRequest.Partitions.size());
     for (size_t partIndex = 0; partIndex < topicKafkaRequest.Partitions.size(); partIndex++) {
@@ -74,6 +76,7 @@ void TKafkaFetchActor::PrepareFetchRequestData(const size_t topicIndex, TVector<
         partPQRequest.Offset = partKafkaRequest.FetchOffset;
         partPQRequest.MaxBytes = partKafkaRequest.PartitionMaxBytes;
     }
+    return partPQRequests;
 }
 
 void TKafkaFetchActor::Handle(NKikimr::TEvPQ::TEvFetchResponse::TPtr& ev, const TActorContext& ctx) {

ydb/core/kafka_proxy/actors/kafka_fetch_actor.h

@@ -37,7 +37,7 @@ class TKafkaFetchActor: public NActors::TActorBootstrapped<TKafkaFetchActor> {
     void Handle(NKikimr::TEvPQ::TEvFetchResponse::TPtr& ev, const TActorContext& ctx);
 
     void SendFetchRequests(const TActorContext& ctx);
-    void PrepareFetchRequestData(const size_t topicIndex, TVector<NKikimr::NPQ::TPartitionFetchRequest>& partPQRequests);
+    TVector<NKikimr::NPQ::TPartitionFetchRequest> PrepareFetchRequestData(const size_t topicIndex);
     void HandleErrorResponse(const NKikimr::TEvPQ::TEvFetchResponse::TPtr& ev, TFetchResponseData::TFetchableTopicResponse& topicResponse);
     void HandleSuccessResponse(const NKikimr::TEvPQ::TEvFetchResponse::TPtr& ev, TFetchResponseData::TFetchableTopicResponse& topicResponse, const TActorContext& ctx);
     void FillRecordsBatch(const NKikimrClient::TPersQueueFetchResponse_TPartResult& partPQResponse, TKafkaRecordBatch& recordsBatch, const TActorContext& ctx);

ydb/core/kafka_proxy/actors/kafka_metadata_actor.cpp

@@ -296,7 +296,7 @@ void TKafkaMetadataActor::SendCreateTopicsRequest(const TString& topicName, ui32
     topicToCreate.NumPartitions = Context->Config.GetTopicCreationDefaultPartitions();
     message->Topics.push_back(topicToCreate);
     TContext::TPtr ContextForTopicCreation;
-    ContextForTopicCreation = std::make_shared<TContext>(TContext(Context->Config));
+    ContextForTopicCreation = std::make_shared<TContext>(TContext(*Context));
     ContextForTopicCreation->ConnectionId = ctx.SelfID;
     ContextForTopicCreation->UserToken = Context->UserToken;
     ContextForTopicCreation->DatabasePath = Context->DatabasePath;

ydb/core/kafka_proxy/actors/kafka_offset_fetch_actor.cpp

@@ -522,7 +522,7 @@ void TKafkaOffsetFetchActor::CreateTopicIfNecessary(const TString& topicName,
     message->Topics.push_back(topicToCreate);
 
     TContext::TPtr ContextForTopicCreation;
-    ContextForTopicCreation = std::make_shared<TContext>(TContext(Context->Config));
+    ContextForTopicCreation = std::make_shared<TContext>(TContext(*Context));
     ContextForTopicCreation->ConnectionId = ctx.SelfID;
     ContextForTopicCreation->UserToken = Context->UserToken;
     ContextForTopicCreation->DatabasePath = Context->DatabasePath;

ydb/core/kafka_proxy/actors/kafka_read_session_actor.cpp

@@ -1,4 +1,5 @@
 #include "kafka_read_session_actor.h"
+#include "kafka_read_session_utils.h"
 
 namespace NKafka {
 static constexpr TDuration WAKEUP_INTERVAL = TDuration::Seconds(1);
@@ -100,7 +101,7 @@ void TKafkaReadSessionActor::HandleJoinGroup(TEvKafka::TEvJoinGroupRequest::TPtr
                 << "_" << "kafka";
 
             if (!supportedProtocolFound) {
-                SendJoinGroupResponseFail(ctx, ev->Get()->CorrelationId, INCONSISTENT_GROUP_PROTOCOL, TStringBuilder() << "unsupported assign protocol. Must be " << SUPPORTED_ASSIGN_STRATEGY);
+                SendJoinGroupResponseFail(ctx, ev->Get()->CorrelationId, INCONSISTENT_GROUP_PROTOCOL, TStringBuilder() << "unsupported assign protocol. Must be " << ASSIGN_STRATEGY_ROUNDROBIN);
                 CloseReadSession(ctx);
                 return;
             }
@@ -231,7 +232,7 @@ void TKafkaReadSessionActor::SendJoinGroupResponseOk(const TActorContext&, ui64
     TJoinGroupResponseData::TPtr response = std::make_shared<TJoinGroupResponseData>();
 
     response->ProtocolType = SUPPORTED_JOIN_GROUP_PROTOCOL;
-    response->ProtocolName = SUPPORTED_ASSIGN_STRATEGY;
+    response->ProtocolName = ASSIGN_STRATEGY_ROUNDROBIN;
     response->ErrorCode = EKafkaErrors::NONE_ERROR;
     response->GenerationId = GenerationId;
     response->MemberId = Session;
@@ -252,7 +253,7 @@ void TKafkaReadSessionActor::SendSyncGroupResponseOk(const TActorContext& ctx, u
     TSyncGroupResponseData::TPtr response = std::make_shared<TSyncGroupResponseData>();
 
     response->ProtocolType = SUPPORTED_JOIN_GROUP_PROTOCOL;
-    response->ProtocolName = SUPPORTED_ASSIGN_STRATEGY;
+    response->ProtocolName = ASSIGN_STRATEGY_ROUNDROBIN;
     response->ErrorCode = EKafkaErrors::NONE_ERROR;
 
     auto assignment = BuildAssignmentAndInformBalancerIfRelease(ctx);
@@ -323,17 +324,28 @@ bool TKafkaReadSessionActor::CheckHeartbeatIsExpired() {
     return now - LastHeartbeatTime > MaxHeartbeatTimeoutMs;
 }
 
-bool TKafkaReadSessionActor::TryFillTopicsToRead(const TMessagePtr<TJoinGroupRequestData> joinGroupRequestData, THashSet<TString>& topics) {
-    auto supportedProtocolFound = false;
-    for (auto protocol: joinGroupRequestData->Protocols) {
-        KAFKA_LOG_D("JOIN_GROUP assign protocol supported by client: " << protocol.Name);
-        if (protocol.Name == SUPPORTED_ASSIGN_STRATEGY) {
-            FillTopicsFromJoinGroupMetadata(protocol.Metadata, topics);
-            supportedProtocolFound = true;
-            break;
+bool TKafkaReadSessionActor::TryFillTopicsToRead(const TMessagePtr<TJoinGroupRequestData> request, THashSet<TString>& topics) {
+    auto validProtocol = request->ProtocolType == SUPPORTED_JOIN_GROUP_PROTOCOL
+        && AnyOf(request->Protocols, [](const TJoinGroupRequestData::TJoinGroupRequestProtocol& p) {
+            return p.Name == ASSIGN_STRATEGY_ROUNDROBIN || p.Name == ASSIGN_STRATEGY_SERVER;
+        });
+    if (!validProtocol) {
+        return false;
+    }
+
+    auto result = GetSubscriptions(*request);
+    for (auto topic: result->Topics) {
+        if (topic.has_value()) {
+            KAFKA_LOG_D("JOIN_GROUP requested topic to read: " << topic);
+
+            auto normalizedTopicName = NormalizePath(Context->DatabasePath, topic.value());
+            OriginalTopicNames[normalizedTopicName] = topic.value();
+            OriginalTopicNames[normalizedTopicName + "/streamImpl"] = topic.value();
+            topics.emplace(normalizedTopicName);
         }
     }
-    return supportedProtocolFound;
+
+    return true;
 }
 
 TConsumerProtocolAssignment TKafkaReadSessionActor::BuildAssignmentAndInformBalancerIfRelease(const TActorContext& ctx) {
@@ -376,26 +388,6 @@ TConsumerProtocolAssignment TKafkaReadSessionActor::BuildAssignmentAndInformBala
     return assignment;
 }
 
-void TKafkaReadSessionActor::FillTopicsFromJoinGroupMetadata(TKafkaBytes& metadata, THashSet<TString>& topics) {
-    TKafkaVersion version = *(TKafkaVersion*)(metadata.value().data() + sizeof(TKafkaVersion));
-
-    TBuffer buffer(metadata.value().data() + sizeof(TKafkaVersion), metadata.value().size_bytes() - sizeof(TKafkaVersion));
-    TKafkaReadable readable(buffer);
-
-    TConsumerProtocolSubscription result;
-    result.Read(readable, version);
-
-    for (auto topic: result.Topics) {
-        if (topic.has_value()) {
-            auto normalizedTopicName = NormalizePath(Context->DatabasePath, topic.value());
-            OriginalTopicNames[normalizedTopicName] = topic.value();
-            OriginalTopicNames[normalizedTopicName + "/streamImpl"] = topic.value();
-            topics.emplace(normalizedTopicName);
-            KAFKA_LOG_D("JOIN_GROUP requested topic to read: " << topic);
-        }
-    }
-}
-
 void TKafkaReadSessionActor::HandlePipeConnected(TEvTabletPipe::TEvClientConnected::TPtr& ev, const TActorContext&) {
     const auto* msg = ev->Get();
     if (msg->Status != NKikimrProto::OK) {

ydb/core/kafka_proxy/actors/kafka_read_session_actor.h

@@ -52,9 +52,6 @@ namespace NKafka {
      *           <----------------
      */
 
-static const TString SUPPORTED_ASSIGN_STRATEGY = "roundrobin";
-static const TString SUPPORTED_JOIN_GROUP_PROTOCOL = "consumer";
-
 class TKafkaReadSessionActor: public NActors::TActorBootstrapped<TKafkaReadSessionActor> {
 
 enum EReadSessionSteps {
@@ -157,7 +154,6 @@ struct TNextRequestError {
     bool CheckHeartbeatIsExpired();
     bool CheckTopicsListAreChanged(const TMessagePtr<TJoinGroupRequestData> joinGroupRequestData);
     bool TryFillTopicsToRead(const TMessagePtr<TJoinGroupRequestData> joinGroupRequestData, THashSet<TString>& topics);
-    void FillTopicsFromJoinGroupMetadata(TKafkaBytes& metadata, THashSet<TString>& topics);
 
 private:
     const TContext::TPtr Context;