PQ tablet add metadata in offset commit

ydb/core/kafka_proxy/actors/kafka_offset_commit_actor.cpp

@@ -112,6 +112,9 @@ void TKafkaOffsetCommitActor::Handle(NGRpcProxy::V1::TEvPQProxy::TEvAuthResultOk
             commit->SetClientId(Message->GroupId.value());
             commit->SetOffset(partitionRequest.CommittedOffset);
             commit->SetStrict(true);
+            if (partitionRequest.CommittedMetadata.has_value()) {
+                commit->SetCommittedMetadata(*partitionRequest.CommittedMetadata);
+            }
 
             PendingResponses++;
             KAFKA_LOG_D("Send commit request for group# " << Message->GroupId.value() <<

ydb/core/kafka_proxy/actors/kafka_offset_fetch_actor.cpp

@@ -93,10 +93,10 @@ class TTopicOffsetActor: public NKikimr::NGRpcProxy::V1::TPQInternalSchemaActor<
         Y_UNUSED(tabletInfo);
         Y_UNUSED(ctx);
         for (auto& partResult : ev->Get()->Record.GetPartResult()) {
-            std::unordered_map<TString, ui32> consumerToOffset;
+            std::unordered_map<TString, TEvKafka::PartitionConsumerOffset> consumerToOffset;
             for (auto& consumerResult : partResult.GetConsumerResult()) {
                 if (consumerResult.GetErrorCode() == NPersQueue::NErrorCode::OK) {
-                    consumerToOffset[consumerResult.GetConsumer()] = consumerResult.GetCommitedOffset();
+                    consumerToOffset.emplace(consumerResult.GetConsumer(), TEvKafka::PartitionConsumerOffset{static_cast<ui64>(consumerResult.GetCommitedOffset()), consumerResult.GetCommittedMetadata()});
                 }
             }
             (*PartitionIdToOffsets)[partResult.GetPartition()] = consumerToOffset;
@@ -144,6 +144,7 @@ class TTopicOffsetActor: public NKikimr::NGRpcProxy::V1::TPQInternalSchemaActor<
         response->TopicName = OriginalTopicName;
         response->Status = NONE_ERROR;
         response->PartitionIdToOffsets = PartitionIdToOffsets;
+
         Send(Requester, response.Release());
         Die(ctx);
     };
@@ -152,8 +153,7 @@ class TTopicOffsetActor: public NKikimr::NGRpcProxy::V1::TPQInternalSchemaActor<
         const TActorId Requester;
         const TString OriginalTopicName;
         const TString UserSID;
-        std::unordered_map<ui32, ui32> PartitionIdToOffset {};
-        std::shared_ptr<std::unordered_map<ui32, std::unordered_map<TString, ui32>>> PartitionIdToOffsets = std::make_shared<std::unordered_map<ui32, std::unordered_map<TString, ui32>>>();
+        std::shared_ptr<std::unordered_map<ui32, std::unordered_map<TString, TEvKafka::PartitionConsumerOffset>>> PartitionIdToOffsets = std::make_shared<std::unordered_map<ui32, std::unordered_map<TString, TEvKafka::PartitionConsumerOffset>>>();
 };
 
 NActors::IActor* CreateKafkaOffsetFetchActor(const TContext::TPtr context, const ui64 correlationId, const TMessagePtr<TOffsetFetchRequestData>& message) {
@@ -175,10 +175,15 @@ TOffsetFetchResponseData::TPtr TKafkaOffsetFetchActor::GetOffsetFetchResponse()
                     TOffsetFetchResponseData::TOffsetFetchResponseGroup::TOffsetFetchResponseTopics::TOffsetFetchResponsePartitions partition;
                     partition.PartitionIndex = requestPartition;
                     if (partitionsToOffsets.get() != nullptr
-                            && partitionsToOffsets->contains(requestPartition)
-                            && (*partitionsToOffsets)[requestPartition].contains(requestGroup.GroupId.value())) {
-                        partition.CommittedOffset = (*partitionsToOffsets)[requestPartition][requestGroup.GroupId.value()];
-                        partition.ErrorCode = NONE_ERROR;
+                            && partitionsToOffsets->contains(requestPartition)) {
+                        auto groupPartitionToOffset = (*partitionsToOffsets)[requestPartition].find(requestGroup.GroupId.value());
+                        if (groupPartitionToOffset != (*partitionsToOffsets)[requestPartition].end()) {
+                            partition.CommittedOffset = groupPartitionToOffset->second.Offset;
+                            partition.Metadata = groupPartitionToOffset->second.Metadata;
+                            partition.ErrorCode = NONE_ERROR;
+                        } else {
+                            partition.ErrorCode = RESOURCE_NOT_FOUND;
+                        }
                     } else {
                         partition.ErrorCode = RESOURCE_NOT_FOUND;
                     }
@@ -204,6 +209,7 @@ TOffsetFetchResponseData::TPtr TKafkaOffsetFetchActor::GetOffsetFetchResponse()
                 NKafka::TOffsetFetchResponseData::TOffsetFetchResponseTopic::TOffsetFetchResponsePartition partition;
                 partition.CommittedOffset = sourcePartition.CommittedOffset;
                 partition.PartitionIndex = sourcePartition.PartitionIndex;
+                partition.Metadata = sourcePartition.Metadata;
                 partition.ErrorCode = sourcePartition.ErrorCode;
                 topic.Partitions.push_back(partition);
             }

ydb/core/kafka_proxy/kafka_events.h

@@ -225,6 +225,11 @@ struct TEvTopicOffsetsResponse : public NActors::TEventLocal<TEvTopicOffsetsResp
     TVector<TPartitionOffsetsInfo> Partitions;
 };
 
