fixed TPersQueueTest.TopicServiceCommitOffset (#45959)

8 files changed, 261 insertions, 3 deletions

diff --git a/ydb/core/persqueue/pqtablet/partition/partition.cpp b/ydb/core/persqueue/pqtablet/partition/partition.cpp
index 7a3e6438c57..f694d77128a 100644
--- a/ydb/core/persqueue/pqtablet/partition/partition.cpp
+++ b/ydb/core/persqueue/pqtablet/partition/partition.cpp
@@ -3726,6 +3726,7 @@ void TPartition::CommitTransaction(TSimpleSharedPtr<TTransaction>& t)
                 userInfo.Generation = 0;
                 userInfo.Step = 0;
                 userInfo.PipeClient = {};
+                FailStaleSessionReadRequests(operation.GetConsumer(), ActorContext());
             }
         }
 
@@ -4435,6 +4436,7 @@ void TPartition::EmulatePostProcessUserAct(const TEvPQ::TEvSetClientInfo& act,
             userInfo.Generation = 0;
             userInfo.Step = 0;
             userInfo.PipeClient = {};
+            FailStaleSessionReadRequests(user, ctx);
         }
 
         PQ_ENSURE(offset <= (ui64)Max<i64>())("Unexpected Offset", offset);
diff --git a/ydb/core/persqueue/pqtablet/partition/partition.h b/ydb/core/persqueue/pqtablet/partition/partition.h
index e55b62710be..7a491f086c1 100644
--- a/ydb/core/persqueue/pqtablet/partition/partition.h
+++ b/ydb/core/persqueue/pqtablet/partition/partition.h
@@ -305,6 +305,7 @@ private:
     void ProcessChangeOwnerRequests(const TActorContext& ctx);
     void ProcessHasDataRequests(const TActorContext& ctx);
     bool ProcessHasDataRequest(const THasDataReq& request, const TActorContext& ctx);
+    void FailStaleSessionReadRequests(const TString& user, const TActorContext& ctx);
     void ProcessRead(const TActorContext& ctx, TReadInfo&& info, const ui64 cookie, bool subscription);
     void ProcessReserveRequests(const TActorContext& ctx);
     void ProcessTimestampRead(const TActorContext& ctx);
diff --git a/ydb/core/persqueue/pqtablet/partition/partition_read.cpp b/ydb/core/persqueue/pqtablet/partition/partition_read.cpp
index ce8df7113b8..be13bd649aa 100644
--- a/ydb/core/persqueue/pqtablet/partition/partition_read.cpp
+++ b/ydb/core/persqueue/pqtablet/partition/partition_read.cpp
@@ -207,6 +207,37 @@ void TPartition::ProcessHasDataRequests(const TActorContext& ctx) {
     }
 }
 
