Read session param for CommitOffset RPC to 24-4 (#16697)

niksaveliev · web-flow · commit c09bd6c8149a · 2025-04-10T21:51:00.000+05:00
diff --git a/ydb/core/persqueue/ut/ut_with_sdk/autoscaling_ut.cpp b/ydb/core/persqueue/ut/ut_with_sdk/autoscaling_ut.cpp
@@ -621,6 +621,265 @@ Y_UNIT_TEST_SUITE(TopicAutoscaling) {
         UNIT_ASSERT_VALUES_EQUAL_C(NYdb::EStatus::BAD_REQUEST, status.GetStatus(), "The consumer cannot commit an offset for inactive, read-to-the-end partitions.");
     }
 
+    Y_UNIT_TEST(PartitionSplit_DistributedTxCommit_CheckSessionResetAfterCommit) {
+        TTopicSdkTestSetup setup = CreateSetup();
+        TTopicClient client = setup.MakeClient();
+
+        TCreateTopicSettings createSettings;
+        createSettings
+            .BeginConfigurePartitioningSettings()
+                .MinActivePartitions(1)
+                .MaxActivePartitions(100)
+                .BeginConfigureAutoPartitioningSettings()
+                    .UpUtilizationPercent(2)
+                    .DownUtilizationPercent(1)
+                    .StabilizationWindow(TDuration::Seconds(2))
+                    .Strategy(EAutoPartitioningStrategy::ScaleUp)
+                .EndConfigureAutoPartitioningSettings()
+            .EndConfigurePartitioningSettings()
+            .BeginAddConsumer()
+                .ConsumerName(TEST_CONSUMER);
+
+        client.CreateTopic(TEST_TOPIC, createSettings).Wait();
+
+        auto msg = TString(1_MB, 'a');
+
+        auto writeSession_1 = CreateWriteSession(client, "producer-1", 0, std::string{TEST_TOPIC}, false);
+        auto writeSession_2 = CreateWriteSession(client, "producer-2", 0, std::string{TEST_TOPIC}, false);
+        auto seqNo = 1;
+        {
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, seqNo++)));
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, seqNo++)));
+            Sleep(TDuration::Seconds(5));
+            auto describe = client.DescribeTopic(TEST_TOPIC).GetValueSync();
+            UNIT_ASSERT_EQUAL(describe.GetTopicDescription().GetPartitions().size(), 1);
+        }
+
+        {
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, seqNo++)));
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, seqNo++)));
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, seqNo++)));
+            UNIT_ASSERT(writeSession_2->Write(Msg(msg, seqNo++)));
+            Sleep(TDuration::Seconds(15));
+            auto describe = client.DescribeTopic(TEST_TOPIC).GetValueSync();
+            UNIT_ASSERT_EQUAL(describe.GetTopicDescription().GetPartitions().size(), 3);
+        }
+
+        auto writeSession_3 = CreateWriteSession(client, "producer-2", 1, std::string{TEST_TOPIC}, false);
+        UNIT_ASSERT(writeSession_3->Write(Msg(TStringBuilder() << "message-" << seqNo, seqNo++)));
+        UNIT_ASSERT(writeSession_3->Write(Msg(TStringBuilder() << "message-" << seqNo, seqNo++)));
+
+        auto reader = client.CreateReadSession(
+            TReadSessionSettings()
+                .AutoPartitioningSupport(true)
+                .AppendTopics(TTopicReadSettings(TEST_TOPIC))
+                .ConsumerName(TEST_CONSUMER));
+
+        TInstant deadlineTime = TInstant::Now() + TDuration::Seconds(5);
+
+        auto commitSent = false;
+        while(deadlineTime > TInstant::Now()) {
+            for (auto event : reader->GetEvents(false)) {
+                if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent>(&event)) {
+                    Cerr << "SESSION EVENT " << x->DebugString() << Endl << Flush;
+                    auto& messages = x->GetMessages();
+                    for (size_t i = 0u; i < messages.size(); ++i) {
+                        auto& message = messages[i];
+                        message.Commit();
+                        Cerr << "SESSION EVENT READ SeqNo: " << message.GetSeqNo() << Endl << Flush;
+                        // check we get this SeqNo two times
+                        if (message.GetSeqNo() == 6) {
+                            if (!commitSent) {
+                                commitSent = true;
+                                Sleep(TDuration::MilliSeconds(300));
+                                auto status = client.CommitOffset(TEST_TOPIC, 0, TEST_CONSUMER, 0).GetValueSync();
+                                UNIT_ASSERT(status.IsSuccess());
+                            } else {
+                                return;
+                            }
+                        }
+                    }
+                    UNIT_ASSERT(writeSession_3->Write(Msg(TStringBuilder() << "message-" << seqNo, seqNo++)));
+                } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TStartPartitionSessionEvent>(&event)) {
+                    x->Confirm();
+                    Cerr << "SESSION EVENT " << x->DebugString() << Endl << Flush;
+                } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TCommitOffsetAcknowledgementEvent>(&event)) {
+                    Cerr << "SESSION EVENT " << x->DebugString() << Endl << Flush;
+                } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TPartitionSessionStatusEvent>(&event)) {
+                    Cerr << "SESSION EVENT " << x->DebugString() << Endl << Flush;
+                } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TStopPartitionSessionEvent>(&event)) {
+                    x->Confirm();
+                    Cerr << "SESSION EVENT " << x->DebugString() << Endl << Flush;
+                } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TPartitionSessionClosedEvent>(&event)) {
+                    Cerr << "SESSION EVENT " << x->DebugString() << Endl << Flush;
+                } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TEndPartitionSessionEvent>(&event)) {
+                    x->Confirm();
+                    Cerr << "SESSION EVENT " << x->DebugString() << Endl << Flush;
+                } else if (auto* sessionClosedEvent = std::get_if<NYdb::NTopic::TSessionClosedEvent>(&event)) {
+                    Cerr << sessionClosedEvent->DebugString() << Endl << Flush;
+                } else {
+                    Cerr << "SESSION EVENT unhandled \n";
+                }
+            }
+            Sleep(TDuration::MilliSeconds(250));
+        }
+
+        UNIT_ASSERT(false);
+    }
+
+    Y_UNIT_TEST(PartitionSplit_DistributedTxCommit_CheckOffsetCommitForDifferentCases_NotSplitedTopic) {
+        TTopicSdkTestSetup setup = CreateSetup();
+        TTopicClient client = setup.MakeClient();
+
+        TCreateTopicSettings createSettings;
+        createSettings
+            .BeginConfigurePartitioningSettings()
+                .MinActivePartitions(1)
+                .MaxActivePartitions(100)
+                .BeginConfigureAutoPartitioningSettings()
+                    .UpUtilizationPercent(2)
+                    .DownUtilizationPercent(1)
+                    .StabilizationWindow(TDuration::Seconds(2))
+                    .Strategy(EAutoPartitioningStrategy::ScaleUp)
+                .EndConfigureAutoPartitioningSettings()
+            .EndConfigurePartitioningSettings()
+            .BeginAddConsumer()
+                .ConsumerName(TEST_CONSUMER);
+
+        client.CreateTopic(TEST_TOPIC, createSettings).Wait();
+
+        auto msg = TString(1_MB, 'a');
+
+        auto writeSession_1 = CreateWriteSession(client, "producer-1", 0, std::string{TEST_TOPIC}, false);
+        auto seqNo = 1;
+        {
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, seqNo++)));
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, seqNo++)));
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, seqNo++)));
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, seqNo++)));
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, seqNo++)));
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, seqNo++)));
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, seqNo++)));
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, seqNo++)));
+            writeSession_1->Close();
+            Sleep(TDuration::Seconds(15));
+            auto describe = client.DescribeTopic(TEST_TOPIC).GetValueSync();
+            UNIT_ASSERT_EQUAL(describe.GetTopicDescription().GetPartitions().size(), 1);
+        }
+
+        auto writeSession_2 = CreateWriteSession(client, "producer-1", 0, std::string{TEST_TOPIC}, false);
+        auto reader = client.CreateReadSession(
+            TReadSessionSettings()
+                .AutoPartitioningSupport(true)
+                .AppendTopics(TTopicReadSettings(TEST_TOPIC))
+                .ConsumerName(TEST_CONSUMER));
+
+        TInstant deadlineTime = TInstant::Now() + TDuration::Seconds(5);
+
+        auto commitSent = false;
+        TString readSessionId = "";
+        while(deadlineTime > TInstant::Now()) {
+            for (auto event : reader->GetEvents(false)) {
+                if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent>(&event)) {
+                    Cerr << "SESSION EVENT " << x->DebugString() << Endl << Flush;
+                    auto& messages = x->GetMessages();
+                    for (size_t i = 0u; i < messages.size(); ++i) {
+                        auto& message = messages[i];
+
+                        if (commitSent) {
+                            // read session not changed
+                            UNIT_ASSERT_EQUAL(readSessionId, message.GetPartitionSession()->GetReadSessionId());
+                        }
+
+                        // check we NOT get this SeqNo two times
+                        if (message.GetSeqNo() == 6) {
+                            if (!commitSent) {
+                                commitSent = true;
+                                Sleep(TDuration::MilliSeconds(300));
+
+                                readSessionId = message.GetPartitionSession()->GetReadSessionId();
+                                TCommitOffsetSettings commitSettings {.ReadSessionId_ = message.GetPartitionSession()->GetReadSessionId()};
+                                auto status = client.CommitOffset(TEST_TOPIC, 0, TEST_CONSUMER, 8, commitSettings).GetValueSync();
+                                UNIT_ASSERT(status.IsSuccess());
+
+                                {
+                                    auto describeConsumerSettings = TDescribeConsumerSettings().IncludeStats(true);
+                                    auto result = client.DescribeConsumer(TEST_TOPIC, TEST_CONSUMER, describeConsumerSettings).GetValueSync();
+                                    UNIT_ASSERT(result.IsSuccess());
+
+                                    auto description = result.GetConsumerDescription();
+
+                                    auto stats = description.GetPartitions().at(0).GetPartitionConsumerStats();
+                                    UNIT_ASSERT(stats);
+                                    UNIT_ASSERT(stats->GetCommittedOffset() == 8);
+                                }
+
+                                // must be ignored, because commit to past
+                                TCommitOffsetSettings commitToPastSettings {.ReadSessionId_ = message.GetPartitionSession()->GetReadSessionId()};
+                                auto commitToPastStatus = client.CommitOffset(TEST_TOPIC, 0, TEST_CONSUMER, 0, commitToPastSettings).GetValueSync();
+                                UNIT_ASSERT(commitToPastStatus.IsSuccess());
+
+                                {
+                                    auto describeConsumerSettings = TDescribeConsumerSettings().IncludeStats(true);
+                                    auto result = client.DescribeConsumer(TEST_TOPIC, TEST_CONSUMER, describeConsumerSettings).GetValueSync();
+                                    UNIT_ASSERT(result.IsSuccess());
+
+                                    auto description = result.GetConsumerDescription();
+
+                                    auto stats = description.GetPartitions().at(0).GetPartitionConsumerStats();
+                                    UNIT_ASSERT(stats);
+                                    UNIT_ASSERT(stats->GetCommittedOffset() == 8);
+                                }
+
+                                TCommitOffsetSettings commitSettingsWrongSession {.ReadSessionId_ = "random_session"};
+                                auto statusWrongSession = client.CommitOffset(TEST_TOPIC, 0, TEST_CONSUMER, 0, commitSettingsWrongSession).GetValueSync();
+                                UNIT_ASSERT(!statusWrongSession.IsSuccess());
+
+                                {
+                                    auto describeConsumerSettings = TDescribeConsumerSettings().IncludeStats(true);
+                                    auto result = client.DescribeConsumer(TEST_TOPIC, TEST_CONSUMER, describeConsumerSettings).GetValueSync();
+                                    UNIT_ASSERT(result.IsSuccess());
+
+                                    auto description = result.GetConsumerDescription();
+
+                                    auto stats = description.GetPartitions().at(0).GetPartitionConsumerStats();
+                                    UNIT_ASSERT(stats);
+                                    UNIT_ASSERT(stats->GetCommittedOffset() == 8);
+                                }
+
+                            } else {
+                                UNIT_ASSERT(false);
+                            }
+                        } else {
+                            message.Commit();
+                        }
+                    }
+                    UNIT_ASSERT(writeSession_2->Write(Msg(msg, seqNo++)));
+                } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TStartPartitionSessionEvent>(&event)) {
+                    x->Confirm();
+                    Cerr << "SESSION EVENT " << x->DebugString() << Endl << Flush;
+                } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TCommitOffsetAcknowledgementEvent>(&event)) {
+                    Cerr << "SESSION EVENT " << x->DebugString() << Endl << Flush;
+                } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TPartitionSessionStatusEvent>(&event)) {
+                    Cerr << "SESSION EVENT " << x->DebugString() << Endl << Flush;
+                } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TStopPartitionSessionEvent>(&event)) {
+                    x->Confirm();
+                    Cerr << "SESSION EVENT " << x->DebugString() << Endl << Flush;
+                } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TPartitionSessionClosedEvent>(&event)) {
+                    Cerr << "SESSION EVENT " << x->DebugString() << Endl << Flush;
+                } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TEndPartitionSessionEvent>(&event)) {
+                    x->Confirm();
+                    Cerr << "SESSION EVENT " << x->DebugString() << Endl << Flush;
+                } else if (auto* sessionClosedEvent = std::get_if<NYdb::NTopic::TSessionClosedEvent>(&event)) {
+                    Cerr << sessionClosedEvent->DebugString() << Endl << Flush;
+                } else {
+                    Cerr << "SESSION EVENT unhandled \n";
+                }
+            }
+            Sleep(TDuration::MilliSeconds(250));
+        }
+    }
+
     Y_UNIT_TEST(ControlPlane_CreateAlterDescribe) {
         auto autoscalingTestTopic = "autoscalit-topic";
         TTopicSdkTestSetup setup = CreateSetup();
diff --git a/ydb/public/api/protos/ydb_topic.proto b/ydb/public/api/protos/ydb_topic.proto
@@ -751,6 +751,9 @@ message CommitOffsetRequest {
 
     // Processed offset.
     int64 offset = 5;
+
+    // Read session identifier from StreamRead RPC.
+    string read_session_id = 6;
 }
 
 // Commit offset response sent from server to client.
diff --git a/ydb/public/sdk/cpp/client/ydb_topic/impl/read_session_impl.h b/ydb/public/sdk/cpp/client/ydb_topic/impl/read_session_impl.h
@@ -599,6 +599,7 @@ class TPartitionStreamImpl : public TAPartitionStream<UseMigrationProtocol> {
     template <bool V = UseMigrationProtocol, class = std::enable_if_t<!V>>
     TPartitionStreamImpl(ui64 partitionStreamId,
                          TString topicPath,
+                         TString readSessionId,
                          i64 partitionId,
                          i64 assignId,
                          i64 readOffset,
@@ -611,6 +612,7 @@ class TPartitionStreamImpl : public TAPartitionStream<UseMigrationProtocol> {
     {
         TAPartitionStream<false>::PartitionSessionId = partitionStreamId;
         TAPartitionStream<false>::TopicPath = std::move(topicPath);
+        TAPartitionStream<false>::ReadSessionId = std::move(readSessionId);
         TAPartitionStream<false>::PartitionId = static_cast<ui64>(partitionId);
         MaxCommittedOffset = static_cast<ui64>(readOffset);
     }
@@ -1296,6 +1298,7 @@ class TSingleClusterReadSessionImpl : public TEnableSelfContext<TSingleClusterRe
     const TString Database;
     const TString SessionId;
     const TString ClusterName;
+    TString ReadSessionId;
     TLog Log;
     ui64 NextPartitionStreamId;
     ui64 PartitionStreamIdStep;
diff --git a/ydb/public/sdk/cpp/client/ydb_topic/impl/read_session_impl.ipp b/ydb/public/sdk/cpp/client/ydb_topic/impl/read_session_impl.ipp
@@ -954,6 +954,7 @@ inline void TSingleClusterReadSessionImpl<true>::OnReadDoneImpl(
     LOG_LAZY(Log, TLOG_INFO, GetLogPrefix() << "Server session id: " << msg.session_id());
 
     RetryState = nullptr;
+    ReadSessionId = msg.session_id();
 
     // Successful init. Do nothing.
     ContinueReadingDataImpl();
@@ -1205,6 +1206,7 @@ inline void TSingleClusterReadSessionImpl<false>::OnReadDoneImpl(
     Y_UNUSED(deferred);
 
     RetryState = nullptr;
+    ReadSessionId = msg.session_id();
 
     LOG_LAZY(Log, TLOG_INFO, GetLogPrefix() << "Server session id: " << msg.session_id());
 
@@ -1304,8 +1306,12 @@ inline void TSingleClusterReadSessionImpl<false>::OnReadDoneImpl(
     Y_ABORT_UNLESS(Lock.IsLocked());
 
     auto partitionStream = MakeIntrusive<TPartitionStreamImpl<false>>(
-        NextPartitionStreamId, msg.partition_session().path(), msg.partition_session().partition_id(),
-        msg.partition_session().partition_session_id(), msg.committed_offset(),
+        NextPartitionStreamId,
+        msg.partition_session().path(),
+        ReadSessionId,
+        msg.partition_session().partition_id(),
+        msg.partition_session().partition_session_id(),
+        msg.committed_offset(),
         SelfContext);
     NextPartitionStreamId += PartitionStreamIdStep;
 
diff --git a/ydb/public/sdk/cpp/client/ydb_topic/impl/topic_impl.h b/ydb/public/sdk/cpp/client/ydb_topic/impl/topic_impl.h
@@ -323,6 +323,9 @@ class TTopicClient::TImpl : public TClientImplCommon<TTopicClient::TImpl> {
         request.set_partition_id(partitionId);
         request.set_consumer(consumerName);
         request.set_offset(offset);
+        if (settings.ReadSessionId_) {
+            request.set_read_session_id(*settings.ReadSessionId_);
+        }
 
         return RunSimple<Ydb::Topic::V1::TopicService, Ydb::Topic::CommitOffsetRequest, Ydb::Topic::CommitOffsetResponse>(
             std::move(request),
diff --git a/ydb/public/sdk/cpp/client/ydb_topic/include/control_plane.h b/ydb/public/sdk/cpp/client/ydb_topic/include/control_plane.h
@@ -749,6 +749,8 @@ struct TDescribePartitionSettings: public TOperationRequestSettings<TDescribePar
 };
 
 // Settings for commit offset request.
-struct TCommitOffsetSettings : public TOperationRequestSettings<TCommitOffsetSettings> {};
+struct TCommitOffsetSettings : public TOperationRequestSettings<TCommitOffsetSettings> {
+    FLUENT_SETTING_OPTIONAL(TString, ReadSessionId);
+};
 
 }  // namespace NYdb::NTopic
diff --git a/ydb/public/sdk/cpp/client/ydb_topic/include/read_events.h b/ydb/public/sdk/cpp/client/ydb_topic/include/read_events.h
@@ -32,6 +32,11 @@ struct TPartitionSession: public TThrRefBase, public TPrintable<TPartitionSessio
         return TopicPath;
     }
 
+    //! Read session id.
+    const TString GetReadSessionId() const {
+        return ReadSessionId;
+    }
+
     //! Partition id.
     ui64 GetPartitionId() const {
         return PartitionId;
@@ -40,6 +45,7 @@ struct TPartitionSession: public TThrRefBase, public TPrintable<TPartitionSessio
 protected:
     ui64 PartitionSessionId;
     TString TopicPath;
+    TString ReadSessionId;
     ui64 PartitionId;
 };
 
diff --git a/ydb/services/persqueue_v1/actors/commit_offset_actor.cpp b/ydb/services/persqueue_v1/actors/commit_offset_actor.cpp
@@ -127,6 +127,9 @@ void TCommitOffsetActor::Handle(TEvPQProxy::TEvAuthResultOk::TPtr& ev, const TAc
     commit->SetClientId(ClientId);
     commit->SetOffset(client_req->offset());
     commit->SetStrict(true);
+    if (!client_req->read_session_id().empty()) {
+        commit->SetSessionId(client_req->read_session_id());
+    }
 
     LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, "strict CommitOffset, partition " << client_req->partition_id()
                         << " committing to position " << client_req->offset() /*<< " prev " << CommittedOffset

Original file line number	Diff line number	Diff line change
`@@ -751,6 +751,9 @@ message CommitOffsetRequest {`
`751`	`751`
`752`	`752`	`// Processed offset.`
`753`	`753`	`int64 offset = 5;`
	`754`	`+`
	`755`	`+ // Read session identifier from StreamRead RPC.`
	`756`	`+ string read_session_id = 6;`
`754`	`757`	`}`
`755`	`758`
`756`	`759`	`// Commit offset response sent from server to client.`