fix request resources deallocation race

cloud/storage/core/libs/rdma/iface/protobuf.cpp

@@ -191,6 +191,19 @@ size_t SerializeError(ui32 code, TStringBuf message, TStringBuf buffer)
     return 0;   // will be interpreted as E_FAIL by ParseError
 }
 
+TString SerializeError(ui32 code, TStringBuf message)
+{
+    auto error = MakeError(code, TString(message));
+
+    TString result;
+    result.resize(error.ByteSizeLong());
+
+    bool succeeded = error.SerializeToArray(result.Detach(), result.size());
+    Y_ENSURE(succeeded, "could not serialize protobuf message");
+
+    return result;
+}
+
 NProto::TError ParseError(TStringBuf buffer)
 {
     NProto::TError error;

cloud/storage/core/libs/rdma/iface/protobuf.h

@@ -101,6 +101,7 @@ class TProtoMessageSerializer
 ////////////////////////////////////////////////////////////////////////////////
 
 size_t SerializeError(ui32 code, TStringBuf message, TStringBuf buffer);
+TString SerializeError(ui32 code, TStringBuf message);
 
 NProto::TError ParseError(TStringBuf buffer);
 

cloud/storage/core/libs/rdma/impl/client.cpp

@@ -61,6 +61,12 @@ constexpr TDuration INSTANT_RECONNECT_DELAY = TDuration::MicroSeconds(1);
 
 ////////////////////////////////////////////////////////////////////////////////
 
+const TString StaleGenerationError = SerializeError(
+    E_RDMA_UNAVAILABLE,
+    "buffer generation changed");
+
+////////////////////////////////////////////////////////////////////////////////
+
 struct TRequest;
 using TRequestPtr = std::unique_ptr<TRequest>;
 
@@ -164,6 +170,8 @@ struct TRequest
     NVerbs::TMemoryWindowPtr InMemoryWindow = NVerbs::NullPtr;
     NVerbs::TMemoryWindowPtr OutMemoryWindow = NVerbs::NullPtr;
 
+    ui64 BufferPoolGeneration = 0;
+
     ERequestState State = ERequestState::Init;
 
     TRequest(
@@ -597,6 +605,7 @@ class TClientEndpoint final
 
     TBufferPool SendBuffers;
     TBufferPool RecvBuffers;
+    std::atomic<ui64> BufferPoolGeneration = 0;
     TMutex AllocationLock;
 
     TPooledBuffer SendBuffer {};
@@ -713,6 +722,7 @@ class TClientEndpoint final
     void InvalidateBuffers(TRequest* req) noexcept;
     void CompleteRequest(ui32 reqId) noexcept;
     void AbortRequest(TRequestPtr req, ui32 err, const TString& msg) noexcept;
+    void RejectStaleGenerationRequest(TRequestPtr req) noexcept;
     void FreeRequest(TRequest* creq) noexcept;
     bool PostSend(TRequest* req, TSendWr* send, ibv_send_wr* wr) noexcept;
     void HandleSendError(TSendWr* send) noexcept;
@@ -853,14 +863,17 @@ void TClientEndpoint::CreateQP()
         recvFlags |= IBV_ACCESS_MW_BIND;
     }
 
-    SendBuffers.Init(Verbs, Connection->pd, sendFlags);
-    RecvBuffers.Init(Verbs, Connection->pd, recvFlags);
+    with_lock (AllocationLock) {
+        ++BufferPoolGeneration;
+        SendBuffers.Init(Verbs, Connection->pd, sendFlags);
+        RecvBuffers.Init(Verbs, Connection->pd, recvFlags);
 
-    SendBuffer = SendBuffers.AcquireBuffer(
-        Config.SendQueueSize * sizeof(TRequestMessage), true);
+        SendBuffer = SendBuffers.AcquireBuffer(
+            Config.SendQueueSize * sizeof(TRequestMessage), true);
 
-    RecvBuffer = RecvBuffers.AcquireBuffer(
-        Config.RecvQueueSize * sizeof(TResponseMessage), true);
+        RecvBuffer = RecvBuffers.AcquireBuffer(
+            Config.RecvQueueSize * sizeof(TResponseMessage), true);
+    }
 
     SendWrs.resize(Config.SendQueueSize);
     RecvWrs.resize(Config.RecvQueueSize);
@@ -990,6 +1003,8 @@ TResultOrError<TClientRequestPtr> TClientEndpoint::AllocateRequest(
 
     try {
         with_lock (AllocationLock) {
+            req->BufferPoolGeneration = BufferPoolGeneration.load();
+
             if (requestBytes) {
                 req->InBuffer = SendBuffers.AcquireBuffer(requestBytes);
             }
@@ -1034,6 +1049,12 @@ ui64 TClientEndpoint::SendRequest(
     auto clientReqId = GetNewReqId();
     req->ClientReqId = clientReqId;
 
+    if (req->BufferPoolGeneration != BufferPoolGeneration.load()) {
+        // Endpoint reconnected between AllocateRequest() and SendRequest()
+        RejectStaleGenerationRequest(std::move(req));
+        return clientReqId;
+    }
+
     if (!CheckState(EEndpointState::Connected)) {
         AbortRequest(std::move(req), E_RDMA_UNAVAILABLE, "endpoint is unavailable");
         return clientReqId;
@@ -1055,6 +1076,20 @@ ui64 TClientEndpoint::SendRequest(
     return clientReqId;
 }
 
+void TClientEndpoint::RejectStaleGenerationRequest(TRequestPtr req) noexcept
+{
+    req->RequestBuffer = {};
+    req->ResponseBuffer = TStringBuf(
+        StaleGenerationError.data(),
+        StaleGenerationError.length());
+
+    auto* handler = req->Handler.get();
+    handler->HandleResponse(
+        std::move(req),
+        RDMA_PROTO_FAIL,
+        StaleGenerationError.length());
+}
+
 bool TClientEndpoint::HandleInputRequests() noexcept
 {
     if (WaitMode == EWaitMode::Poll) {
@@ -1100,6 +1135,16 @@ void TClientEndpoint::HandleQueuedRequests() noexcept
         Y_ABORT_UNLESS(req);
         Counters->RequestDequeued();
 
+        if (req->BufferPoolGeneration != BufferPoolGeneration.load()) {
+            // The caller's thread could have been preempted between the
+            // generation check in SendRequest() and InputRequests.Enqueue()
+            // for the whole duration of a reconnect cycle - req->InBuffer/
+            // OutBuffer may already belong to a torn down pool generation.
+            SendQueue.Push(send);
+            RejectStaleGenerationRequest(std::move(req));
+            continue;
+        }
+
         if (req->State != ERequestState::Enqueued) {
             RDMA_ERROR(
                 "request " << req->ReqId << " has unexpected state "
@@ -1265,8 +1310,10 @@ void TClientEndpoint::AbortRequest(
     ui32 err,
     const TString& msg) noexcept
 {
-    // destroying memory window automatically invalidates it. this ensures no
-    // remote write can succeed after this point
+    if (req->InMemoryWindow) {
+        req->InMemoryWindow.reset();
+        Counters->ReleaseMemoryWindow();
+    }
     if (req->OutMemoryWindow) {
         req->OutMemoryWindow.reset();
         Counters->ReleaseMemoryWindow();
@@ -1957,6 +2004,17 @@ bool TClientEndpoint::FlushHanging() const
 void TClientEndpoint::FreeRequest(TRequest* req) noexcept
 {
     with_lock (AllocationLock) {
+        const bool sameBufferPoolGeneration =
+            req->BufferPoolGeneration == BufferPoolGeneration.load();
+
+        if (!sameBufferPoolGeneration) {
+            // Late request destruction after reconnect: InBuffer/OutBuffer
+            // were acquired eagerly in AllocateRequest() and belong to an
+            // older pool generation, so they cannot be safely released via
+            // the current SendBuffers/RecvBuffers.
+            return;
+        }
+
         // destroy memory windows that haven't been properly invalidated
         if (req->InMemoryWindow) {
             req->InMemoryWindow.reset();

cloud/storage/core/libs/rdma/impl/client_ut.cpp

@@ -1600,4 +1600,178 @@ TEST(TRdmaClientTest, ShouldBindAndInvalidateBuffers)
         }
 }
 
+TEST(TRdmaClientTest, ShouldRejectRequestFromStaleBufferPoolGeneration)
+{
+        auto testContext = MakeIntrusive<NVerbs::TTestContext>();
+        testContext->AllowConnect = true;
+
+        auto verbs = NVerbs::CreateTestVerbs(testContext);
+        auto monitoring = CreateMonitoringServiceStub();
+        auto clientConfig = std::make_shared<TClientConfig>();
+
+        auto logging = CreateLoggingService(
+            "console",
+            TLogSettings{TLOG_RESOURCES});
+
+        auto client = CreateTestClient(
+            verbs,
+            logging,
+            monitoring,
+            clientConfig);
+
+        client->Start();
+        Y_DEFER {
+            client->Stop();
+        };
+
+        auto endpoint = client->StartEndpoint("::", 10020).ExtractValueSync();
+
+        testContext->PostSend = [&](ibv_qp* qp, ibv_send_wr* wr) {
+            PostSend<TRequestMessage>(testContext, qp, wr);
+        };
+
+        size_t requestBytes = 1024;
+        size_t responseBytes = 1024;
+
+        struct TResponse
+        {
+            bool Received = false;
+            ui32 Status = 0;
+            size_t Bytes = 0;
+            TString Buffer;
+        };
+
+        TResponse response;
+
+        auto ctx = std::make_unique<TRequestContext>();
+        ctx->Handler = [&](TStringBuf requestBuffer,
+            TStringBuf responseBuffer,
+            ui32 status,
+            size_t bytes)
+        {
+            Y_UNUSED(requestBuffer);
+            response =
+                TResponse{
+                    .Received = true,
+                    .Status = status,
+                    .Bytes = bytes,
+                    .Buffer = TString(responseBuffer.Head(bytes))};
+        };
+
+        // allocate a request (InBuffer/OutBuffer) against the current
+        // generation, but do not send it yet
+        auto request = endpoint->AllocateRequest(
+            std::make_shared<TClientHandler>(),
+            std::move(ctx),
+            requestBytes,
+            responseBytes);
+        ASSERT_FALSE(HasError(request.GetError()));
+
+        Disconnect(testContext);
+
+        ui64 recv;
+        do {
+            recv = AtomicGet(testContext->PostRecvCounter);
+        } while (recv != 2 * clientConfig->QueueSize);
+
+        endpoint->SendRequest(
+            request.ExtractResult(),
+            MakeIntrusive<TCallContextBase>(0u));
+
+        ASSERT_TRUE(response.Received);
+        ASSERT_EQ(static_cast<ui32>(RDMA_PROTO_FAIL), response.Status);
+        ASSERT_GT(response.Bytes, 0u);
+
+        NProto::TError error =
+            ParseError(TStringBuf(response.Buffer).Head(response.Bytes));
+        ASSERT_EQ(static_cast<ui32>(E_RDMA_UNAVAILABLE), error.GetCode());
+    }
+
+    TEST(TRdmaClientTest, ShouldEagerlyDestroyBothMemoryWindowsOnAbortRequest)
+    {
+        auto testContext = MakeIntrusive<NVerbs::TTestContext>();
+        testContext->AllowConnect = true;
+
+        auto verbs = NVerbs::CreateTestVerbs(testContext);
+        auto monitoring = CreateMonitoringServiceStub();
+        auto clientConfig = std::make_shared<TClientConfig>();
+        clientConfig->UseMemoryWindows = true;
+
+        auto logging = CreateLoggingService(
+            "console",
+            TLogSettings{TLOG_RESOURCES});
+
+        auto client = CreateTestClient(
+            verbs,
+            logging,
+            monitoring,
+            clientConfig);
+
+        client->Start();
+        Y_DEFER {
+            client->Stop();
+        };
+
+        std::atomic<int> bound = 0;
+        std::atomic<int> destroyed = 0;
+
+        testContext->PostSend = [&](ibv_qp* qp, ibv_send_wr* wr) {
+            Y_UNUSED(qp);
+            if (wr->opcode == IBV_WR_BIND_MW) {
+                bound++;
+            }
+        };
+
+        testContext->DestroyMemoryWindow = [&](ibv_mw* mw) {
+            Y_UNUSED(mw);
+            destroyed++;
+        };
+
+        auto endpoint = client->StartEndpoint("::", 10020).ExtractValueSync();
+
+        struct THoldingHandler: IClientHandler
+        {
+            TClientRequestPtr Held;
+            TManualEvent Received;
+
+            void HandleResponse(
+                TClientRequestPtr req,
+                ui32 status,
+                size_t responseBytes) override
+            {
+                Y_UNUSED(status);
+                Y_UNUSED(responseBytes);
+                Held = std::move(req);
+                Received.Signal();
+            }
+        };
+
+        auto handler = std::make_shared<THoldingHandler>();
+
+        auto request = endpoint->AllocateRequest(
+            handler,
+            std::make_unique<TNullContext>(),
+            1024,
+            1024);
+        ASSERT_FALSE(HasError(request.GetError()));
+
+        auto reqId = endpoint->SendRequest(
+            request.ExtractResult(),
+            MakeIntrusive<TCallContextBase>(0u));
+
+        // wait until the request's memory windows have been acquired and
+        // the first bind work request posted
+        while (bound.load() < 1) {
+            SpinLockPause();
+        }
+
+        endpoint->CancelRequest(reqId);
+
+        handler->Received.Wait();
+        ASSERT_TRUE(handler->Held);
+        ASSERT_EQ(2, destroyed.load());
+
+        handler->Held.reset();
+    }
+
 }   // namespace NCloud::NStorage::NRdma