[coro_io][ibverbs] fix async_event_watcher (#1144)

Author: poor-circle

include/ylt/coro_io/ibverbs/ib_device.hpp

@@ -25,6 +25,7 @@
 #include <system_error>
 #include <unordered_map>
 
+#include "asio/dispatch.hpp"
 #include "asio/ip/address.hpp"
 #include "asio/posix/stream_descriptor.hpp"
 #include "ib_buffer.hpp"
@@ -334,6 +335,19 @@ class ib_device_t : public std::enable_shared_from_this<ib_device_t> {
     return buffer_pool_;
   }
 
+  struct async_event_watcher_manager_t
+      : public std::unordered_map<ib_device_t*, std::weak_ptr<ib_device_t>> {
+    ~async_event_watcher_manager_t() {
+      for (auto& [_, dev] : *this) {
+        if (auto ptr = dev.lock(); ptr) {
+          ptr->stop_async_event_watcher();
+        }
+      }
+    }
+  };
+
+  ~ib_device_t() { stop_async_event_watcher(); }
+
  private:
   int ipv6_addr_v4mapped(const struct in6_addr* a) {
     return ((a->s6_addr32[0] | a->s6_addr32[1]) |
@@ -362,6 +376,12 @@ class ib_device_t : public std::enable_shared_from_this<ib_device_t> {
     return default_gid_index;
   }
 
+  template <typename T>
+  static bool weak_ptrs_equal(const std::weak_ptr<T>& a,
+                              const std::weak_ptr<T>& b) {
+    return !a.owner_before(b) && !b.owner_before(a);
+  }
+
   void start_async_event_watcher() {
     if (!ctx_) {
       return;
@@ -373,32 +393,35 @@ class ib_device_t : public std::enable_shared_from_this<ib_device_t> {
           << "Error, failed to change file descriptor of async event queue\n";
       return;
     }
-
-    auto fd_deleter = [](asio::posix::stream_descriptor* fd) {
-      fd->release();
-      delete fd;
-    };
-    auto fd =
-        std::unique_ptr<asio::posix::stream_descriptor, decltype(fd_deleter)>(
-            new asio::posix::stream_descriptor(
-                coro_io::get_global_executor()->get_asio_executor(),
-                ctx_->async_fd));
+    auto executor = coro_io::get_global_executor();
+    async_event_watcher_fd_ = std::unique_ptr<asio::posix::stream_descriptor>(
+        new asio::posix::stream_descriptor(executor->get_asio_executor(),
+                                           ctx_->async_fd));
     auto listen_event = [](std::weak_ptr<ib_device_t> dev,
-                           auto fd) -> async_simple::coro::Lazy<void> {
+                           coro_io::ExecutorWrapper<>* executor)
+        -> async_simple::coro::Lazy<void> {
       std::error_code ec;
-      auto name = std::string{dev.lock()->name()};
+      auto self = dev.lock();
+      if (self == nullptr) {
+        co_return;
+      }
+      auto self_raw_ptr = self.get();
+      (*executor->get_data_with_default<async_event_watcher_manager_t>(
+          "ylt_ib_watch"))[self_raw_ptr] = dev;
+      auto name = std::string{self->name()};
       ELOG_INFO << "start_async_event_watcher of device:" << name << " start";
       while (!ec) {
         coro_io::callback_awaitor<std::error_code> awaitor;
-        ec = co_await awaitor.await_resume([&fd](auto handler) {
-          fd->async_wait(asio::posix::stream_descriptor::wait_read,
-                         [handler](const std::error_code& ec) mutable {
-                           handler.set_value_then_resume(ec);
-                         });
+        ec = co_await awaitor.await_resume([&self](auto handler) {
+          self->async_event_watcher_fd_->async_wait(
+              asio::posix::stream_descriptor::wait_read,
+              [handler](const std::error_code& ec) mutable {
+                handler.set_value_then_resume(ec);
+              });
+          self = nullptr;
         });
-
         if (!ec) {
-          auto self = dev.lock();
+          self = dev.lock();
           if (!self) {
             ELOG_DEBUG
                 << "ib_device async event stop listening by close device:"
@@ -412,14 +435,28 @@ class ib_device_t : public std::enable_shared_from_this<ib_device_t> {
                      << ec.message() << ",device:" << name;
         }
       }
+      auto* table = executor->get_data<async_event_watcher_manager_t>(
+          "ib_device_async_event_watcher");
+      if (table) {
+        auto iter = table->find(self_raw_ptr);
+        if (iter != table->end() && weak_ptrs_equal(iter->second, dev)) {
+          table->erase(iter);
+        }
+      }
     };
-    listen_event(shared_from_this(), std::move(fd)).start([](auto&& ec) {
-    });
+    listen_event(shared_from_this(), executor).via(executor).detach();
+  }
+
+  void stop_async_event_watcher() {
+    std::error_code ec;
+    [[maybe_unused]] auto _ = async_event_watcher_fd_->cancel(ec);
+    async_event_watcher_fd_->release();
   }
 
   std::string name_;
   std::unique_ptr<ibv_context, ib_deleter> ctx_;
   std::unique_ptr<ibv_pd, ib_deleter> pd_;
+  std::unique_ptr<asio::posix::stream_descriptor> async_event_watcher_fd_;
   std::shared_ptr<ib_buffer_pool_t> buffer_pool_;
   std::atomic<bool> support_inline_data_ = true;
   ibv_port_attr attr_;

include/ylt/coro_io/io_context_pool.hpp

@@ -28,6 +28,7 @@
 #include <iostream>
 #include <memory>
 #include <mutex>
+#include <span>
 #include <thread>
 #include <type_traits>
 #include <utility>
@@ -249,8 +250,8 @@ class io_context_pool {
     promise_.set_value();
   }
 
-  void stop() {
-    std::call_once(flag_, [this] {
+  void stop(bool force = false) {
+    std::call_once(flag_, [this, force] {
       bool has_run_or_stop = false;
       bool ok = has_run_or_stop_.compare_exchange_strong(has_run_or_stop, true);
 
@@ -262,15 +263,16 @@ class io_context_pool {
 
       work_.clear();
 
-      if (ok) {
-        // clear all unfinished work
-        for (auto &e : io_contexts_) {
-          e->run();
+      for (auto &e : io_contexts_) {
+        if (ok) {
+          e->poll();  // clear all unfinished work
+        }
+        if (force) {
+          e->stop();
         }
-        return;
       }
-
-      promise_.get_future().wait();
+      if (!ok)
+        promise_.get_future().wait();
     });
   }
 
@@ -281,6 +283,10 @@ class io_context_pool {
 
   std::size_t pool_size() const noexcept { return io_contexts_.size(); }
 
+  std::span<std::unique_ptr<coro_io::ExecutorWrapper<>>> get_all_executor() {
+    return executors;
+  }
+
   bool has_stop() const { return work_.empty(); }
 
   size_t current_io_context() { return next_io_context_ - 1; }

include/ylt/coro_rpc/impl/context.hpp

@@ -44,13 +44,13 @@ class context_base {
   typename rpc_protocol::req_header &get_req_head() { return self_->req_head_; }
 
   bool check_status() {
-    auto old_flag = self_->status_.exchange(context_status::start_response);
-    if (old_flag != context_status::init)
+    auto expected = context_status::init;
+    if (!self_->status_.compare_exchange_strong(expected,
+                                                context_status::start_response))
       AS_UNLIKELY {
         ELOG_ERROR << "response message more than one time";
         return false;
       }
-
     if (self_->has_closed())
       AS_UNLIKELY {
         ELOG_DEBUG << "response_msg failed: connection has been closed";

src/coro_io/tests/ibverbs/ib_socket_pressure_test.cpp

@@ -36,6 +36,8 @@
 struct config_t {
   std::size_t buffer_size = 256 * 1024;
   std::size_t request_size = 20 * 1024 * 1024 + 1;
+  std::size_t recv_buffer_cnt = 8;
+  std::size_t send_buffer_cnt = 4;
   int concurrency = 50;
   int test_type = 0;
   int enable_log = 0;
@@ -44,8 +46,9 @@ struct config_t {
   int port = 58110;
   int test_time = 10;
 };
-YLT_REFL(config_t, buffer_size, request_size, concurrency, test_type,
-         enable_log, enable_server, enable_client, port, test_time);
+YLT_REFL(config_t, buffer_size, request_size, recv_buffer_cnt, send_buffer_cnt,
+         concurrency, test_type, enable_log, enable_server, enable_client, port,
+         test_time);
 
 config_t config;
 std::shared_ptr<coro_io::ib_device_t> g_dev;
@@ -193,7 +196,11 @@ async_simple::coro::Lazy<std::error_code> echo_accept() {
 
   ELOG_INFO << "tcp listening";
   while (true) {
-    coro_io::ib_socket_t soc(coro_io::get_global_executor(), {.device = g_dev});
+    coro_io::ib_socket_t soc(
+        coro_io::get_global_executor(),
+        {.recv_buffer_cnt = (uint16_t)config.recv_buffer_cnt,
+         .send_buffer_cnt = (uint16_t)config.send_buffer_cnt,
+         .device = g_dev});
     auto ec = co_await coro_io::async_accept(acceptor, soc);
 
     if (ec) [[unlikely]] {
@@ -309,7 +316,10 @@ async_simple::coro::Lazy<std::error_code> echo_client_read_some(
 }
 
 async_simple::coro::Lazy<std::error_code> echo_connect() {
-  coro_io::ib_socket_t soc(coro_io::get_global_executor(), {.device = g_dev});
+  coro_io::ib_socket_t soc(coro_io::get_global_executor(),
+                           {.recv_buffer_cnt = (uint16_t)config.recv_buffer_cnt,
+                            .send_buffer_cnt = (uint16_t)config.send_buffer_cnt,
+                            .device = g_dev});
   auto ec = co_await coro_io::async_connect(soc, config.enable_client,
                                             std::to_string(config.port));
   if (ec) [[unlikely]] {