[ICRDMA] Support for RDMA memory region in pdisk, NBYDB-1578 (ydb-platform#27580)

Common RDMA aware RCBUF allocation infrastructure and pdisk support.

Author: dcherednik

ydb/core/blobstorage/base/bufferwithgaps.h

@@ -45,12 +45,16 @@ namespace NKikimr {
             , IsCommited(false)
         {}
 
-        TBufferWithGaps(ui32 offset, ui32 size, ui32 tailroom)
-            : Data(TRcBuf::UninitializedPageAligned(size, tailroom))
+        TBufferWithGaps(ui32 offset, TRcBuf&& data)
+            : Data(std::move(data))
             , Offset(offset)
             , IsCommited(false)
         {}
 
+        TBufferWithGaps(ui32 offset, ui32 size, ui32 tailroom)
+            : TBufferWithGaps(offset, TRcBuf::UninitializedPageAligned(size, tailroom))
+        {}
+
         TBufferWithGaps(TBufferWithGaps &&) = default;
         TBufferWithGaps &operator=(TBufferWithGaps &&) = default;
 

ydb/core/blobstorage/pdisk/blobstorage_pdisk_completion_impl.cpp

@@ -311,7 +311,7 @@ void TCompletionChunkReadPart::Release(TActorSystem *actorSystem) {
 }
 
 TCompletionChunkRead::TCompletionChunkRead(TPDisk *pDisk, TIntrusivePtr<TChunkRead> &read, std::function<void()> onDestroy,
-            ui64 chunkNonce)
+            ui64 chunkNonce, IRcBufAllocator* alloc)
     : TCompletionAction()
     , PDisk(pDisk)
     , Read(read)
@@ -339,7 +339,7 @@ TCompletionChunkRead::TCompletionChunkRead(TPDisk *pDisk, TIntrusivePtr<TChunkRe
         ? read->Size
         : read->Size + read->Offset % sectorSize;
     size_t tailroom = AlignUp<size_t>(newSize, sectorSize) - newSize;
-    CommonBuffer = TBufferWithGaps(read->Offset, newSize, tailroom);
+    CommonBuffer = TBufferWithGaps(read->Offset, alloc->AllocPageAlignedRcBuf(newSize, tailroom));
 }
 
 TCompletionChunkRead::~TCompletionChunkRead() {

ydb/core/blobstorage/pdisk/blobstorage_pdisk_completion_impl.h

@@ -222,7 +222,7 @@ class TCompletionChunkRead : public TCompletionAction {
     const ui64 DoubleFreeCanary;
 public:
     TCompletionChunkRead(TPDisk *pDisk, TIntrusivePtr<TChunkRead> &read, std::function<void()> onDestroy,
-            ui64 chunkNonce);
+            ui64 chunkNonce, IRcBufAllocator* alloc);
     void Exec(TActorSystem *actorSystem) override;
     ~TCompletionChunkRead();
     void ReplyError(TActorSystem *actorSystem, TString reason);

ydb/core/blobstorage/pdisk/blobstorage_pdisk_impl.cpp

@@ -3187,7 +3187,7 @@ bool TPDisk::PreprocessRequest(TRequestBase *request) {
                 --state.OperationsInProgress;
                 --inFlight->ChunkReads;
             };
-            read->FinalCompletion = new TCompletionChunkRead(this, read, std::move(onDestroy), state.Nonce);
+            read->FinalCompletion = new TCompletionChunkRead(this, read, std::move(onDestroy), state.Nonce, PCtx->ActorSystem->GetRcBufAllocator());
 
             static_cast<TChunkRead*>(request)->SelfPointer = std::move(read);
 

ydb/core/blobstorage/pdisk/blobstorage_pdisk_ut.cpp

@@ -10,6 +10,8 @@
 #include <ydb/core/driver_lib/version/ut/ut_helpers.h>
 #include <ydb/core/testlib/actors/test_runtime.h>
 
+#include <ydb/library/actors/interconnect/rdma/mem_pool.h>
+
 #include <util/system/hp_timer.h>
 
 namespace NKikimr {
@@ -1562,9 +1564,10 @@ Y_UNIT_TEST_SUITE(TPDiskTest) {
         return data;
     }
 
-    void ChunkWriteDifferentOffsetAndSizeImpl(bool plainDataChunks) {
+    void ChunkWriteDifferentOffsetAndSizeImpl(bool plainDataChunks, bool rdmaAlloc) {
         TActorTestContext testCtx({
             .PlainDataChunks = plainDataChunks,
+            .UseRdmaAllocator = rdmaAlloc,
         });
         Cerr << "plainDataChunks# " << plainDataChunks << Endl;
 
@@ -1609,8 +1612,8 @@ Y_UNIT_TEST_SUITE(TPDiskTest) {
         }
     }
     Y_UNIT_TEST(ChunkWriteDifferentOffsetAndSize) {
-        for (int i = 0; i <= 1; ++i) {
-            ChunkWriteDifferentOffsetAndSizeImpl(i);
+        for (ui32 i = 0; i <= 3; ++i) {
+            ChunkWriteDifferentOffsetAndSizeImpl(i & 1, i & 2);
         }
     }
 
@@ -3033,4 +3036,86 @@ Y_UNIT_TEST_SUITE(TPDiskPrefailureDiskTest) {
     }
 }
 
+Y_UNIT_TEST_SUITE(RDMA) {
+    void TestChunkReadWithRdmaAllocator(bool plainDataChunks) {
+        TActorTestContext testCtx({
+            .PlainDataChunks = plainDataChunks,
+            .UseRdmaAllocator = true,
+        });
+        TVDiskMock vdisk(&testCtx);
+
+        vdisk.InitFull();
+        vdisk.ReserveChunk();
+        vdisk.CommitReservedChunks();
+        auto chunk = *vdisk.Chunks[EChunkState::COMMITTED].begin();
+
+        for (i64 writeSize: {1_KB - 123, 64_KB, 128_KB + 8765}) {
+            for (i64 readOffset: {0_KB, 0_KB + 123, 4_KB, 8_KB + 123}) {
+                for (i64 readSize: {writeSize, writeSize - 567, writeSize - i64(8_KB), writeSize - i64(8_KB) + 987}) {
+                    if (readOffset < 0 || readSize < 0 || readOffset + readSize > writeSize) {
+                        continue;
+                    }
+                    Cerr << "chunkBufSize: " << writeSize << ", readOffset: " << readOffset << ", readSize: " << readSize << Endl;
+                    auto parts = MakeIntrusive<NPDisk::TEvChunkWrite::TAlignedParts>(PrepareData(writeSize));
+                    testCtx.Send(new NPDisk::TEvChunkWrite(
+                        vdisk.PDiskParams->Owner, vdisk.PDiskParams->OwnerRound,
+                        chunk, 0, parts, nullptr, false, 0));
+                    auto write = testCtx.Recv<NPDisk::TEvChunkWriteResult>();
+
+                    testCtx.Send(new NPDisk::TEvChunkRead(
+                        vdisk.PDiskParams->Owner, vdisk.PDiskParams->OwnerRound,
+                        chunk, readOffset, readSize, 0, nullptr));
+                    auto read = testCtx.Recv<NPDisk::TEvChunkReadResult>();
+                    TRcBuf readBuf = read->Data.ToString();
+                    NInterconnect::NRdma::TMemRegionSlice memReg = NInterconnect::NRdma::TryExtractFromRcBuf(readBuf);
+                    UNIT_ASSERT_C(!memReg.Empty(), "Failed to extract RDMA memory region from RcBuf");
+                    UNIT_ASSERT_VALUES_EQUAL_C(memReg.GetSize(), readSize, "Unexpected size of RDMA memory region");
+                }
+            }
+        }
+    }
+
+    Y_UNIT_TEST(TestChunkReadWithRdmaAllocatorEncryptedChunks) {
+        TestChunkReadWithRdmaAllocator(false);
+    }
+    Y_UNIT_TEST(TestChunkReadWithRdmaAllocatorPlainChunks) {
+        TestChunkReadWithRdmaAllocator(true);
+    }
+
+    Y_UNIT_TEST(TestRcBuf) {
+        ui32 size = 129961;
+        ui32 offset = 123;
+        ui32 tailRoom = 1111;
+        ui32 totalSize = size + tailRoom;
+        UNIT_ASSERT_VALUES_EQUAL(totalSize, 131072);
+
+        auto alloc1 = [](ui32 size, ui32 headRoom, ui32 tailRoom) {
+            TRcBuf buf = TRcBuf::UninitializedPageAligned(size + headRoom + tailRoom);
+            buf.TrimFront(size + tailRoom);
+            buf.TrimBack(size);
+            Cerr << "alloc1: " << buf.Size() << " " << buf.Tailroom() << " " << buf.UnsafeTailroom() << Endl;
+            return buf;
+        };
+        auto memPool = NInterconnect::NRdma::CreateDummyMemPool();
+        auto alloc2 = [memPool](ui32 size, ui32 headRoom, ui32 tailRoom) -> TRcBuf {
+            TRcBuf buf = memPool->AllocRcBuf(size + headRoom + tailRoom, NInterconnect::NRdma::IMemPool::EMPTY).value();
+            buf.TrimFront(size + tailRoom);
+            buf.TrimBack(size);
+            Cerr << "alloc2: " << buf.Size() << " " << buf.Tailroom() << " " << buf.UnsafeTailroom() << Endl;
+            return buf;
+        };
+
+        auto buf1 = TBufferWithGaps(offset, alloc1(size, 0, tailRoom));
+        auto buf2 = TBufferWithGaps(offset, alloc2(size, 0, tailRoom));
+
+        Cerr << "buf1: " << buf1.PrintState() << " " << buf1.Size() << " " << buf1.SizeWithTail() << Endl;
+        Cerr << "buf2: " << buf2.PrintState() << " " << buf2.Size() << " " << buf2.SizeWithTail() << Endl;
+
+        UNIT_ASSERT_VALUES_EQUAL_C(buf1.Size(), buf2.Size(), "Buffers should have the same size");
+
+        buf1.RawDataPtr(0, totalSize);
+        buf2.RawDataPtr(0, totalSize);
+    }
+}
+
 } // namespace NKikimr

ydb/core/blobstorage/pdisk/blobstorage_pdisk_ut_env.h

@@ -35,6 +35,7 @@ struct TActorTestContext {
         bool ReadOnly = false;
         bool InitiallyZeroed = false; // Only for sector map. Zero first 1MiB on start.
         bool PlainDataChunks = false;
+        bool UseRdmaAllocator = false;
     };
 
 private:
@@ -94,7 +95,7 @@ struct TActorTestContext {
     }
 
     TActorTestContext(TSettings settings)
-        : Runtime(new TTestActorRuntime(1, true))
+        : Runtime(new TTestActorRuntime(1, 1, true, settings.UseRdmaAllocator))
         , TestCtx(settings.UseSectorMap, settings.DiskMode, settings.DiskSize, settings.UsePath)
         , Settings(settings)
     {

ydb/core/blobstorage/ut_blobstorage/get.cpp

@@ -1,6 +1,8 @@
 #include <ydb/core/blobstorage/ut_blobstorage/lib/env.h>
 #include <ydb/core/blobstorage/ut_blobstorage/lib/common.h>
 
+#include <ydb/library/actors/interconnect/rdma/mem_pool.h>
+
 Y_UNIT_TEST_SUITE(Get) {
 
     void SendGet(
@@ -108,4 +110,54 @@ Y_UNIT_TEST_SUITE(Get) {
         // check that now TEvGet returns BLOCKED for blocked generation with reader params
         SendGet(test, originalBlobId, data, NKikimrProto::BLOCKED, TEvBlobStorage::TEvGet::TReaderTabletData(tabletId, tabletGeneration));
     }
+
+    Y_UNIT_TEST(TestRdmaRegiseredMemory) {
+        TEnvironmentSetup env(true);
+        TTestInfo test = InitTest(env);
+
+        auto groups = env.GetGroups();
+        auto groupInfo = env.GetGroupInfo(groups.front());
+
+        ui64 tabletId = 1;
+        ui32 tabletGeneration = 1;
+
+        constexpr ui32 size = 1000;
+        TString data(size, 'a');
+        TLogoBlobID originalBlobId(tabletId, tabletGeneration, 0, 0, size, 0);
+
+        SendPut(test, originalBlobId, data, NKikimrProto::OK);
+        env.CompactVDisk(groupInfo->GetActorId(0));
+
+        auto memPool = NInterconnect::NRdma::CreateDummyMemPool();
+        env.Runtime->FilterFunction = [&](ui32, std::unique_ptr<IEventHandle>& ev) {
+            if (ev->GetTypeRewrite() == TEvBlobStorage::EvChunkReadResult) {
+                Cerr << "TEvChunkReadResult: " << ev->ToString() << Endl;
+                auto* res = ev->CastAsLocal<NPDisk::TEvChunkReadResult>();
+                auto buf = res->Data.ToString();
+                auto newBuf = memPool->AllocRcBuf(buf.size(), NInterconnect::NRdma::IMemPool::EMPTY);
+                UNIT_ASSERT(newBuf);
+                std::memcpy(newBuf->UnsafeGetDataMut(), buf.GetData(), buf.GetSize());
+                res->Data = TBufferWithGaps(0, std::move(*newBuf));
+                res->Data.Commit();
+            } else if (ev->GetTypeRewrite() == TEvBlobStorage::EvVGetResult) {
+                Cerr << "TEvVGetResult: " << ev->ToString() << Endl;
+                auto* res = ev->CastAsLocal<TEvBlobStorage::TEvVGetResult>();
+                auto payload = res->GetPayload();
+                UNIT_ASSERT_VALUES_EQUAL(payload.size(), 1);
+                for (auto& rope : payload) {
+                    for (auto it = rope.Begin(); it != rope.End(); ++it) {
+                        const TRcBuf& chunk = it.GetChunk();
+                        auto memReg = NInterconnect::NRdma::TryExtractFromRcBuf(chunk);
+                        UNIT_ASSERT_C(!memReg.Empty(), "unable to extract mem region from chunk");
+                        UNIT_ASSERT_VALUES_EQUAL_C(memReg.GetSize(), size, "invalid size for memReg");
+                        UNIT_ASSERT_C(memReg.GetLKey(0) != 0, "invalid lkey");
+                        UNIT_ASSERT_C(memReg.GetRKey(0) != 0, "invalid rkey");
+                    }
+                }
+            }
+            return true;
+        };
+
+        SendGet(test, originalBlobId, data, NKikimrProto::OK, TEvBlobStorage::TEvGet::TReaderTabletData(tabletId, tabletGeneration));
+    }
 }

ydb/core/testlib/actor_helpers.cpp

@@ -2,10 +2,13 @@
 
 namespace NKikimr {
 
-TActorSystemStub::TActorSystemStub()
+TActorSystemStub::TActorSystemStub(std::shared_ptr<IRcBufAllocator> alloc)
     : AppData(0, 0, 0, 0, {}, nullptr, nullptr, nullptr, nullptr)
 {
     THolder<NActors::TActorSystemSetup> setup(new NActors::TActorSystemSetup);
+    if (alloc) {
+        setup->RcBufAllocator = alloc;
+    }
     System.Reset(new NActors::TActorSystem(setup, &AppData));
     Mailbox.Reset(new NActors::TMailbox());
     ExecutorThread.Reset(new NActors::TExecutorThread(0, System.Get(), nullptr, "thread"));

ydb/core/testlib/actor_helpers.h

@@ -16,7 +16,7 @@ struct TActorSystemStub {
     NActors::TActivationContext* PrevCtx;
     TAppData AppData;
 
-    TActorSystemStub();
+    TActorSystemStub(std::shared_ptr<IRcBufAllocator> alloc = {});
     ~TActorSystemStub();
 };
 

ydb/core/testlib/actors/test_runtime.cpp

@@ -69,9 +69,9 @@ namespace NActors {
         Initialize();
     }
 
-    TTestActorRuntime::TTestActorRuntime(ui32 nodeCount, ui32 dataCenterCount, bool useRealThreads)
+    TTestActorRuntime::TTestActorRuntime(ui32 nodeCount, ui32 dataCenterCount, bool useRealThreads, bool useRdmaAllocator)
         : TPortManager(false)
-        , TTestActorRuntimeBase{nodeCount, dataCenterCount, useRealThreads}
+        , TTestActorRuntimeBase{nodeCount, dataCenterCount, useRealThreads, useRdmaAllocator}
     {
         Initialize();
     }