Add TMessage::GetBrokenData method (#23656)

Author: qyryq

.github/last_commit.txt

@@ -1 +1 @@
-66e1273f270f166cf9a11288b0b2310f55e6de82
+af4a59730e62ce87813dd141cab505ff94c465d1

include/ydb-cpp-sdk/client/topic/read_events.h

@@ -102,6 +102,7 @@ struct TReadSessionEvent {
             virtual ~TMessageBase() = default;
 
             virtual const std::string& GetData() const;
+            virtual const std::string& GetBrokenData() const;
 
             virtual void Commit() = 0;
 
@@ -144,6 +145,8 @@ struct TReadSessionEvent {
             //! User data.
             //! Throws decompressor exception if decompression failed.
             const std::string& GetData() const override;
+            //! Throws exception if decompression succeeded.
+            const std::string& GetBrokenData() const override;
 
             //! Commits single message.
             void Commit() override;

src/client/topic/impl/read_session_event.cpp

@@ -88,6 +88,10 @@ const std::string& TMessageBase::GetData() const {
     return Data;
 }
 
+const std::string& TMessageBase::GetBrokenData() const {
+    return Data;
+}
+
 uint64_t TMessageBase::GetOffset() const {
     return Information.Offset;
 }
@@ -175,6 +179,13 @@ const std::string& TMessage::GetData() const {
     return TMessageBase::GetData();
 }
 
+const std::string& TMessage::GetBrokenData() const {
+    if (DecompressionException) {
+        return TMessageBase::GetData();
+    }
+    ythrow yexception() << "Can not get broken data after successful decompression";
+}
+
 bool TMessage::HasException() const {
     return DecompressionException != nullptr;
 }

src/client/topic/impl/read_session_impl.ipp

@@ -1362,7 +1362,6 @@ inline void TSingleClusterReadSessionImpl<false>::StopPartitionSessionImpl(
 
     if (graceful) {
         auto committedOffset = partitionStream->GetMaxCommittedOffset();
-        LOG_LAZY(Log, TLOG_DEBUG, GetLogPrefix() << "XXXXX PushEvent 1422 TStopPartitionSessionEvent");
         pushRes = EventsQueue->PushEvent(
             partitionStream,
             // TODO(qyryq) Is it safe to use GetMaxCommittedOffset here instead of StopPartitionSessionRequest.commmitted_offset?
@@ -1375,7 +1374,6 @@ inline void TSingleClusterReadSessionImpl<false>::StopPartitionSessionImpl(
         released.set_partition_session_id(partitionStream->GetAssignId());
         WriteToProcessorImpl(std::move(req));
         PartitionStreams.erase(partitionSessionId);
-        LOG_LAZY(Log, TLOG_DEBUG, GetLogPrefix() << "XXXXX PushEvent 1435 TPartitionSessionClosedEvent");
         pushRes = EventsQueue->PushEvent(
             partitionStream,
             TReadSessionEvent::TPartitionSessionClosedEvent(partitionStream, TReadSessionEvent::TPartitionSessionClosedEvent::EReason::Lost),
@@ -3064,16 +3062,15 @@ void TDataDecompressionInfo<UseMigrationProtocol>::TDecompressionTask::operator(
                         data.set_data(TStringType{decompressed});
                     }
                 }
-
-                DecompressedSize += data.data().size();
             } catch (...) {
                 parent->PutDecompressionError(std::current_exception(), messages.Batch, i);
-                data.clear_data(); // Free memory, because we don't count it.
 
                 if (auto session = parent->CbContext->LockShared()) {
                     session->GetLog() << TLOG_INFO << "Error decompressing data: " << CurrentExceptionMessage();
                 }
             }
+
+            DecompressedSize += data.data().size();
         }
     }
 

tests/integration/topic/basic_usage.cpp

@@ -747,6 +747,113 @@ TEST_F(BasicUsage, TEST_NAME(TWriteSession_WriteEncoded)) {
     }
 }
 
+TEST_F(BasicUsage, TEST_NAME(TWriteSession_WriteEncoded_Broken)) {
+    // Write a broken compressed message.
+    // GetData should throw an exception.
+    // GetBrokenData should return the broken data.
+
+    // Write a correct compressed message.
+    // GetData should return the correct data.
+    // GetBrokenData should throw an exception.
+
+    auto driver = MakeDriver();
+
+    TTopicClient client(driver);
+
+    auto settings = TWriteSessionSettings()
+        .Path(GetTopicPath())
+        .MessageGroupId(TEST_MESSAGE_GROUP_ID);
+
+    auto writer = client.CreateWriteSession(settings);
+    std::string brokenPacked = "some broken data";
+
+    {
+        auto event = *writer->GetEvent(true);
+        ASSERT_TRUE(std::holds_alternative<TWriteSessionEvent::TReadyToAcceptEvent>(event));
+        writer->WriteEncoded(
+            std::move(std::get<TWriteSessionEvent::TReadyToAcceptEvent>(event).ContinuationToken),
+            TWriteMessage::CompressedMessage(brokenPacked, ECodec::GZIP, 100)
+        );
+    }
+
+    std::string correctMessage = "message";
+    TString packed;
+    {
+        TStringOutput so(packed);
+        TZLibCompress oss(&so, ZLib::GZip, 6);
+        oss << correctMessage;
+    }
+    {
+        auto event = *writer->GetEvent(true);
+        ASSERT_TRUE(std::holds_alternative<TWriteSessionEvent::TReadyToAcceptEvent>(event));
+        writer->WriteEncoded(
+            std::move(std::get<TWriteSessionEvent::TReadyToAcceptEvent>(event).ContinuationToken),
+            TWriteMessage::CompressedMessage(packed, ECodec::GZIP, correctMessage.size())
+        );
+    }
+
+    std::uint32_t acks = 0;
+    while (acks < 2)  {
+        auto event = *writer->GetEvent(true);
+        if (auto e = std::get_if<TWriteSessionEvent::TAcksEvent>(&event)) {
+            acks += e->Acks.size();
+        } else {
+            continue;
+        }
+    }
+
+    ASSERT_EQ(acks, 2u);
+
+    auto readSettings = TReadSessionSettings()
+        .ConsumerName(GetConsumerName())
+        .AppendTopics(GetTopicPath())
+        // .DirectRead(EnableDirectRead)
+        ;
+    std::shared_ptr<IReadSession> readSession = client.CreateReadSession(readSettings);
+    std::uint32_t readMessageCount = 0;
+    while (readMessageCount < 2) {
+        std::cerr << "Get event on client\n";
+        auto event = *readSession->GetEvent(true);
+        std::visit(TOverloaded {
+            [&](TReadSessionEvent::TDataReceivedEvent& event) {
+                for (auto& message: event.GetMessages()) {
+                    if (readMessageCount == 0) {
+                        ASSERT_THROW(message.GetData(), std::exception);
+                        std::string data = message.GetBrokenData();
+                        ASSERT_TRUE(brokenPacked == data);
+                    } else {
+                        ASSERT_THROW(message.GetBrokenData(), std::exception);
+                        std::string data = message.GetData();
+                        ASSERT_TRUE(correctMessage == data);
+                    }
+                    ++readMessageCount;
+                }
+            },
+            [&](TReadSessionEvent::TCommitOffsetAcknowledgementEvent&) {
+                FAIL();
+            },
+            [&](TReadSessionEvent::TStartPartitionSessionEvent& event) {
+                event.Confirm();
+            },
+            [&](TReadSessionEvent::TStopPartitionSessionEvent& event) {
+                event.Confirm();
+            },
+            [&](TReadSessionEvent::TEndPartitionSessionEvent& event) {
+                event.Confirm();
+            },
+            [&](TReadSessionEvent::TPartitionSessionStatusEvent&) {
+                FAIL() << "Test does not support lock sessions yet";
+            },
+            [&](TReadSessionEvent::TPartitionSessionClosedEvent&) {
+                FAIL() << "Test does not support lock sessions yet";
+            },
+            [&](TSessionClosedEvent&) {
+                FAIL() << "Session closed";
+            }
+        }, event);
+    }
+}
+
 namespace {
     enum class EExpectedTestResult {
         SUCCESS,