fix buffered stream interleaving for proto format

commit_hash:66e6b75f97ea253ddbf4c97a180f9f971c59f93b

Author: ermolovd

yt/cpp/mapreduce/io/proto_table_writer.cpp

@@ -231,4 +231,21 @@ void TLenvalProtoSingleTableWriter::AddRow(const Message& row, size_t tableIndex
 
 ////////////////////////////////////////////////////////////////////////////////
 
+void LenvalEncodeProto(IZeroCopyOutput* output, const ::google::protobuf::Message& message)
+{
+    i32 size = message.ByteSizeLong();
+    output->Write(&size, sizeof(size));
+
+    TProtobufZeroCopyOutputStream adapter(output);
+    auto result = message.SerializeToZeroCopyStream(&adapter);
+
+    adapter.ThrowOnError();
+
+    if (!result) {
+        ythrow yexception() << "Failed to serialize protobuf message";
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 } // namespace NYT

yt/cpp/mapreduce/io/proto_table_writer.h

@@ -79,4 +79,8 @@ class TLenvalProtoSingleTableWriter
 
 ////////////////////////////////////////////////////////////////////////////////
 
+void LenvalEncodeProto(IZeroCopyOutput* output, const ::google::protobuf::Message& message);
+
+////////////////////////////////////////////////////////////////////////////////
+
 } // namespace NYT

yt/yt/core/misc/protobuf_helpers.cpp

@@ -617,6 +617,46 @@ TProtobufOutputStreamAdaptor::TProtobufOutputStreamAdaptor(IOutputStream* stream
     , CopyingOutputStreamAdaptor(this)
 { }
 
+TProtobufZeroCopyOutputStream::TProtobufZeroCopyOutputStream(IZeroCopyOutput* stream)
+    : Stream_(stream)
+{ }
+
+bool TProtobufZeroCopyOutputStream::Next(void** data, int* size)
+{
+    try {
+        size_t sizetSize = Stream_->Next(data);
+        constexpr int maxSize = std::numeric_limits<int>::max();
+        if (sizetSize > maxSize) {
+            Stream_->Undo(sizetSize - maxSize);
+            sizetSize = maxSize;
+        }
+        *size = sizetSize;
+    } catch (const std::exception&) {
+        Error_ = std::current_exception();
+        return false;
+    }
+    ByteCount_ += *size;
+    return true;
+}
+
+void TProtobufZeroCopyOutputStream::BackUp(int count)
+{
+    ByteCount_ -= count;
+    Stream_->Undo(count);
+}
+
+int64_t TProtobufZeroCopyOutputStream::ByteCount() const
+{
+    return ByteCount_;
+}
+
+void TProtobufZeroCopyOutputStream::ThrowOnError() const
+{
+    if (Error_) {
+        std::rethrow_exception(Error_);
+    }
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace NYT

yt/yt/core/misc/protobuf_helpers.h

@@ -548,6 +548,24 @@ class TProtobufOutputStreamAdaptor
     explicit TProtobufOutputStreamAdaptor(IOutputStream* stream);
 };
 
+class TProtobufZeroCopyOutputStream
+    : public ::google::protobuf::io::ZeroCopyOutputStream
+{
+public:
+    explicit TProtobufZeroCopyOutputStream(IZeroCopyOutput* stream);
+
+    bool Next(void** data, int* size) override;
+    void BackUp(int count) override;
+    int64_t ByteCount() const override;
+
+    void ThrowOnError() const;
+
+private:
+    IZeroCopyOutput* const Stream_;
+    std::exception_ptr Error_;
+    int64_t ByteCount_ = 0;
+};
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace NYT