+void TPartition::FailStaleSessionReadRequests(const TString& user, const TActorContext& ctx) {
+    const auto now = ctx.Now();
+
+    auto forgetSubscription = [&](const TString& clientId) {
+        if (InitDone && !clientId.empty()) {
+            UsersInfoStorage->GetOrCreate(clientId, ctx).ForgetSubscription(GetEndOffset(), now);
+        }
+    };
+
+    for (auto it = HasDataRequests.begin(); it != HasDataRequests.end();) {
+        if (it->ClientId != user) {
+            ++it;
+            continue;
+        }
+        auto response = MakeHasDataInfoResponse(0, it->Cookie);
+        response->Record.SetSessionInvalidated(true);
+        ctx.Send(it->Sender, response.Release());
+
+        forgetSubscription(it->ClientId);
+        it = HasDataRequests.erase(it);
+    }
+
+    for (auto dlIt = HasDataDeadlines.begin(); dlIt != HasDataDeadlines.end();) {
+        if (dlIt->Request.ClientId == user) {
+            dlIt = HasDataDeadlines.erase(dlIt);
+        } else {
+            ++dlIt;
+        }
+    }
+}
+
 void TPartition::Handle(TEvPersQueue::TEvHasDataInfo::TPtr& ev, const TActorContext& ctx) {
     auto& record = ev->Get()->Record;
     PQ_ENSURE(record.HasSender());
@@ -217,6 +248,21 @@ void TPartition::Handle(TEvPersQueue::TEvHasDataInfo::TPtr& ev, const TActorCont
     auto readTimestamp = GetReadFrom(record.GetMaxTimeLagMs(), record.GetReadTimestampMs(), TInstant::Zero(), ctx);
     TActorId sender = ActorIdFromProto(record.GetSender());
 
+    if (InitDone && !record.GetSessionId().empty()) {
+        auto& userInfo = UsersInfoStorage->GetOrCreate(record.GetClientId(), ctx);
+        if (!userInfo.NoConsumer) {
+            // Prefer pending session: CommitTransaction clears it before persist applies to UsersInfoStorage.
+            const TUserInfoBase* pending = GetPendingUserIfExists(record.GetClientId());
+            const TString& session = pending ? pending->Session : userInfo.Session;
+            if (session != record.GetSessionId()) {
+                auto response = MakeHasDataInfoResponse(0, cookie);
+                response->Record.SetSessionInvalidated(true);
+                ctx.Send(sender, response.Release());
+                return;
+            }
+        }
+    }
+
     THasDataReq req{++HasDataReqNum, (ui64)record.GetOffset(), sender, cookie,
         record.HasClientId() && InitDone ? record.GetClientId() : "", readTimestamp};
 
diff --git a/ydb/core/persqueue/public/fetcher/fetch_request_actor.cpp b/ydb/core/persqueue/public/fetcher/fetch_request_actor.cpp
index 05951b7ff2d..c0ca8a02e5b 100644
--- a/ydb/core/persqueue/public/fetcher/fetch_request_actor.cpp
+++ b/ydb/core/persqueue/public/fetcher/fetch_request_actor.cpp
@@ -360,6 +360,12 @@ public:
             return;
         }
 
+        if (record.GetSessionInvalidated()) {
+            AddResult(partitionIndex, EPartitionStatus::DataReceived, NPersQueue::NErrorCode::READ_ERROR_NO_SESSION);
+            ProceedFetchRequest(ctx);
+            return;
+        }
+
         status = EPartitionStatus::HasDataReceived;
 
         auto& partition = Settings.Partitions[partitionIndex];
diff --git a/ydb/core/persqueue/ut/partition_ut.cpp b/ydb/core/persqueue/ut/partition_ut.cpp
index a8cb0d359ac..3730f15b9b7 100644
--- a/ydb/core/persqueue/ut/partition_ut.cpp
+++ b/ydb/core/persqueue/ut/partition_ut.cpp
@@ -318,7 +318,8 @@ protected:
                            const TString& consumer,
                            ui64 begin,
                            ui64 end,
-                           const TActorId& suppPartitionId = {});
+                           const TActorId& suppPartitionId = {},
+                           bool killReadSession = false);
     void WaitCalcPredicateResult(const TCalcPredicateMatcher& matcher = TCalcPredicateMatcher::EmptyMatcher());
 
     void SendCommitTx(ui64 step, ui64 txId, const TSendCommitTxOptions& options = {});
@@ -1132,13 +1133,14 @@ void TPartitionFixture::SendCalcPredicate(ui64 step,
                                           const TString& consumer,
                                           ui64 begin,
                                           ui64 end,
-                                          const TActorId& suppPartitionId)
+                                          const TActorId& suppPartitionId,
+                                          bool killReadSession)
 {
     auto event = MakeHolder<TEvPQ::TEvTxCalcPredicate>(step, txId);
     if (suppPartitionId) {
         event->SupportivePartitionActor = suppPartitionId;
     } else {
-        event->AddOperation(consumer, begin, end);
+        event->AddOperation(consumer, begin, end, false, killReadSession);
     }
 
     Ctx->Runtime->SingleSys()->Send(new IEventHandle(ActorId, Ctx->Edge, event.Release()));
@@ -2383,6 +2385,114 @@ Y_UNIT_TEST_F(CorrectRange_Commit, TPartitionFixture)
     WaitCommitTxDone({.TxId=txId, .Partition=TPartitionId(partition)});
 }
 
