[ICRDMA] Handle allocation failure in case of event serialization (#29029)

Author: dcherednik

ydb/library/actors/core/event.h

@@ -12,6 +12,10 @@
 #include <util/system/hp_timer.h>
 #include <util/generic/maybe.h>
 
+namespace NInterconnect::NRdma {
+    class IMemPool;
+}
+
 namespace NActors {
     class TChunkSerializer;
     class IActor;
@@ -39,7 +43,7 @@ namespace NActors {
         }
         virtual ui32 Type() const = 0;
         virtual bool SerializeToArcadiaStream(TChunkSerializer*) const = 0;
-        virtual std::optional<TRope> SerializeToRope(std::function<TRcBuf(ui32 size)>) const {
+        virtual std::optional<TRope> SerializeToRope(NInterconnect::NRdma::IMemPool*) const {
             return std::nullopt;
         }
         virtual bool IsSerializable() const = 0;

ydb/library/actors/core/event_pb.cpp

@@ -333,30 +333,40 @@ namespace NActors {
         return true;
     }
 
-    std::optional<TRope> SerializeToRopeImpl(std::function<TRcBuf(ui32 size)> alloc, const TVector<TRope> &payload) {
+    std::optional<TRope> SerializeToRopeImpl(const google::protobuf::MessageLite& msg, const TVector<TRope>& payload, NInterconnect::NRdma::IMemPool* pool) {
         TRope result;
         auto sz = CalculateSerializedHeaderSizeImpl(payload);
-        if (!sz) {
-            return result;
-        }
-        TRcBuf headerBuf = alloc(sz);
-        if (!headerBuf) {
-            return {};
+        if (sz) {
+            std::optional<TRcBuf> headerBuf = pool->AllocRcBuf(sz, NInterconnect::NRdma::IMemPool::EMPTY);
+            if (!headerBuf) {
+                return {};
+            }
+            char* data = headerBuf->GetDataMut();
+            auto append = [&data](const char *p, size_t len) {
+                std::memcpy(data, p, len);
+                data += len;
+                return true;
+            };
+            SerializeHeaderCommon(payload, append);
+            result.Insert(result.End(), std::move(headerBuf.value()));
+
+            auto appendRope = [&](TRope rope) {
+                result.Insert(result.End(), std::move(rope));
+                return true;
+            };
+            SerializePayloadCommon(payload, appendRope);
         }
-        char* data = headerBuf.GetDataMut();
-        auto append = [&data](const char *p, size_t len) {
-            std::memcpy(data, p, len);
-            data += len;
-            return true;
-        };
-        SerializeHeaderCommon(payload, append);
-        result.Insert(result.End(), headerBuf);
 
-        auto appendRope = [&](TRope rope) {
-            result.Insert(result.End(), std::move(rope));
-            return true;
-        };
-        SerializePayloadCommon(payload, appendRope);
+        {
+            ui32 size = msg.ByteSizeLong();
+            std::optional<TRcBuf> recordsSerializedBuf = pool->AllocRcBuf(size, NInterconnect::NRdma::IMemPool::EMPTY);
+            if (!recordsSerializedBuf) {
+                return {};
+            }
+            bool serializationDone = msg.SerializePartialToArray(recordsSerializedBuf->GetDataMut(), size);
+            Y_ABORT_UNLESS(serializationDone);
+            result.Insert(result.End(), std::move(recordsSerializedBuf.value()));
+        }
 
         return result;
     }

ydb/library/actors/core/event_pb.h

@@ -152,7 +152,7 @@ namespace NActors {
     void ParseExtendedFormatPayload(TRope::TConstIterator &iter, size_t &size, TVector<TRope> &payload, size_t &totalPayloadSize);
     bool SerializeToArcadiaStreamImpl(TChunkSerializer* chunker, const TVector<TRope> &payload);
     ui32 CalculateSerializedHeaderSizeImpl(const TVector<TRope> &payload);
-    std::optional<TRope> SerializeToRopeImpl(std::function<TRcBuf(ui32 size)> alloc, const TVector<TRope> &payload);
+    std::optional<TRope> SerializeToRopeImpl(const google::protobuf::MessageLite& msg, const TVector<TRope>& payload, NInterconnect::NRdma::IMemPool* pool);
     ui32 CalculateSerializedSizeImpl(const TVector<TRope> &payload, ssize_t recordSize);
     TEventSerializationInfo CreateSerializationInfoImpl(size_t preserializedSize, bool allowExternalDataChannel, const TVector<TRope> &payload, ssize_t recordSize);
 
@@ -205,20 +205,8 @@ namespace NActors {
             return CalculateSerializedSizeImpl(Payload, Record.ByteSize());
         }
 
-        std::optional<TRope> SerializeToRope(std::function<TRcBuf(ui32 size)> alloc) const override {
-            std::optional<TRope> result = SerializeToRopeImpl(alloc, Payload);
-            if (!result) {
-                return {};
-            }
-            ui32 size = Record.ByteSizeLong();
-            TRcBuf recordsSerializedBuf = alloc(size);
-            if (!recordsSerializedBuf) {
-                return {};
-            }
-            bool serializationDone = Record.SerializePartialToArray(recordsSerializedBuf.GetDataMut(), size);
-            Y_ABORT_UNLESS(serializationDone);
-            result->Insert(result->End(), std::move(recordsSerializedBuf));
-            return result;
+        std::optional<TRope> SerializeToRope(NInterconnect::NRdma::IMemPool* pool) const override {
+            return NActors::SerializeToRopeImpl(Record, Payload, pool);
         }
 
         static TEv* Load(const TEventSerializedData *input) {

ydb/library/actors/interconnect/interconnect_channel.cpp

@@ -328,9 +328,7 @@ namespace NActors {
         ui32* checksum, ssize_t rdmaDeviceIndex)
     {
         if (!event.Buffer && event.Event) {
-            std::optional<TRope> rope = event.Event->SerializeToRope(
-                [&](ui32 size) -> TRcBuf { return RdmaMemPool->AllocRcBuf(size, NInterconnect::NRdma::IMemPool::EMPTY).value(); }
-            );
+            std::optional<TRope> rope = event.Event->SerializeToRope(RdmaMemPool.get());
             if (!rope) {
                 return false; // serialization failed
             }

ydb/library/actors/interconnect/ut_rdma/rdma_xdc_ut.cpp

@@ -1,5 +1,6 @@
 #include <ydb/library/actors/core/event_pb.h>
 #include <ydb/library/actors/interconnect/rdma/ut/utils/utils.h>
+#include <ydb/library/actors/interconnect/rdma/mem_pool.h>
 
 #include <library/cpp/testing/gtest/gtest.h>
 
@@ -254,10 +255,10 @@ TEST_F(XdcRdmaTest, SerializeToRope) {
         totalXdcSize += len;
     }
 
-    auto allocRcBuf = [](ui32 size) {
-        return TRcBuf::Uninitialized(size);
-    };
-    auto serializedRope = ev->SerializeToRope(allocRcBuf);
+    auto mempool = NInterconnect::NRdma::CreateSlotMemPool(nullptr);
+
+    auto serializedRope = ev->SerializeToRope(mempool.get());
+
     ASSERT_TRUE(serializedRope.has_value());
     auto rope = serializedRope->ConvertToString();
     // 6 1 -120 39 88x5000 8 123 18 11 104 101 108 108 111 32 119 111 114 108 100