+struct PartitionConsumerOffset {
+    ui64 Offset;
+    TString Metadata;
+};
+
 struct TEvCommitedOffsetsResponse : public NActors::TEventLocal<TEvCommitedOffsetsResponse, EvTopicOffsetsResponse>
                                   , public NKikimr::NGRpcProxy::V1::TLocalResponseBase
 {
@@ -233,7 +238,8 @@ struct TEvCommitedOffsetsResponse : public NActors::TEventLocal<TEvCommitedOffse
 
     TString TopicName;
     EKafkaErrors Status;
-    std::shared_ptr<std::unordered_map<ui32, std::unordered_map<TString, ui32>>> PartitionIdToOffsets;
+
+    std::shared_ptr<std::unordered_map<ui32, std::unordered_map<TString, PartitionConsumerOffset>>> PartitionIdToOffsets;
 };
 
 struct TEvTopicModificationResponse : public NActors::TEventLocal<TEvTopicModificationResponse, EvCreateTopicsResponse>
@@ -336,7 +342,7 @@ struct TProducerInstanceId {
     auto operator<=>(TProducerInstanceId const&) const = default;
 };
 
-/* 
+/*
 Event sent from TIintProducerActor to TKafkaTransactionRouter to notify that producer id will be obtained by client
  */
 struct TEvSaveTxnProducerRequest : public NActors::TEventLocal<TEvSaveTxnProducerRequest, EvSaveTxnProducerRequest> {
@@ -349,7 +355,7 @@ struct TEvSaveTxnProducerRequest : public NActors::TEventLocal<TEvSaveTxnProduce
     const TProducerInstanceId ProducerState;
 };
 