+Y_UNIT_TEST_F(KillReadSessionFailsPendingHasData, TPartitionFixture)
+{
+    const TPartitionId partition{0};
+    const ui64 end = 10;
+    const TString client = "client";
+    const TString session = "session";
+    const ui64 step = 12345;
+    const ui64 txId = 67890;
+    const ui64 hasDataCookie = 42;
+
+    CreatePartition({.Partition=partition, .Begin=0, .End=end, .PlanStep=step, .TxId=10000});
+    CreateSession(client, session);
+
+    {
+        auto event = MakeHolder<TEvPersQueue::TEvHasDataInfo>();
+        event->Record.SetPartition(partition.InternalPartitionId);
+        event->Record.SetOffset(end);
+        event->Record.SetDeadline((TInstant::Now() + TDuration::Minutes(1)).MilliSeconds());
+        event->Record.SetCookie(hasDataCookie);
+        event->Record.SetClientId(client);
+        event->Record.SetSessionId(session);
+        ActorIdToProto(Ctx->Edge, event->Record.MutableSender());
+        Ctx->Runtime->SingleSys()->Send(new IEventHandle(ActorId, Ctx->Edge, event.Release()));
+    }
+
+    UNIT_ASSERT_C(
+        !Ctx->Runtime->GrabEdgeEvent<TEvPersQueue::TEvHasDataInfoResponse>(TDuration::MilliSeconds(100)),
+        "HasData must stay pending while offset == EndOffset");
+
+    SendCalcPredicate(step, txId, client, 0, 2, {}, true);
+    WaitCalcPredicateResult({.Step=step, .TxId=txId, .Partition=TPartitionId(partition), .Predicate=true});
+    SendCommitTx(step, txId);
+
+    bool gotSessionInvalidated = false;
+    bool gotCommitDone = false;
+    while (!gotSessionInvalidated || !gotCommitDone) {
+        TAutoPtr<IEventHandle> handle;
+        auto events = Ctx->Runtime->GrabEdgeEvents<
+            TEvPersQueue::TEvHasDataInfoResponse,
+            TEvKeyValue::TEvRequest,
+            TEvPQ::TEvTxDone>(handle, TDuration::Seconds(5));
+
+        if (auto* response = std::get<TEvPersQueue::TEvHasDataInfoResponse*>(events)) {
+            UNIT_ASSERT_VALUES_EQUAL(response->Record.GetCookie(), hasDataCookie);
+            UNIT_ASSERT(response->Record.GetSessionInvalidated());
+            gotSessionInvalidated = true;
+        } else if (std::get<TEvKeyValue::TEvRequest*>(events)) {
+            SendCmdWriteResponse(NMsgBusProxy::MSTATUS_OK);
+        } else if (auto* done = std::get<TEvPQ::TEvTxDone*>(events)) {
+            UNIT_ASSERT_VALUES_EQUAL(done->TxId, txId);
+            gotCommitDone = true;
+        } else {
+            UNIT_FAIL("timeout waiting for SessionInvalidated HasData response and TxDone");
+        }
+    }
+}
+
+// Commit without session clears pending.Session before UsersInfoStorage is updated.
+// HasData with the old SessionId must be rejected via the pending-session check.
+Y_UNIT_TEST_F(HasDataRejectedByPendingSessionAfterCommitWithoutSession, TPartitionFixture)
+{
+    const TPartitionId partition{0};
+    const ui64 end = 10;
+    const TString client = "client";
+    const TString session = "session";
+    const ui64 hasDataCookie = 99;
+
+    CreatePartition({.Partition=partition, .Begin=0, .End=end});
+    CreateSession(client, session);
+
+    // Commit without session id → pending.Session cleared, FailStale, then persist.
+    {
+        auto event = MakeHolder<TEvPQ::TEvSetClientInfo>(
+            /*cookie=*/7, client, /*offset=*/2, /*session=*/"",
+            /*partitionSessionId=*/0, /*gen=*/0, /*step=*/0, TActorId{});
+        event->Strict = true;
+        Ctx->Runtime->SingleSys()->Send(new IEventHandle(ActorId, Ctx->Edge, event.Release()));
+    }
+
+    {
+        auto kv = Ctx->Runtime->GrabEdgeEvent<TEvKeyValue::TEvRequest>(TDuration::Seconds(5));
+        UNIT_ASSERT(kv);
+        // Hold the write response: UsersInfoStorage still has the old session.
+    }
+
+    {
+        auto event = MakeHolder<TEvPersQueue::TEvHasDataInfo>();
+        event->Record.SetPartition(partition.InternalPartitionId);
+        event->Record.SetOffset(end);
+        event->Record.SetDeadline((TInstant::Now() + TDuration::Minutes(1)).MilliSeconds());
+        event->Record.SetCookie(hasDataCookie);
+        event->Record.SetClientId(client);
+        event->Record.SetSessionId(session);
+        ActorIdToProto(Ctx->Edge, event->Record.MutableSender());
+        Ctx->Runtime->SingleSys()->Send(new IEventHandle(ActorId, Ctx->Edge, event.Release()));
+    }
+
+    {
+        auto response = Ctx->Runtime->GrabEdgeEvent<TEvPersQueue::TEvHasDataInfoResponse>(TDuration::Seconds(5));
+        UNIT_ASSERT(response);
+        UNIT_ASSERT_VALUES_EQUAL(response->Record.GetCookie(), hasDataCookie);
+        UNIT_ASSERT(response->Record.GetSessionInvalidated());
+    }
+
+    SendCmdWriteResponse(NMsgBusProxy::MSTATUS_OK);
+    WaitProxyResponse({.Cookie=7, .Status=NMsgBusProxy::MSTATUS_OK});
+}
+
 Y_UNIT_TEST_F(CorrectRange_Multiple_Transactions, TPartitionFixture)
 {
     const TPartitionId partition{3};
diff --git a/ydb/core/protos/pqevents_global.proto b/ydb/core/protos/pqevents_global.proto
index fb8775aad02..17f543f426b 100644
--- a/ydb/core/protos/pqevents_global.proto
+++ b/ydb/core/protos/pqevents_global.proto
@@ -175,6 +175,7 @@ message THasDataInfo {
 
     optional int32 MaxTimeLagMs = 7; // optional, default = infinity, why we use int instead of uint?
     optional uint64 ReadTimestampMs = 8; //optional, default = 0
+    optional string SessionId = 9;
 }
 
 message THasDataInfoResponse { //signal
@@ -185,6 +186,7 @@ message THasDataInfoResponse { //signal
     optional bool ReadingFinished = 5;
     repeated uint32 AdjacentPartitionIds = 6;
     repeated uint32 ChildPartitionIds = 7;
+    optional bool SessionInvalidated = 8;
 }
 
 message TPartitionClientInfo {
diff --git a/ydb/services/persqueue_v1/actors/partition_actor.cpp b/ydb/services/persqueue_v1/actors/partition_actor.cpp
index f18987892cc..37f96af084a 100644
--- a/ydb/services/persqueue_v1/actors/partition_actor.cpp
+++ b/ydb/services/persqueue_v1/actors/partition_actor.cpp
@@ -1482,6 +1482,7 @@ void TPartitionActor::WaitDataInPartition(const TActorContext& ctx) {
     ui64 deadline = (ctx.Now() + WAIT_DATA - WAIT_DELTA).MilliSeconds();
     event->Record.SetDeadline(deadline);
     event->Record.SetClientId(ClientId);
+    event->Record.SetSessionId(Session);
     if (MaxTimeLagMs) {
         event->Record.SetMaxTimeLagMs(MaxTimeLagMs);
     }
@@ -1520,6 +1521,20 @@ void TPartitionActor::Handle(TEvPersQueue::TEvHasDataInfoResponse::TPtr& ev, con
     if (!WaitForData)
         return;
 
+    if (record.GetSessionInvalidated()) {
+        YDB_LOG_DEBUG_CTX(ctx, "Session invalidated while waiting for data, close read session",
+            {"PQLOGPREFIX", PQ_LOG_PREFIX},
+            {"partition", Partition},
+            {"session", Session});
+        WaitForData = false;
+        WaitDataInfly.clear();
+        Counters.Errors.Inc();
+        ctx.Send(ParentId, new TEvPQProxy::TEvCloseSession(
+            TStringBuilder() << "status is not ok: no such session '" << Session << "'",
+            ConvertOldCode(NPersQueue::NErrorCode::READ_ERROR_NO_SESSION)));
+        return;
+    }
+
     if (Counters.WaitsForData) {
         Counters.WaitsForData.Inc();
     }
diff --git a/ydb/services/persqueue_v1/persqueue_ut.cpp b/ydb/services/persqueue_v1/persqueue_ut.cpp
index e9f9f6af54b..53cd9d5ad8b 100644
--- a/ydb/services/persqueue_v1/persqueue_ut.cpp
+++ b/ydb/services/persqueue_v1/persqueue_ut.cpp
@@ -2289,6 +2289,82 @@ Y_UNIT_TEST_SUITE(TPersQueueTest) {
         }
     }
 
+    // CommitOffset while the partition actor is waiting for data via HasData must close the
+    // stream with SESSION_EXPIRED without the client sending another Read.
+    Y_UNIT_TEST(TopicServiceCommitOffsetWhileWaitingForData) {
+        auto server = SetupLbFederationServerWithTopic("topic_wait_hasdata", 1, {"user"});
+        server->EnableLogs({ NKikimrServices::PQ_METACACHE, NKikimrServices::PQ_READ_PROXY });
+        server->EnableLogs({ NKikimrServices::KQP_PROXY }, NLog::EPriority::PRI_EMERG);
+        server->EnableLogs({ NKikimrServices::FLAT_TX_SCHEMESHARD }, NLog::EPriority::PRI_ERROR);
+        server->EnableLogs({ NKikimrServices::PERSQUEUE }, NLog::EPriority::PRI_DEBUG);
+
+        const TString topicPath = "account/topic_wait_hasdata";
+
+        auto Channel_ = grpc::CreateChannel("localhost:" + ToString(server->GrpcPort), grpc::InsecureChannelCredentials());
+        auto TopicStubP_ = Ydb::Topic::V1::TopicService::NewStub(Channel_);
+
+        grpc::ClientContext readContext;
+        readContext.set_deadline(std::chrono::system_clock::now() + std::chrono::seconds(30));
+        auto readStream = TopicStubP_->StreamRead(&readContext);
+        UNIT_ASSERT(readStream);
+
+        // Empty topic: after StartPartitionSession the partition actor enters WaitForData / HasData.
+        {
+            Ydb::Topic::StreamReadMessage::FromClient req;
+            Ydb::Topic::StreamReadMessage::FromServer resp;
+
+            req.mutable_init_request()->add_topics_read_settings()->set_path(topicPath);
+            req.mutable_init_request()->set_consumer("user");
+
+            if (!readStream->Write(req)) {
+                ythrow yexception() << "write fail";
+            }
+            UNIT_ASSERT(readStream->Read(&resp));
+            UNIT_ASSERT(resp.server_message_case() == Ydb::Topic::StreamReadMessage::FromServer::kInitResponse);
+
+            UNIT_ASSERT(readStream->Read(&resp));
+            UNIT_ASSERT(resp.server_message_case() == Ydb::Topic::StreamReadMessage::FromServer::kStartPartitionSessionRequest);
+            UNIT_ASSERT_VALUES_EQUAL(resp.start_partition_session_request().partition_session().path(), topicPath);
+
+            const i64 assignId = resp.start_partition_session_request().partition_session().partition_session_id();
+            req.Clear();
+            req.mutable_start_partition_session_response()->set_partition_session_id(assignId);
+            req.mutable_start_partition_session_response()->set_read_offset(0);
+            if (!readStream->Write(req)) {
+                ythrow yexception() << "write fail";
+            }
+        }
+
+        // Let HasData reach the tablet and stay pending (no data on empty topic).
+        Sleep(TDuration::Seconds(1));
+
+        {
+            Ydb::Topic::CommitOffsetRequest req;
+            Ydb::Topic::CommitOffsetResponse resp;
+
+            req.set_path(topicPath);
+            req.set_consumer("user");
+            req.set_offset(0);
+            grpc::ClientContext rcontext;
+
+            auto status = TopicStubP_->CommitOffset(&rcontext, req, &resp);
+
+            Cerr << resp << "\n";
+            UNIT_ASSERT(status.ok());
+            UNIT_ASSERT_VALUES_EQUAL(resp.operation().status(), Ydb::StatusIds::SUCCESS);
+        }
+
+        {
+            Ydb::Topic::StreamReadMessage::FromServer resp;
+            UNIT_ASSERT(readStream->Read(&resp));
+            Cerr << "=== Got response (expect session expired via HasData invalidate): "
+                 << resp.ShortDebugString() << Endl;
+            UNIT_ASSERT_VALUES_EQUAL(resp.status(), Ydb::StatusIds::SESSION_EXPIRED);
+            UNIT_ASSERT_GE(resp.issues_size(), 1);
+            UNIT_ASSERT_STRING_CONTAINS(resp.issues(0).message(), "no such session");
+        }
+    }
+
     Y_UNIT_TEST(TopicServiceCommitOffsetBadOffsets) {
         auto server = SetupLbFederationServer();
         server->EnableLogs({ NKikimrServices::PQ_METACACHE, NKikimrServices::PQ_READ_PROXY });