Fix generation mismatch handler

qyryq · qyryq · commit a0fcd506ae2d · 2025-05-06T13:41:02.000Z
Set CmdForgetReadResult.DirectReadId

Correct bytes bookkeeping on tablet restarts

Stop adding messages to response when LastOffset is reached

Fix for direct read (on behalf of FloatingCrowbar)
diff --git a/ydb/core/persqueue/dread_cache_service/caching_service.cpp b/ydb/core/persqueue/dread_cache_service/caching_service.cpp
@@ -72,15 +72,16 @@ class TPQDirectReadCacheService : public TActorBootstrapped<TPQDirectReadCacheSe
             CloseSession(ev->Sender, Ydb::PersQueue::ErrorCode::ErrorCode::BAD_REQUEST, "Unknown session");
             return;
         }
+
+        auto sender = ev->Sender;
         if (sessionIter->second.Generation != ev->Get()->Generation) {
             ctx.Send(
-                    sessionIter->second.Client->ProxyId,
-                    new TEvPQProxy::TEvDirectReadDestroyPartitionSession(key, Ydb::PersQueue::ErrorCode::ErrorCode::ERROR, "Generation mismatch")
+                sender,
+                new TEvPQProxy::TEvDirectReadDestroyPartitionSession(key, Ydb::PersQueue::ErrorCode::ErrorCode::ERROR, "Generation mismatch")
             );
             return;
         }
 
-        auto sender = ev->Sender;
         auto startingReadId = ev->Get()->StartingReadId;
 
         // Let the proxy respond with StartDirectReadPartitionSessionResponse right away,
@@ -344,6 +345,7 @@ class TPQDirectReadCacheService : public TActorBootstrapped<TPQDirectReadCacheSe
         auto* directReadMessage = message->mutable_direct_read_response();
         directReadMessage->set_direct_read_id(readId);
         directReadMessage->set_partition_session_id(partSessionId);
