ydb_topic: split write messages on different codecs (#3428)

Author: qyryq

ydb/public/sdk/cpp/client/ydb_topic/impl/write_session_impl.cpp

@@ -1192,6 +1192,11 @@ size_t TWriteSessionImpl::WriteBatchImpl() {
         TBlock block{};
         for (; block.OriginalSize < MaxBlockSize && i != CurrentBatch.Messages.size(); ++i) {
             auto& currMessage = CurrentBatch.Messages[i];
+
+            // If MaxBlockSize or MaxBlockMessageCount values are ever changed from infinity and 1 correspondingly,
+            // create a new block, if the existing one is non-empty AND (adding another message will overflow it OR
+            //                                                           its codec is different from the codec of the next message).
+
             auto id = currMessage.Id;
             auto createTs = currMessage.CreatedAt;
 
@@ -1283,22 +1288,25 @@ void TWriteSessionImpl::UpdateTokenIfNeededImpl() {
 void TWriteSessionImpl::SendImpl() {
     Y_ABORT_UNLESS(Lock.IsLocked());
 
-    // External cycle splits ready blocks into multiple gRPC messages. Current gRPC message size hard limit is 64MiB
-    while(IsReadyToSendNextImpl()) {
+    // External cycle splits ready blocks into multiple gRPC messages. Current gRPC message size hard limit is 64MiB.
+    while (IsReadyToSendNextImpl()) {
         TClientMessage clientMessage;
         auto* writeRequest = clientMessage.mutable_write_request();
-
-        // Sent blocks while we can without messages reordering
+        ui32 prevCodec = 0;
+        // Send blocks while we can without messages reordering.
         while (IsReadyToSendNextImpl() && clientMessage.ByteSizeLong() < GetMaxGrpcMessageSize()) {
             const auto& block = PackedMessagesToSend.top();
             Y_ABORT_UNLESS(block.Valid);
+            if (writeRequest->messages_size() > 0 && prevCodec != block.CodecID) {
+                break;
+            }
+            prevCodec = block.CodecID;
             writeRequest->set_codec(static_cast<i32>(block.CodecID));
             Y_ABORT_UNLESS(block.MessageCount == 1);
             for (size_t i = 0; i != block.MessageCount; ++i) {
                 Y_ABORT_UNLESS(!OriginalMessagesToSend.empty());
 
                 auto& message = OriginalMessagesToSend.front();
-
                 auto* msgData = writeRequest->add_messages();
 
                 if (message.Tx) {
@@ -1309,27 +1317,24 @@ void TWriteSessionImpl::SendImpl() {
                 msgData->set_seq_no(GetSeqNoImpl(message.Id));
                 *msgData->mutable_created_at() = ::google::protobuf::util::TimeUtil::MillisecondsToTimestamp(message.CreatedAt.MilliSeconds());
 
-                if (!message.MessageMeta.empty()) {
-                    for (auto& [k, v] : message.MessageMeta) {
-                        auto* pair = msgData->add_metadata_items();
-                        pair->set_key(k);
-                        pair->set_value(v);
-                    }
+                for (auto& [k, v] : message.MessageMeta) {
+                    auto* pair = msgData->add_metadata_items();
+                    pair->set_key(k);
+                    pair->set_value(v);
                 }
                 SentOriginalMessages.emplace(std::move(message));
                 OriginalMessagesToSend.pop();
 
                 msgData->set_uncompressed_size(block.OriginalSize);
-                if (block.Compressed)
+                if (block.Compressed) {
                     msgData->set_data(block.Data.data(), block.Data.size());
-                else {
+                } else {
                     for (auto& buffer: block.OriginalDataRefs) {
                         msgData->set_data(buffer.data(), buffer.size());
                     }
                 }
             }
 
-
             TBlock moveBlock;
             moveBlock.Move(block);
             SentPackedMessage.emplace(std::move(moveBlock));

ydb/public/sdk/cpp/client/ydb_topic/ut/basic_usage_ut.cpp

@@ -698,6 +698,115 @@ Y_UNIT_TEST_SUITE(BasicUsage) {
         UNIT_ASSERT_VALUES_EQUAL(stats->GetEndOffset(), count);
 
     }
+
+    Y_UNIT_TEST(TWriteSession_WriteEncoded) {
+        // This test was adapted from ydb_persqueue tests.
+        // It writes 4 messages: 2 with default codec, 1 with explicitly set GZIP codec, 1 with RAW codec.
+        // The last message MUST be sent in a separate WriteRequest, as it has a codec field applied for all messages in the request.
+        // This separation currently happens in TWriteSessionImpl::SendImpl method.
+
+        auto setup = std::make_shared<TTopicSdkTestSetup>(TEST_CASE_NAME);
+        auto client = setup->MakeClient();
+        auto settings = TWriteSessionSettings()
+            .Path(TEST_TOPIC)
+            .MessageGroupId(TEST_MESSAGE_GROUP_ID);
+
+        size_t batchSize = 100000000;
+        settings.BatchFlushInterval(TDuration::Seconds(1000)); // Batch on size, not on time.
+        settings.BatchFlushSizeBytes(batchSize);
+        auto writer = client.CreateWriteSession(settings);
+        TString message = "message";
+        TString packed;
+        {
+            TStringOutput so(packed);
+            TZLibCompress oss(&so, ZLib::GZip, 6);
+            oss << message;
+        }
+
+        Cerr << message << " " << packed << "\n";
+
+        {
+            auto event = *writer->GetEvent(true);
+            UNIT_ASSERT(!writer->WaitEvent().Wait(TDuration::Seconds(1)));
+            auto ev = writer->WaitEvent();
+            UNIT_ASSERT(std::holds_alternative<TWriteSessionEvent::TReadyToAcceptEvent>(event));
+            auto continueToken = std::move(std::get<TWriteSessionEvent::TReadyToAcceptEvent>(event).ContinuationToken);
+            writer->Write(std::move(continueToken), message);
+            UNIT_ASSERT(ev.Wait(TDuration::Seconds(1)));
+        }
+        {
+            auto event = *writer->GetEvent(true);
+            UNIT_ASSERT(std::holds_alternative<TWriteSessionEvent::TReadyToAcceptEvent>(event));
+            auto continueToken = std::move(std::get<TWriteSessionEvent::TReadyToAcceptEvent>(event).ContinuationToken);
+            writer->Write(std::move(continueToken), "");
+        }
+        {
+            auto event = *writer->GetEvent(true);
+            UNIT_ASSERT(std::holds_alternative<TWriteSessionEvent::TReadyToAcceptEvent>(event));
+            auto continueToken = std::move(std::get<TWriteSessionEvent::TReadyToAcceptEvent>(event).ContinuationToken);
+            writer->WriteEncoded(std::move(continueToken), packed, ECodec::GZIP, message.size());
+        }
+
+        ui32 acks = 0, tokens = 0;
+        while(acks < 4 || tokens < 2)  {
+            auto event = *writer->GetEvent(true);
+            if (std::holds_alternative<TWriteSessionEvent::TAcksEvent>(event)) {
+                acks += std::get<TWriteSessionEvent::TAcksEvent>(event).Acks.size();
+            }
+            if (std::holds_alternative<TWriteSessionEvent::TReadyToAcceptEvent>(event)) {
+                if (tokens == 0) {
+                    auto continueToken = std::move(std::get<TWriteSessionEvent::TReadyToAcceptEvent>(event).ContinuationToken);
+                    writer->WriteEncoded(std::move(continueToken), "", ECodec::RAW, 0);
+                }
+                ++tokens;
+            }
+            Cerr << "GOT EVENT " << acks << " " << tokens << "\n";
+        }
+        UNIT_ASSERT(!writer->WaitEvent().Wait(TDuration::Seconds(5)));
+
+        UNIT_ASSERT_VALUES_EQUAL(acks, 4);
+        UNIT_ASSERT_VALUES_EQUAL(tokens, 2);
+
+        auto readSettings = TReadSessionSettings()
+            .ConsumerName(TEST_CONSUMER)
+            .AppendTopics(TEST_TOPIC);
+        std::shared_ptr<IReadSession> readSession = client.CreateReadSession(readSettings);
+        ui32 readMessageCount = 0;
+        while (readMessageCount < 4) {
+            Cerr << "Get event on client\n";
+            auto event = *readSession->GetEvent(true);
+            std::visit(TOverloaded {
+                [&](TReadSessionEvent::TDataReceivedEvent& event) {
+                    for (auto& message: event.GetMessages()) {
+                        TString sourceId = message.GetMessageGroupId();
+                        ui32 seqNo = message.GetSeqNo();
+                        UNIT_ASSERT_VALUES_EQUAL(readMessageCount + 1, seqNo);
+                        ++readMessageCount;
+                        UNIT_ASSERT_VALUES_EQUAL(message.GetData(), (seqNo % 2) == 1 ? "message" : "");
+                    }
+                },
+                [&](TReadSessionEvent::TCommitOffsetAcknowledgementEvent&) {
+                    UNIT_FAIL("no commits in test");
+                },
+                [&](TReadSessionEvent::TStartPartitionSessionEvent& event) {
+                    event.Confirm();
+                },
+                [&](TReadSessionEvent::TStopPartitionSessionEvent& event) {
+                    event.Confirm();
+                },
+                [&](TReadSessionEvent::TPartitionSessionStatusEvent&) {
+                    UNIT_FAIL("Test does not support lock sessions yet");
+                },
+                [&](TReadSessionEvent::TPartitionSessionClosedEvent&) {
+                    UNIT_FAIL("Test does not support lock sessions yet");
+                },
+                [&](TSessionClosedEvent&) {
+                    UNIT_FAIL("Session closed");
+                }
+
+            }, event);
+        }
+    }
 } // Y_UNIT_TEST_SUITE(BasicUsage)
 
 Y_UNIT_TEST_SUITE(TSettingsValidation) {