-/* 
+/*
 Event sent from TKafkaTransactionRouter to TIintProducerActor to notify that new transactional id was succesfully saved
 
 OK if this transactional producer was not found or older version was found

ydb/core/kafka_proxy/ut/kafka_test_client.cpp

@@ -82,7 +82,9 @@ TMessagePtr<TInitProducerIdResponseData> TKafkaTestClient::InitProducerId(const
     return WriteAndRead<TInitProducerIdResponseData>(header, request);
 }
 
-TMessagePtr<TOffsetCommitResponseData> TKafkaTestClient::OffsetCommit(TString groupId, std::unordered_map<TString, std::vector<std::pair<ui64,ui64>>> topicsToPartions) {
+
+
+TMessagePtr<TOffsetCommitResponseData> TKafkaTestClient::OffsetCommit(TString groupId, std::unordered_map<TString, std::vector<TConsumerOffset>> topicsToPartions) {
     Cerr << ">>>>> TOffsetCommitRequestData\n";
 
     TRequestHeaderData header = Header(NKafka::EApiKey::OFFSET_COMMIT, 1);
@@ -96,8 +98,9 @@ TMessagePtr<TOffsetCommitResponseData> TKafkaTestClient::OffsetCommit(TString gr
 
         for (auto partitionAndOffset : topicToPartitions.second) {
             NKafka::TOffsetCommitRequestData::TOffsetCommitRequestTopic::TOffsetCommitRequestPartition partition;
-            partition.PartitionIndex = partitionAndOffset.first;
-            partition.CommittedOffset = partitionAndOffset.second;
+            partition.PartitionIndex = partitionAndOffset.PartitionIndex;
+            partition.CommittedOffset = partitionAndOffset.Offset;
+            partition.CommittedMetadata = partitionAndOffset.Metadata;
             topic.Partitions.push_back(partition);
         }
         request.Topics.push_back(topic);
@@ -663,4 +666,4 @@ void TKafkaTestClient::FillTopicsFromJoinGroupMetadata(TKafkaBytes& metadata, TH
             topics.emplace(topic.value());
         }
     }
-}
+}

ydb/core/kafka_proxy/ut/kafka_test_client.h

@@ -37,6 +37,13 @@ struct TReadInfo {
     i32 GenerationId;
 };
 
+struct TConsumerOffset {
+    ui64 PartitionIndex;
+    ui64 Offset;
+    TString Metadata;
+};
+
+
 class TKafkaTestClient {
     public:
         TKafkaTestClient(ui16 port, const TString clientName = "TestClient");
@@ -61,7 +68,7 @@ class TKafkaTestClient {
 
         TMessagePtr<TInitProducerIdResponseData> InitProducerId(const TString& transactionalId = "");
 
-        TMessagePtr<TOffsetCommitResponseData> OffsetCommit(TString groupId, std::unordered_map<TString, std::vector<std::pair<ui64,ui64>>> topicsToPartions);
+        TMessagePtr<TOffsetCommitResponseData> OffsetCommit(TString groupId, std::unordered_map<TString, std::vector<TConsumerOffset>> topicsToPartions);
 
         TMessagePtr<TProduceResponseData> Produce(const TString& topicName, ui32 partition, const TKafkaRecordBatch& batch);
 
@@ -128,4 +135,4 @@ class TKafkaTestClient {
 
         ui32 Correlation;
         TString ClientName;
-    };
+    };

ydb/core/kafka_proxy/ut/ut_protocol.cpp

@@ -1,3 +1,5 @@
+
+
 #include <library/cpp/testing/unittest/registar.h>
 
 #include "kafka_test_client.h"
@@ -279,6 +281,8 @@ void CreateTopic(NYdb::NTopic::TTopicClient& pqClient, TString& topicName, ui32
 
 }
 
+
+
 Y_UNIT_TEST_SUITE(KafkaProtocol) {
     // this test imitates kafka producer behaviour:
     // 1. get api version,
@@ -996,6 +1000,8 @@ Y_UNIT_TEST_SUITE(KafkaProtocol) {
         TString headerKey = "header-key";
         TString headerValue = "header-value";
 
+        TString commitedMetaData = "additional-info";
+
         NYdb::NTopic::TTopicClient pqClient(*testServer.Driver);
         CreateTopic(pqClient, firstTopicName, minActivePartitions, {firstConsumerName, secondConsumerName});
         CreateTopic(pqClient, secondTopicName, minActivePartitions, {firstConsumerName, secondConsumerName});
@@ -1043,10 +1049,10 @@ Y_UNIT_TEST_SUITE(KafkaProtocol) {
 
         {
             // Check commit
-            std::unordered_map<TString, std::vector<std::pair<ui64,ui64>>> offsets;
-            std::vector<std::pair<ui64, ui64>> partitionsAndOffsets;
+            std::unordered_map<TString, std::vector<TConsumerOffset>> offsets;
+            std::vector<TConsumerOffset> partitionsAndOffsets;
             for (ui64 i = 0; i < minActivePartitions; ++i) {
-                partitionsAndOffsets.emplace_back(std::make_pair(i, recordsCount));
+                partitionsAndOffsets.emplace_back(TConsumerOffset{i, static_cast<ui64>(recordsCount), commitedMetaData});
             }
             offsets[firstTopicName] = partitionsAndOffsets;
             offsets[shortTopicName] = partitionsAndOffsets;
@@ -1074,12 +1080,16 @@ Y_UNIT_TEST_SUITE(KafkaProtocol) {
             topicsToPartions[firstTopicName] = std::vector<i32>{0, 1, 2 , 3 };
             auto msg = client.OffsetFetch(firstConsumerName, topicsToPartions);
             UNIT_ASSERT_VALUES_EQUAL(msg->Groups.size(), 1);
+            UNIT_ASSERT_VALUES_EQUAL(msg->Groups.size(), 1);
             UNIT_ASSERT_VALUES_EQUAL(msg->Groups[0].Topics.size(), 1);
             const auto& partitions = msg->Groups[0].Topics[0].Partitions;
             UNIT_ASSERT_VALUES_EQUAL(partitions.size(), 4);
             auto partition0 = std::find_if(partitions.begin(), partitions.end(), [](const auto& partition) { return partition.PartitionIndex == 0; });
             UNIT_ASSERT_VALUES_UNEQUAL(partition0, partitions.end());
             UNIT_ASSERT_VALUES_EQUAL(partition0->CommittedOffset, 5);
+            for (auto p = partitions.begin(); p != partitions.end(); p++) {
+                UNIT_ASSERT_VALUES_EQUAL(p->Metadata, "");
+            }
         }
 
         {
@@ -1097,10 +1107,10 @@ Y_UNIT_TEST_SUITE(KafkaProtocol) {
 
         {
             // Check commit with nonexistent topic
-            std::unordered_map<TString, std::vector<std::pair<ui64,ui64>>> offsets;
-            std::vector<std::pair<ui64, ui64>> partitionsAndOffsets;
+            std::unordered_map<TString, std::vector<TConsumerOffset>> offsets;
+            std::vector<TConsumerOffset> partitionsAndOffsets;
             for (ui64 i = 0; i < minActivePartitions; ++i) {
-                partitionsAndOffsets.emplace_back(std::make_pair(i, recordsCount));
+                partitionsAndOffsets.emplace_back(TConsumerOffset{i, static_cast<ui64>(recordsCount), commitedMetaData});
             }
             offsets[firstTopicName] = partitionsAndOffsets;
             offsets[notExistsTopicName] = partitionsAndOffsets;
@@ -1130,10 +1140,10 @@ Y_UNIT_TEST_SUITE(KafkaProtocol) {
 
         {
             // Check commit with nonexistent consumer
-            std::unordered_map<TString, std::vector<std::pair<ui64,ui64>>> offsets;
-            std::vector<std::pair<ui64, ui64>> partitionsAndOffsets;
+            std::unordered_map<TString, std::vector<TConsumerOffset>> offsets;
+            std::vector<TConsumerOffset> partitionsAndOffsets;
             for (ui64 i = 0; i < minActivePartitions; ++i) {
-                partitionsAndOffsets.emplace_back(std::make_pair(i, recordsCount));
+                partitionsAndOffsets.emplace_back(TConsumerOffset{i, static_cast<ui64>(recordsCount), commitedMetaData});
             }
             offsets[firstTopicName] = partitionsAndOffsets;
 

ydb/core/persqueue/events/internal.h

@@ -362,7 +362,7 @@ struct TEvPQ {
 
         TEvSetClientInfo(const ui64 cookie, const TString& clientId, const ui64 offset, const TString& sessionId, const ui64 partitionSessionId,
                             const ui32 generation, const ui32 step, const TActorId& pipeClient,
-                            ESetClientInfoType type = ESCI_OFFSET, ui64 readRuleGeneration = 0, bool strict = false)
+                            ESetClientInfoType type = ESCI_OFFSET, ui64 readRuleGeneration = 0, bool strict = false, const TString& сommittedMetadata = "")
         : Cookie(cookie)
         , ClientId(clientId)
         , Offset(offset)
@@ -374,6 +374,7 @@ struct TEvPQ {
         , ReadRuleGeneration(readRuleGeneration)
         , Strict(strict)
         , PipeClient(pipeClient)
+        , CommittedMetadata(сommittedMetadata)
         {
         }
 
@@ -388,6 +389,7 @@ struct TEvPQ {
         ui64 ReadRuleGeneration;
         bool Strict;
         TActorId PipeClient;
+        TString CommittedMetadata;
     };
 
 

ydb/core/persqueue/partition.cpp

@@ -2761,6 +2761,7 @@ void TPartition::OnProcessTxsAndUserActsWriteComplete(const TActorContext& ctx)
             userInfo.Generation = actual->Generation;
             userInfo.Step = actual->Step;
             userInfo.Offset = actual->Offset;
+            userInfo.CommittedMetadata = actual->CommittedMetadata;
             if (userInfo.Offset <= (i64)StartOffset) {
                 userInfo.AnyCommits = false;
             }
@@ -3244,6 +3245,7 @@ void TPartition::EmulatePostProcessUserAct(const TEvPQ::TEvSetClientInfo& act,
 {
     const TString& user = act.ClientId;
     ui64 offset = act.Offset;
+    const TString& committedMetadata = act.CommittedMetadata ? act.CommittedMetadata : userInfo.CommittedMetadata; // тут константная ссылка?
     const TString& session = act.SessionId;
     ui32 generation = act.Generation;
     ui32 step = act.Step;
@@ -3259,6 +3261,7 @@ void TPartition::EmulatePostProcessUserAct(const TEvPQ::TEvSetClientInfo& act,
         userInfo.Generation = userInfo.Step = 0;
         userInfo.Offset = 0;
         userInfo.AnyCommits = false;
+        userInfo.CommittedMetadata = "";
 
         PQ_LOG_D("Topic '" << TopicName() << "' partition " << Partition << " user " << user
                     << " drop done"
@@ -3275,13 +3278,16 @@ void TPartition::EmulatePostProcessUserAct(const TEvPQ::TEvSetClientInfo& act,
         userInfo.Generation = userInfo.Step = 0;
         userInfo.Offset = 0;
         userInfo.AnyCommits = false;
+        userInfo.CommittedMetadata = "";
 
         if (userInfo.Important) {
             userInfo.Offset = StartOffset;
         }
     } else {
         if (createSession || dropSession) {
             offset = userInfo.Offset;
+            // committedMetadata = userInfo.CommittedMetadata;
+
             auto *ui = UsersInfoStorage->GetIfExists(userInfo.User);
             auto ts = ui ? GetTime(*ui, userInfo.Offset) : std::make_pair<TInstant, TInstant>(TInstant::Zero(), TInstant::Zero());
 
@@ -3313,6 +3319,8 @@ void TPartition::EmulatePostProcessUserAct(const TEvPQ::TEvSetClientInfo& act,
         );
 
         userInfo.Offset = offset;
+        userInfo.CommittedMetadata = committedMetadata;
+        // userInfo.CommittedMetadata = metadata;
         if (userInfo.Offset <= (i64)StartOffset) {
             userInfo.AnyCommits = false;
         }
@@ -3419,7 +3427,7 @@ void TPartition::AddCmdWrite(NKikimrClient::TKeyValueRequest& request,
                              ui64 offset, ui32 gen, ui32 step, const TString& session,
                              ui64 readOffsetRewindSum,
                              ui64 readRuleGeneration,
-                             bool anyCommits)
+                             bool anyCommits, const TString& committedMetadata)
 {
     TBuffer idata;
     {
@@ -3431,6 +3439,7 @@ void TPartition::AddCmdWrite(NKikimrClient::TKeyValueRequest& request,
         userData.SetOffsetRewindSum(readOffsetRewindSum);
         userData.SetReadRuleGeneration(readRuleGeneration);
         userData.SetAnyCommits(anyCommits);
+        userData.SetCommittedMetadata(committedMetadata);
 
         TString out;
         Y_PROTOBUF_SUPPRESS_NODISCARD userData.SerializeToString(&out);
@@ -3492,7 +3501,7 @@ void TPartition::AddCmdWriteUserInfos(NKikimrClient::TKeyValueRequest& request)
                         userInfo->Session,
                         ui ? ui->ReadOffsetRewindSum : 0,
                         userInfo->ReadRuleGeneration,
-                        userInfo->AnyCommits);
+                        userInfo->AnyCommits, userInfo->CommittedMetadata);
         } else {
             AddCmdDeleteRange(request,
                               ikey, ikeyDeprecated);
@@ -3534,6 +3543,7 @@ TUserInfoBase& TPartition::GetOrCreatePendingUser(const TString& user,
             newPendingUserIt->second.Generation = userIt->Generation;
             newPendingUserIt->second.Step = userIt->Step;
             newPendingUserIt->second.Offset = userIt->Offset;
+            newPendingUserIt->second.CommittedMetadata = userIt->CommittedMetadata;
             newPendingUserIt->second.ReadRuleGeneration = userIt->ReadRuleGeneration;
             newPendingUserIt->second.Important = userIt->Important;
             newPendingUserIt->second.ReadFromTimestamp = userIt->ReadFromTimestamp;

ydb/core/persqueue/partition.h

@@ -363,7 +363,7 @@ class TPartition : public TActorBootstrapped<TPartition> {
                      ui64 offset, ui32 gen, ui32 step, const TString& session,
                      ui64 readOffsetRewindSum,
                      ui64 readRuleGeneration,
-                     bool anyCommits);
+                     bool anyCommits, const TString& committedMetadata);
     void AddCmdWriteTxMeta(NKikimrClient::TKeyValueRequest& request);
     void AddCmdWriteUserInfos(NKikimrClient::TKeyValueRequest& request);
     void AddCmdWriteConfig(NKikimrClient::TKeyValueRequest& request);

ydb/core/persqueue/partition_read.cpp

@@ -152,7 +152,7 @@ bool TPartition::ProcessHasDataRequest(const THasDataReq& request, const TActorC
 
             auto now = ctx.Now();
             auto& userInfo = UsersInfoStorage->GetOrCreate(request.ClientId, ctx);
-            userInfo.UpdateReadOffset((i64)EndOffset - 1, now, now, now, true);          
+            userInfo.UpdateReadOffset((i64)EndOffset - 1, now, now, now, true);
         }
     } else if (request.Offset < EndOffset) {
         sendResponse(GetSizeLag(request.Offset), false);
@@ -252,7 +252,7 @@ void TPartition::InitUserInfoForImportantClients(const TActorContext& ctx) {
         }
         if (!userInfo) {
             userInfo = &UsersInfoStorage->Create(
-                    ctx, consumer.GetName(), 0, true, "", 0, 0, 0, 0, 0, TInstant::Zero(), {}, false
+                    ctx, consumer.GetName(), 0, true, "", 0, 0, 0, 0, 0, TInstant::Zero(), {}, false, ""
             );
         }
         if (userInfo->Offset < (i64)StartOffset)