+        directReadMessage->set_bytes_size(response->GetPartitionResponse().GetCmdPrepareReadResult().GetBytesSizeEstimate());
 
         auto ok = VaildatePartitionResponse(proxyClient, *response);
         if (!ok) {
diff --git a/ydb/core/persqueue/partition_read.cpp b/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);
@@ -515,6 +515,11 @@ TReadAnswer TReadInfo::FormAnswer(
                     << " size " << header.GetPayloadSize() << " from pos " << pos << " cbcount " << batch.Blobs.size());
 
             for (size_t i = pos; i < batch.Blobs.size(); ++i) {
+                if (0 < LastOffset && LastOffset <= Offset) {
+                    needStop = true;
+                    break;
+                }
+
                 TClientBlob &res = batch.Blobs[i];
                 VERIFY_RESULT_BLOB(res, i);
 
@@ -640,7 +645,7 @@ TVector<TRequestedBlob> TPartition::GetReadRequestFromBody(
         }
         while (it != DataKeysBody.end()
                && (size < maxSize && count < maxCount || count == 0) //count== 0 grants that blob with offset from ReadFromTimestamp will be readed
-               && (lastOffset == 0 || it->Key.GetOffset() <= lastOffset)
+               && (lastOffset == 0 || it->Key.GetOffset() < lastOffset)
         ) {
             size += sz;
             count += cnt;
@@ -688,7 +693,7 @@ TVector<TClientBlob> TPartition::GetReadRequestFromHead(
             Y_ABORT_UNLESS(pno == blobs[i].GetPartNo());
             bool skip = offset < startOffset || offset == startOffset &&
                 blobs[i].GetPartNo() < partNo;
-            if (lastOffset != 0 && lastOffset < offset) {
+            if (0 < lastOffset && lastOffset <= offset) {
                 break;
             }
             if (blobs[i].IsLastPart()) {
diff --git a/ydb/core/persqueue/pq_impl.cpp b/ydb/core/persqueue/pq_impl.cpp
@@ -227,6 +227,7 @@ class TReadProxy : public TActorBootstrapped<TReadProxy> {
                 }
             }
         }
+
         if (!partResp->GetResult().empty()) {
             const auto& lastRes = partResp->GetResult(partResp->GetResult().size() - 1);
             if (lastRes.HasPartNo() && lastRes.GetPartNo() + 1 < lastRes.GetTotalParts()) { //last res is not full
@@ -260,7 +261,10 @@ class TReadProxy : public TActorBootstrapped<TReadProxy> {
         }
         if (isDirectRead) {
             auto* prepareResponse = Response->Record.MutablePartitionResponse()->MutableCmdPrepareReadResult();
-            prepareResponse->SetBytesSizeEstimate(readResult.GetSizeEstimate());
+            auto sizeEstimate = Request.GetPartitionRequest().GetCmdRead().GetSizeEstimate();
+            sizeEstimate = sizeEstimate ? sizeEstimate : PreparedResponse->GetPartitionResponse().ByteSize();
+            PreparedResponse->MutablePartitionResponse()->MutableCmdPrepareReadResult()->SetBytesSizeEstimate(sizeEstimate);
+            prepareResponse->SetBytesSizeEstimate(sizeEstimate);
             prepareResponse->SetDirectReadId(DirectReadKey.ReadId);
             prepareResponse->SetReadOffset(readResult.GetRealReadOffset());
             prepareResponse->SetLastOffset(readResult.GetLastOffset());
@@ -2318,15 +2322,18 @@ void TPersQueue::HandleReadRequest(
                 return;
             }
         }
+
         THolder<TEvPQ::TEvRead> event =
             MakeHolder<TEvPQ::TEvRead>(responseCookie, cmd.GetOffset(), cmd.GetLastOffset(),
                                        cmd.HasPartNo() ? cmd.GetPartNo() : 0,
                                        count,
                                        cmd.HasSessionId() ? cmd.GetSessionId() : "",
                                        cmd.GetClientId(),
-                                       cmd.HasTimeoutMs() ? cmd.GetTimeoutMs() : 0, bytes,
+                                       cmd.HasTimeoutMs() ? cmd.GetTimeoutMs() : 0,
+                                       bytes,
                                        cmd.HasMaxTimeLagMs() ? cmd.GetMaxTimeLagMs() : 0,
-                                       cmd.HasReadTimestampMs() ? cmd.GetReadTimestampMs() : 0, clientDC,
+                                       cmd.HasReadTimestampMs() ? cmd.GetReadTimestampMs() : 0,
+                                       clientDC,
                                        cmd.GetExternalOperation(),
                                        pipeClient);
 
@@ -2407,8 +2414,7 @@ void TPersQueue::HandleForgetReadRequest(
     THolder<TEvPQ::TEvProxyResponse> forgetDoneEvent = MakeHolder<TEvPQ::TEvProxyResponse>(responseCookie);
     forgetDoneEvent->Response->SetStatus(NMsgBusProxy::MSTATUS_OK);
     forgetDoneEvent->Response->SetErrorCode(NPersQueue::NErrorCode::OK);
-
-    forgetDoneEvent->Response->MutablePartitionResponse()->MutableCmdForgetReadResult();
+    forgetDoneEvent->Response->MutablePartitionResponse()->MutableCmdForgetReadResult()->SetDirectReadId(key.ReadId);
     ctx.Send(SelfId(), forgetDoneEvent.Release());
 
     PQ_LOG_D("Forget direct read id " << key.ReadId << " for session " << key.SessionId);
diff --git a/ydb/core/persqueue/ut/common/pq_ut_common.cpp b/ydb/core/persqueue/ut/common/pq_ut_common.cpp
@@ -1029,6 +1029,9 @@ void BeginCmdRead(const TPQCmdReadSettings& settings, TTestContext& tc)
     if (settings.PartitionSessionId > 0) {
         read->SetPartitionSessionId(settings.PartitionSessionId);
     }
+    if (settings.LastOffset > 0) {
+        read->SetLastOffset(settings.LastOffset);
+    }
     if (settings.Pipe) {
         ActorIdToProto(settings.Pipe, req->MutablePipeClient());
     }
diff --git a/ydb/core/persqueue/ut/common/pq_ut_common.h b/ydb/core/persqueue/ut/common/pq_ut_common.h
@@ -449,11 +449,12 @@ struct TPQCmdReadSettings : public TPQCmdSettingsBase {
     ui32 MaxTimeLagMs = 0;
     ui32 ReadTimestampMs = 0;
     ui64 DirectReadId = 0;
+    i64 LastOffset = 0;
     TActorId Pipe;
     TPQCmdReadSettings() = default;
     TPQCmdReadSettings(const TString& session, ui32 partition, i64 offset, ui32 count, ui32 size, ui32 resCount, bool timeout = false,
                        TVector<i32> offsets = {}, const ui32 maxTimeLagMs = 0, const ui64 readTimestampMs = 0,
-                       const TString user = "user")
+                       const TString user = "user", const i64 lastOffset = 0)
 
         : TPQCmdSettingsBase{partition, user, session, 0, offset, false}
         , Count(count)
@@ -463,6 +464,7 @@ struct TPQCmdReadSettings : public TPQCmdSettingsBase {
         , Offsets (offsets)
         , MaxTimeLagMs(maxTimeLagMs)
         , ReadTimestampMs(readTimestampMs)
+        , LastOffset(lastOffset)
     {}
 };
 
diff --git a/ydb/core/persqueue/ut/pq_ut.cpp b/ydb/core/persqueue/ut/pq_ut.cpp
@@ -20,6 +20,75 @@ const static TString TOPIC_NAME = "rt3.dc1--topic";
 
 Y_UNIT_TEST_SUITE(TPQTest) {
 
+Y_UNIT_TEST(TestCmdReadWithLastOffset) {
+    TTestContext tc;
+    tc.EnableDetailedPQLog = true;
+    RunTestWithReboots(tc.TabletIds, [&]() {
+        return tc.InitialEventsFilter.Prepare();
+    }, [&](const TString& dispatchName, std::function<void(TTestActorRuntime&)> setup, bool& activeZone) {
+        activeZone = false;
+        TFinalizer finalizer(tc);
+        tc.Prepare(dispatchName, setup, activeZone);
+        activeZone = false;
+        tc.Runtime->SetScheduledLimit(1000);
+        tc.Runtime->RegisterService(MakePQDReadCacheServiceActorId(), tc.Runtime->Register(
+                CreatePQDReadCacheService(new NMonitoring::TDynamicCounters()))
+        );
+
+        PQTabletPrepare({.partitions = 1, .writeSpeed = 100_KB}, {{"user1", true}}, tc);
+        TVector<std::pair<ui64, TString>> data;
+        i64 messageCount = 100;
+        for (i64 i = 1; i <= messageCount; ++i) {
+            data.push_back({i, TString(100_KB, 'a')});
+        }
+        CmdWrite(0, "sourceid0", data, tc, false, {}, false, "", -1, 0, false, false, true);
+        TString sessionId = "session1";
+        TString user = "user1";
+        TPQCmdSettings sessionSettings{0, user, sessionId};
+        sessionSettings.PartitionSessionId = 1;
+        sessionSettings.KeepPipe = true;
+        TPQCmdReadSettings readSettings{
+            /*session=*/ sessionId,
+            /*partition=*/ 0,
+            /*offset=*/ 0,
+            /*count=*/ static_cast<ui32>(messageCount),
+            /*size=*/ 16_MB,
+            /*resCount=*/ 0,
+        };
+        readSettings.PartitionSessionId = 1;
+        readSettings.User = user;
+
+        activeZone = false;
+        Cerr << "Create session\n";
+        auto pipe = CmdCreateSession(sessionSettings, tc);
+        readSettings.Pipe = pipe;
+
+        for (i64 offset = 0; offset < messageCount; offset += 10) {
+            for (i64 lastOffset = 0; lastOffset <= messageCount; lastOffset += 10) {
+                readSettings.Offset = offset;
+                readSettings.LastOffset = lastOffset;
+                readSettings.ResCount = lastOffset < offset ? 0 : static_cast<ui32>(lastOffset - offset);
+                BeginCmdRead(readSettings, tc);
+
+                TAutoPtr<IEventHandle> handle;
+                auto* result = tc.Runtime->GrabEdgeEvent<TEvPersQueue::TEvResponse>(handle);
+
+                UNIT_ASSERT_C(result->Record.GetPartitionResponse().HasCmdReadResult(), result->Record.GetPartitionResponse().DebugString());
+                auto res = result->Record.GetPartitionResponse().GetCmdReadResult();
+
+                if (lastOffset) {
+                    UNIT_ASSERT_VALUES_EQUAL(readSettings.ResCount, res.ResultSize());
+                }
+
+                for (size_t i = 0; i < res.ResultSize(); ++i) {
+                    UNIT_ASSERT_EQUAL(res.GetResult(i).GetOffset(), offset + i);
+                    UNIT_ASSERT_EQUAL(res.GetResult(i).GetData(), data[offset + i].second);
+                }
+            }
+        }
+    });
+}
+
 Y_UNIT_TEST(TestDirectReadHappyWay) {
     TTestContext tc;
     tc.EnableDetailedPQLog = true;
diff --git a/ydb/core/protos/msgbus_pq.proto b/ydb/core/protos/msgbus_pq.proto
@@ -36,6 +36,7 @@ message TPersQueuePartitionRequest {
         optional uint64 PartitionSessionId = 16;
         optional int64 LastOffset = 17;
 
+        optional uint64 SizeEstimate = 18;
     }
 
     message TCmdPublishDirectRead {
diff --git a/ydb/core/testlib/test_pq_client.h b/ydb/core/testlib/test_pq_client.h
@@ -1455,6 +1455,7 @@ class TFlatMsgBusPQClient : public NFlatTests::TFlatMsgBusClient {
     struct CreateTopicNoLegacyParams {
         TString Name;
         ui32 PartsCount;
+        TDuration RetentionPeriod = TDuration::Hours(18);
         bool DoWait = true;
         bool CanWrite = true;
         std::optional<TString> Dc = std::nullopt;
@@ -1477,6 +1478,7 @@ class TFlatMsgBusPQClient : public NFlatTests::TFlatMsgBusClient {
         auto settings = NYdb::NPersQueue::TCreateTopicSettings().PartitionsCount(params.PartsCount).ClientWriteDisabled(!params.CanWrite);
         settings.FederationAccount(params.Account);
         settings.SupportedCodecs(params.Codecs);
+        settings.RetentionPeriod(params.RetentionPeriod);
         //settings.MaxPartitionWriteSpeed(50_MB);
         //settings.MaxPartitionWriteBurst(50_MB);
         TVector<NYdb::NPersQueue::TReadRuleSettings> rrSettings;
diff --git a/ydb/public/api/protos/ydb_topic.proto b/ydb/public/api/protos/ydb_topic.proto
@@ -688,6 +688,10 @@ message StreamDirectReadMessage {
 
         // Messages data
         StreamReadMessage.ReadResponse.PartitionData partition_data = 3;
+
+        // Total size in bytes of this response as calculated by server.
+        // See ReadRequest comment above.
+        int64 bytes_size = 4;
     }
 }
 
diff --git a/ydb/services/persqueue_v1/actors/partition_actor.cpp b/ydb/services/persqueue_v1/actors/partition_actor.cpp
@@ -685,6 +685,16 @@ void TPartitionActor::Handle(TEvPersQueue::TEvResponse::TPtr& ev, const TActorCo
         }
         return;
     }
+
+    if (!(
+            result.HasCmdReadResult() || result.HasCmdPrepareReadResult() || result.HasCmdPublishReadResult()
+            || result.HasCmdForgetReadResult() || result.HasCmdRestoreDirectReadResult()
+    )) {
+        // this is commit response
+        CommitDone(result.GetCookie(), ctx);
+        return;
+    }
+
     switch (DirectReadRestoreStage) {
         case EDirectReadRestoreStage::None:
             break;
@@ -696,14 +706,19 @@ void TPartitionActor::Handle(TEvPersQueue::TEvResponse::TPtr& ev, const TActorCo
             return;
         case EDirectReadRestoreStage::Prepare:
             Y_ABORT_UNLESS(RestoredDirectReadId != 0);
-
+            if (!result.HasCmdPrepareReadResult()) {
+                LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " Invalid response on direct read restore for " << Partition << ": expect PrepareReadResult");
+            }
             Y_ABORT_UNLESS(result.HasCmdPrepareReadResult());
             Y_ABORT_UNLESS(DirectReadsToRestore.begin()->first == result.GetCmdPrepareReadResult().GetDirectReadId());
             DirectReadsToRestore.erase(DirectReadsToRestore.begin());
             {
                 auto sent = SendNextRestorePublishRequest();
-                if (!sent) // Nothing to publish
+                if (!sent) {
+                    // Read was not published previously and thus no response sent to session. Need to keep it
+                    UnpublishedDirectReads.insert(result.GetCmdPrepareReadResult().GetDirectReadId());
                     sent = SendNextRestorePrepareOrForget();
+                }
                 if (!sent)
                     OnDirectReadsRestored();
             }
@@ -728,11 +743,6 @@ void TPartitionActor::Handle(TEvPersQueue::TEvResponse::TPtr& ev, const TActorCo
             return;
     }
 
-    if (!(result.HasCmdReadResult() || result.HasCmdPrepareReadResult() || result.HasCmdPublishReadResult() || result.HasCmdForgetReadResult())) { // this is commit response
-        CommitDone(result.GetCookie(), ctx);
-        return;
-    }
-
     if (result.HasCmdForgetReadResult()) {
         // ignore it
         return;
@@ -1164,8 +1174,8 @@ bool TPartitionActor::SendNextRestorePrepareOrForget() {
         DirectReadRestoreStage = EDirectReadRestoreStage::Prepare;
         Y_ABORT_UNLESS(dr.GetReadOffset() <= dr.GetLastOffset());
 
-        auto request = MakeReadRequest(dr.GetReadOffset(), dr.GetLastOffset(), std::numeric_limits<i32>::max(),
-                                    std::numeric_limits<i32>::max(), 0, 0, dr.GetDirectReadId());
+        auto request = MakeReadRequest(dr.GetReadOffset(), dr.GetLastOffset() + 1, std::numeric_limits<i32>::max(),
+                                    std::numeric_limits<i32>::max(), 0, 0, dr.GetDirectReadId(), dr.GetBytesSizeEstimate());
 
         if (!PipeClient) //Pipe will be recreated soon
             return true;
@@ -1205,6 +1215,16 @@ void TPartitionActor::OnDirectReadsRestored() {
     const auto& ctx = ActorContext();
     LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " " << Partition
                             << "Restore direct reads done, continue working");
+
+    if (InitDone) {
+        ctx.Send(ParentId, new TEvPQProxy::TEvUpdateSession(Partition, NodeId, TabletGeneration));
+    }
+
+    for (auto id: UnpublishedDirectReads) {
+        SendPublishDirectRead(id, ActorContext());
+    }
+    UnpublishedDirectReads.clear();
+
     ResendRecentRequests();
 }
 
@@ -1310,7 +1330,7 @@ void TPartitionActor::Handle(TEvPersQueue::TEvHasDataInfoResponse::TPtr& ev, con
 
 
 NKikimrClient::TPersQueueRequest TPartitionActor::MakeReadRequest(
-        ui64 readOffset, ui64 lastOffset, ui64 maxCount, ui64 maxSize, ui64 maxTimeLagMs, ui64 readTimestampMs, ui64 directReadId
+        ui64 readOffset, ui64 lastOffset, ui64 maxCount, ui64 maxSize, ui64 maxTimeLagMs, ui64 readTimestampMs, ui64 directReadId, ui64 sizeEstimate
 ) const {
     NKikimrClient::TPersQueueRequest request;
 
@@ -1326,6 +1346,7 @@ NKikimrClient::TPersQueueRequest TPartitionActor::MakeReadRequest(
     read->SetSessionId(Session);
     if (DirectRead) {
         read->SetDirectReadId(directReadId);
+        read->SetSizeEstimate(sizeEstimate);
     }
     if (maxCount) {
         read->SetCount(maxCount);
@@ -1366,7 +1387,7 @@ void TPartitionActor::Handle(TEvPQProxy::TEvRead::TPtr& ev, const TActorContext&
     auto request = MakeReadRequest(ReadOffset, 0, req->MaxCount, req->MaxSize, req->MaxTimeLagMs, req->ReadTimestampMs, DirectReadId);
     RequestInfly = true;
     CurrentRequest = request;
-    
+
     if (!PipeClient) //Pipe will be recreated soon
         return;
 
diff --git a/ydb/services/persqueue_v1/actors/partition_actor.h b/ydb/services/persqueue_v1/actors/partition_actor.h
@@ -145,7 +145,7 @@ class TPartitionActor : public NActors::TActorBootstrapped<TPartitionActor> {
     void InitLockPartition(const NActors::TActorContext& ctx);
     void InitStartReading(const NActors::TActorContext& ctx);
     void RestartDirectReadSession();
-    void OnDirectReadsRestored();
+    void    OnDirectReadsRestored();
     [[nodiscard]] bool SendNextRestorePrepareOrForget();
     [[nodiscard]] bool SendNextRestorePublishRequest();
     void ResendRecentRequests();
@@ -161,7 +161,7 @@ class TPartitionActor : public NActors::TActorBootstrapped<TPartitionActor> {
     NKikimrClient::TPersQueueRequest MakeCreateSessionRequest(bool initial) const;
     NKikimrClient::TPersQueueRequest MakeReadRequest(ui64 readOffset, ui64 lastOffset, ui64 maxCount,
                                                                       ui64 maxSize, ui64 maxTimeLagMs, ui64 readTimestampMs,
-                                                                      ui64 directReadId) const;
+                                                                      ui64 directReadId, ui64 sizeEstimate = 0) const;
 
 
 private:
@@ -246,6 +246,7 @@ class TPartitionActor : public NActors::TActorBootstrapped<TPartitionActor> {
 
     std::map<ui64, NKikimrClient::TPersQueuePartitionResponse::TCmdPrepareDirectReadResult> DirectReadsToRestore;
     std::set<ui64> DirectReadsToPublish;
+    std::set<ui64> UnpublishedDirectReads;
     std::set<ui64> DirectReadsToForget;
 
     enum class EDirectReadRestoreStage {
diff --git a/ydb/services/persqueue_v1/persqueue_ut.cpp b/ydb/services/persqueue_v1/persqueue_ut.cpp

Original file line number	Diff line number	Diff line change
`@@ -1029,6 +1029,9 @@ void BeginCmdRead(const TPQCmdReadSettings& settings, TTestContext& tc)`
`1029`	`1029`	`if (settings.PartitionSessionId > 0) {`
`1030`	`1030`	`read->SetPartitionSessionId(settings.PartitionSessionId);`
`1031`	`1031`	`}`
	`1032`	`+ if (settings.LastOffset > 0) {`
	`1033`	`+ read->SetLastOffset(settings.LastOffset);`
	`1034`	`+ }`
`1032`	`1035`	`if (settings.Pipe) {`
`1033`	`1036`	`ActorIdToProto(settings.Pipe, req->MutablePipeClient());`
`1034`	`1037`	`}`
Original file line number	Diff line number	Diff line change
`@@ -36,6 +36,7 @@ message TPersQueuePartitionRequest {`
`36`	`36`	`optional uint64 PartitionSessionId = 16;`
`37`	`37`	`optional int64 LastOffset = 17;`
`38`	`38`
	`39`	`+ optional uint64 SizeEstimate = 18;`
`39`	`40`	`}`
`40`	`41`
`41`	`42`	`message TCmdPublishDirectRead {`
Original file line number	Diff line number	Diff line change
`@@ -688,6 +688,10 @@ message StreamDirectReadMessage {`
`688`	`688`
`689`	`689`	`// Messages data`
`690`	`690`	`StreamReadMessage.ReadResponse.PartitionData partition_data = 3;`
	`691`	`+`
	`692`	`+ // Total size in bytes of this response as calculated by server.`
	`693`	`+ // See ReadRequest comment above.`
	`694`	`+ int64 bytes_size = 4;`
`691`	`695`	`}`
`692`	`696`	`}`
`693`	`697`