Optimization of the algorithm for determining the end of partition processing

ydb/core/persqueue/events/global.h

@@ -275,11 +275,12 @@ namespace TEvPersQueue {
     struct TEvReadingPartitionFinishedRequest : public TEventPB<TEvReadingPartitionFinishedRequest, NKikimrPQ::TEvReadingPartitionFinishedRequest, EvReadingPartitionFinished> {
         TEvReadingPartitionFinishedRequest() = default;
 
-        TEvReadingPartitionFinishedRequest(const TString& consumer, ui32 partitionId, bool scaleAwareSDK, bool startedReadingFromEndOffset) {
+        TEvReadingPartitionFinishedRequest(const TString& consumer, ui32 partitionId, bool scaleAwareSDK, bool startedReadingFromEndOffset, ui64 readTimestampMs = 0) {
             Record.SetConsumer(consumer);
             Record.SetPartitionId(partitionId);
             Record.SetScaleAwareSDK(scaleAwareSDK);
             Record.SetStartedReadingFromEndOffset(startedReadingFromEndOffset);
+            Record.SetReadTimestampMs(readTimestampMs);
         }
     };
 

ydb/core/persqueue/partition.cpp

@@ -750,8 +750,11 @@ void TPartition::Handle(TEvPQ::TEvPartitionStatus::TPtr& ev, const TActorContext
 
     }
 
+    result.SetPartitionStatus(PartitionConfig->GetStatus());
+
     result.SetStartOffset(StartOffset);
     result.SetEndOffset(EndOffset);
+    result.SetEndWriteTimestampMs(EndWriteTimestamp.MilliSeconds());
 
     if (filterConsumers) {
         for (TString consumer : requiredConsumers) {

ydb/core/persqueue/read_balancer.cpp

@@ -565,9 +565,17 @@ void TPersQueueReadBalancer::Handle(TEvPersQueue::TEvStatusResponse::TPtr& ev, c
 
     for (const auto& partRes : record.GetPartResult()) {
         ui32 partitionId = partRes.GetPartition();
-        if (!PartitionsInfo.contains(partitionId)) {
+
+        auto p = PartitionsInfo.find(partitionId);
+        if (p == PartitionsInfo.end()) {
             continue;
         }
+        auto& partitionInfo = p->second;
+
+        if (partRes.HasPartitionStatus() && partRes.HasEndWriteTimestampMs()) {
+            partitionInfo.Status = partRes.GetPartitionStatus();
+            partitionInfo.EndWriteTimestamp = TInstant::MilliSeconds(partRes.GetEndWriteTimestampMs());
+        }
 
         if (SplitMergeEnabled(TabletConfig) && PartitionsScaleManager) {
             PartitionsScaleManager->HandleScaleStatusChange(partitionId, partRes.GetScaleStatus(), ctx);

ydb/core/persqueue/read_balancer.h

@@ -177,6 +177,8 @@ class TPersQueueReadBalancer : public TActor<TPersQueueReadBalancer>, public TTa
 public:
     struct TPartitionInfo {
         ui64 TabletId;
+        NKikimrPQ::ETopicPartitionStatus Status = NKikimrPQ::ETopicPartitionStatus::Active;
+        TInstant EndWriteTimestamp;
     };
 
 private:

ydb/core/persqueue/read_balancer__balancing.cpp

@@ -1161,7 +1161,23 @@ void TConsumer::FinishReading(TEvPersQueue::TEvReadingPartitionFinishedRequest::
             ScheduleBalance(ctx);
         }
     } else if (!partition.IsInactive()) {
-        auto delay = std::min<size_t>(1ul << partition.Iteration, Balancer.GetLifetimeSeconds()); // TODO use split/merge time
+        auto now = TInstant::Now();
+
+        auto* partitionInfo = GetPartitionInfo(partitionId);
+        bool hasEndWriteTimestamp = partitionInfo && partitionInfo->Status == NKikimrPQ::ETopicPartitionStatus::Inactive;
+        auto endWriteTimestamp = hasEndWriteTimestamp ? partitionInfo->EndWriteTimestamp : now;
+
+        size_t delay;
+        if (r.GetReadTimestampMs() && hasEndWriteTimestamp) {
+            TInstant readTimestamp = TInstant::MilliSeconds(r.GetReadTimestampMs());
+            if (readTimestamp >= endWriteTimestamp) {
+                delay = 1;
+            } else {
+                delay = (endWriteTimestamp - readTimestamp).Seconds() + 1;
+            }
+        } else {
+            delay = std::min<size_t>(1ul << partition.Iteration, Balancer.GetLifetimeSeconds() - (now - endWriteTimestamp).Seconds() + 1);
+        }
 
         PQ_LOG_D("Reading of the partition " << partitionId << " was finished by " << r.GetConsumer()
                 << ". Scheduled release of the partition for re-reading. Delay=" << delay << " seconds,"

ydb/core/protos/pqconfig.proto

@@ -833,6 +833,8 @@ message TStatusResponse {
         optional uint64 Cookie = 34;
 
         optional EScaleStatus ScaleStatus = 35;
+        optional ETopicPartitionStatus PartitionStatus = 36;
+        optional uint64 EndWriteTimestampMs = 37;
     }
 
     message TConsumerResult {
@@ -1072,6 +1074,7 @@ message TEvReadingPartitionFinishedRequest {
     optional uint32 PartitionId = 2;
     optional bool ScaleAwareSDK = 3;
     optional bool StartedReadingFromEndOffset = 4;
+    optional uint64 ReadTimestampMs = 5; //optional, default = 0
 };
 
 // The consumer's reading of the partition is started (from ReadSession)

ydb/services/persqueue_v1/actors/read_session_actor.cpp

@@ -2333,8 +2333,13 @@ void TReadSessionActor<UseMigrationProtocol>::Handle(TEvPQProxy::TEvReadingFinis
         return;
     }
 
+    auto [_, maxLag, readTimestampMs] = GetReadFrom(it->second.FullConverter, ctx);
+    if (!readTimestampMs && maxLag) {
+        readTimestampMs = (TInstant::Now() - TDuration::MilliSeconds(maxLag)).MilliSeconds();
+    }
+
     auto& topic = it->second;
-    NTabletPipe::SendData(ctx, topic.PipeClient, new TEvPersQueue::TEvReadingPartitionFinishedRequest(ClientId, msg->PartitionId, AutoPartitioningSupport, msg->FirstMessage));
+    NTabletPipe::SendData(ctx, topic.PipeClient, new TEvPersQueue::TEvReadingPartitionFinishedRequest(ClientId, msg->PartitionId, AutoPartitioningSupport, msg->FirstMessage, readTimestampMs));
 
     if constexpr (!UseMigrationProtocol) {
         if (AutoPartitioningSupport) {