Merge 5cb08cd into d8e4b72

Author: ydbot

ydb/core/persqueue/pqrb/read_balancer__balancing.cpp

@@ -386,8 +386,8 @@ void TPartitionFamily::InactivatePartition(ui32 partitionId) {
 }
 
  void TPartitionFamily::ChangePartitionCounters(ssize_t active, ssize_t inactive) {
-    Y_VERIFY_DEBUG((ssize_t)ActivePartitionCount + active >= 0);
-    Y_VERIFY_DEBUG((ssize_t)InactivePartitionCount + inactive >= 0);
+    Y_VERIFY_DEBUG((ssize_t)ActivePartitionCount + active >= 0, "ActivePartitionCount: %lu, active: %ld", ActivePartitionCount, active);
+    Y_VERIFY_DEBUG((ssize_t)InactivePartitionCount + inactive >= 0, "InactivePartitionCount: %lu, inactive: %ld", InactivePartitionCount, inactive);
 
     ActivePartitionCount += active;
     InactivePartitionCount += inactive;
@@ -1155,11 +1155,12 @@ void TConsumer::FinishReading(TEvPersQueue::TEvReadingPartitionFinishedRequest::
 
     auto& partition = Partitions[partitionId];
 
-    if (partition.SetFinishedState(r.GetScaleAwareSDK(), r.GetStartedReadingFromEndOffset())) {
+    const bool wasInactive = partition.IsInactive();
+    if (partition.SetFinishedState(r.GetScaleAwareSDK(), r.GetStartedReadingFromEndOffset()) || wasInactive) {
         PQ_LOG_D("Reading of the partition " << partitionId << " was finished by " << r.GetConsumer()
                 << ", firstMessage=" << r.GetStartedReadingFromEndOffset() << ", " << GetSdkDebugString0(r.GetScaleAwareSDK()));
 
-        if (ProccessReadingFinished(partitionId, false, ctx)) {
+        if (ProccessReadingFinished(partitionId, wasInactive, ctx)) {
             ScheduleBalance(ctx);
         }
     } else if (!partition.IsInactive()) {
@@ -1835,6 +1836,62 @@ void TBalancer::ProcessPendingStats(const TActorContext& ctx) {
     PendingUpdates.clear();
 }
 
+<<<<<<< HEAD
+=======
+void TBalancer::Handle(TEvPersQueue::TEvBalancingSubscribe::TPtr& ev, const TActorContext& ctx) {
+    auto& record = ev->Get()->Record;
+    PQ_LOG_D("Handle TEvPersQueue::TEvBalancingSubscribe " << record.ShortDebugString());
+
+    auto sender = ActorIdFromProto(record.GetSourceActor());
+    auto status = Consumers.contains(record.GetConsumer()) ?
+        NKikimrPQ::TEvBalancingSubscribeNotify::BALANCING : NKikimrPQ::TEvBalancingSubscribeNotify::FREE;
+    Notify(sender, record.GetConsumer(), status, ctx);
+
+    Subscriptions[ev->Sender].emplace_back(std::move(sender), std::move(*record.MutableConsumer()));
+}
+
+void TBalancer::Handle(TEvPersQueue::TEvBalancingUnsubscribe::TPtr& ev, const TActorContext&) {
+    auto& record = ev->Get()->Record;
+    PQ_LOG_D("Handle TEvPersQueue::TEvBalancingUnsubscribe " << record.ShortDebugString());
+
+    auto sender = ActorIdFromProto(record.GetSourceActor());
+    auto& consumer = record.GetConsumer();
+
+    auto it = Subscriptions.find(ev->Sender);
+    if (it == Subscriptions.end()) {
+        return;
+    }
+
+    std::vector<TSubscription>& subscriptions = it->second;
+    std::vector<TSubscription> actualSubscriptions;
+    actualSubscriptions.resize(subscriptions.size());
+
+    for (auto& [existsSender, existsConsumer] : subscriptions) {
+        if (sender == existsSender && consumer == existsConsumer) {
+            continue;
+        }
+
+        actualSubscriptions.emplace_back(std::move(existsSender), std::move(existsConsumer));
+    }
+
+    subscriptions = std::move(actualSubscriptions);
+}
+
+void TBalancer::Notify(const TString& consumer, NKikimrPQ::TEvBalancingSubscribeNotify::EStatus status, const TActorContext& ctx) {
+    for (auto& [_, subscriptions] : Subscriptions) {
+        for (auto& subscription : subscriptions) {
+            if (subscription.Consumer == consumer) {
+                Notify(subscription.Sender, consumer, status, ctx);
+            }
+        }
+    }
+}
+
+void TBalancer::Notify(const TActorId subscriber, const TString& consumer, NKikimrPQ::TEvBalancingSubscribeNotify::EStatus status, const TActorContext& ctx) {
+    ctx.Send(subscriber, new TEvPersQueue::TEvBalancingSubscribeNotify(TabletGeneration(), ++NotifyCookie, TopicPath(), consumer, status));
+}
+
+>>>>>>> 45f5616be25 (EXT-1793 Fix when reading from the end of uncommitted partition (#30854))
 TString TBalancer::LogPrefix() const {
     return TStringBuilder() << "[" << TopicActor.TabletID() << "][" << Topic() << "] ";
 }

ydb/core/persqueue/pqrb/read_balancer__balancing.h

@@ -16,7 +16,7 @@ struct TConsumer;
 class TBalancer;
 
 struct TPartition {
-    // Client had commited rad offset equals EndOffset of the partition
+    // Client had commited read offset equals EndOffset of the partition
     bool Commited = false;
     // ReadSession reach EndOffset of the partition
     bool ReadingFinished = false;

ydb/core/persqueue/ut/ut_with_sdk/autoscaling_ut.cpp

@@ -622,6 +622,64 @@ Y_UNIT_TEST_SUITE(TopicAutoscaling) {
         readSession2->Close();
     }
 
+    Y_UNIT_TEST(PartitionSplit_ManySessions_NoCommits_AutoscaleAwareSDK) {
+        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();
+
+        TConsumerSettings<TCreateTopicSettings> consumers(createSettings, setup.GetConsumerName(TEST_CONSUMER));
+        createSettings.AppendConsumers(consumers);
+        client.CreateTopic(TEST_TOPIC, createSettings).Wait();
+
+        auto msg = TString(1_MB, 'a');
+
+        auto writeSession_1 = CreateWriteSession(client, "producer-1", 0, TString{TEST_TOPIC}, false);
+        auto writeSession_2 = CreateWriteSession(client, "producer-2", 0, TString{TEST_TOPIC}, false);
+
+        {
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, 1)));
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, 2)));
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, 3)));
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, 4)));
+            UNIT_ASSERT(writeSession_1->Write(Msg(msg, 5)));
+            UNIT_ASSERT(writeSession_2->Write(Msg(msg, 6)));
+            Sleep(TDuration::Seconds(15));
+            auto describe = client.DescribeTopic(TEST_TOPIC).GetValueSync();
+            UNIT_ASSERT_VALUES_EQUAL(describe.GetTopicDescription().GetPartitions().size(), 3);
+        }
+
+        auto readSession1 = CreateTestReadSession({ .Name="Session-1", .Setup=setup, .Sdk = SdkVersion::Topic, .ExpectedMessagesCount = 6, .AutoCommit = false, .AutoPartitioningSupport = true });
+        readSession1->WaitAndAssertPartitions({0, 1, 2}, "Must read all partition");
+        readSession1->WaitAllMessages();
+
+        auto readSession2 = CreateTestReadSession({ .Name="Session-2", .Setup=setup, .Sdk = SdkVersion::Topic, .ExpectedMessagesCount = 0, .AutoCommit = false, .AutoPartitioningSupport = true });
+
+        Sleep(TDuration::Seconds(1));
+
+        readSession1->Close();
+
+        auto yetAnotherReadSession1 = CreateTestReadSession({ .Name="YetAnotherSession-1", .Setup=setup, .Sdk = SdkVersion::Topic, .ExpectedMessagesCount = 0, .AutoCommit = false, .AutoPartitioningSupport = true });
+        yetAnotherReadSession1->SetOffset(0, 6); // read from end offset
+
+        Sleep(TDuration::Seconds(1));
+
+        readSession2->Close();
+
+        yetAnotherReadSession1->WaitAndAssertPartitions({0, 1, 2}, "Must read all children partitions");
+    }
+
     Y_UNIT_TEST(ControlPlane_CreateAlterDescribe) {
         auto autoscalingTestTopic = "autoscalit-topic";
         TTopicSdkTestSetup setup = CreateSetup();