Support RDMA SRQ

src/brpc/rdma/block_pool.cpp

@@ -27,6 +27,7 @@
 #include "butil/thread_local.h"
 #include "butil/memory/scope_guard.h"
 #include "brpc/rdma/block_pool.h"
+#include "brpc/rdma/srq_receive_pool.h"
 
 namespace brpc {
 namespace rdma {
@@ -127,7 +128,7 @@ static inline Region* GetRegion(const void* buf) {
 uint32_t GetRegionId(const void* buf) {
     Region* r = GetRegion(buf);
     if (!r) {
-        return 0;
+        return SrqReceivePool::RegionId(buf);
     }
     return r->id;
 }
@@ -461,6 +462,9 @@ int DeallocBlock(void* buf) {
 
     Region* r = GetRegion(buf);
     if (!r) {
+        if (SrqReceivePool::Recycle(buf)) {
+            return 0;
+        }
         errno = ERANGE;
         return -1;
     }

src/brpc/rdma/rdma_async_event.cpp

@@ -0,0 +1,105 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#if BRPC_WITH_RDMA
+#include "brpc/rdma/rdma_async_event.h"
+#include <mutex>
+#include "bthread/mutex.h"
+#include "butil/containers/flat_map.h"
+#include "butil/logging.h"
+
+namespace brpc {
+namespace rdma {
+extern void (*IbvAckAsyncEvent)(ibv_async_event*);
+
+namespace {
+struct AsyncEventRegistry {
+    // Serializes registration, removal and callback capture only. Event
+    // acknowledgement, verbs destruction and callback execution run outside
+    // this lock.
+    bthread::Mutex mutex;
+    // Look up verbs objects by identity. External objects may contain arbitrary
+    // context pointers that must not be dereferenced.
+    butil::FlatMap<const void*, RdmaAsyncEventContext*> contexts;
+};
+
+AsyncEventRegistry& Registry() {
+    // RDMA atexit cleanup still unregisters objects, so the registry must
+    // outlive static destruction.
+    static auto* registry = new AsyncEventRegistry;
+    return *registry;
+}
+}  // namespace
+
+void RegisterRdmaAsyncEventContext(const void* object,
+                                   RdmaAsyncEventContext* context) {
+    CHECK(object);
+    CHECK(context);
+    auto* registry = &Registry();
+    std::lock_guard<bthread::Mutex> lock(registry->mutex);
+    CHECK(registry->contexts.seek(object) == nullptr);
+    registry->contexts[object] = context;
+}
+
+void UnregisterRdmaAsyncEventContext(const void* object) {
+    auto* registry = &Registry();
+    std::lock_guard<bthread::Mutex> lock(registry->mutex);
+    registry->contexts.erase(object);
+}
+
+bool DispatchRdmaAsyncEvent(ibv_async_event* event) {
+    const void* object = nullptr;
+    switch (event->event_type) {
+    case IBV_EVENT_QP_REQ_ERR:
+    case IBV_EVENT_QP_ACCESS_ERR:
+    case IBV_EVENT_QP_FATAL:
+    case IBV_EVENT_QP_LAST_WQE_REACHED:
+    case IBV_EVENT_COMM_EST:
+    case IBV_EVENT_SQ_DRAINED:
+        object = event->element.qp;
+        break;
+    case IBV_EVENT_CQ_ERR:
+        object = event->element.cq;
+        break;
+    case IBV_EVENT_SRQ_ERR:
+    case IBV_EVENT_SRQ_LIMIT_REACHED:
+        object = event->element.srq;
+        break;
+    default:
+        return false;
+    }
+    std::function<void()> callback;
+    {
+        auto* registry = &Registry();
+        std::lock_guard<bthread::Mutex> lock(registry->mutex);
+        auto* context = registry->contexts.seek(object);
+        if (context != nullptr) {
+            callback = (*context)->PrepareAsyncEvent(event->event_type);
+        }
+    }
+    // Acknowledge events for unregistered objects too, without accessing their
+    // context. A callback may destroy its object. Acknowledge first so
+    // destruction cannot wait for this callback to acknowledge its own event.
+    IbvAckAsyncEvent(event);
+    if (callback) {
+        callback();
+    }
+    return true;
+}
+}  // namespace rdma
+}  // namespace brpc
+#endif  // BRPC_WITH_RDMA

src/brpc/rdma/rdma_async_event.h

@@ -0,0 +1,48 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#ifndef BRPC_RDMA_ASYNC_EVENT_H
+#define BRPC_RDMA_ASYNC_EVENT_H
+#if BRPC_WITH_RDMA
+#include <functional>
+#include <infiniband/verbs.h>
+
+namespace brpc {
+namespace rdma {
+// Prepares callbacks only for registered QPs, CQs and SRQs. Called under the
+// registry lock: capture references or event values only. Do not wait for
+// transport locks, destroy verbs objects or modify the registry. The captured
+// callback runs after unlocking and acknowledging the event, so it may initiate
+// destruction without waiting for its own event acknowledgement.
+class RdmaAsyncEventContext {
+public:
+    virtual ~RdmaAsyncEventContext() = default;
+    virtual std::function<void()> PrepareAsyncEvent(ibv_event_type type) = 0;
+};
+// Register after object creation and before publishing it to send/receive
+// paths. The registry does not own the context.
+void RegisterRdmaAsyncEventContext(const void* object, RdmaAsyncEventContext* context);
+// Unregister before verbs destruction. Waits for any in-flight
+// PrepareAsyncEvent call, but not for callbacks already captured. Destruction
+// and its event-acknowledgement wait must run without the registry lock. Even
+// if destruction fails, the object no longer receives new callbacks.
+void UnregisterRdmaAsyncEventContext(const void* object);
+bool DispatchRdmaAsyncEvent(ibv_async_event* event);
+}  // namespace rdma
+}  // namespace brpc
+#endif
+#endif  // BRPC_RDMA_ASYNC_EVENT_H