io context pool

include/rest_rpc/io_context_pool.hpp

@@ -0,0 +1,53 @@
+#pragma once
+#include <asio.hpp>
+#include <vector>
+
+namespace rest_rpc {
+class io_context_pool {
+public:
+  explicit io_context_pool(size_t pool_size) {
+    if (pool_size == 0) {
+      pool_size = 1;
+    }
+
+    for (size_t i = 0; i < pool_size; i++) {
+      auto io_ctx = std::make_shared<asio::io_context>();
+      works_.push_back(asio::io_context::work(*io_ctx));
+      io_contexts_.emplace_back(io_ctx);
+    }
+  }
+
+  ~io_context_pool() { stop(); }
+
+  void run() {
+    std::call_once(run_flag_, [this] {
+      std::vector<std::thread> threads;
+      for (auto &ctx : io_contexts_) {
+        threads.push_back(std::thread([&ctx] { ctx->run(); }));
+      }
+
+      for (auto &thd : threads) {
+        thd.join();
+      }
+    });
+  }
+
+  void stop() {
+    std::call_once(stop_flag_, [this] { works_.clear(); });
+  }
+
+  size_t size() const { return io_contexts_.size(); }
+
+  std::shared_ptr<asio::io_context> &get_io_context() {
+    size_t i = next_.fetch_add(1, std::memory_order::relaxed);
+    return io_contexts_[i % io_contexts_.size()];
+  }
+
+private:
+  std::vector<std::shared_ptr<asio::io_context>> io_contexts_;
+  std::vector<asio::io_context::work> works_;
+  std::once_flag run_flag_;
+  std::once_flag stop_flag_;
+  std::atomic<size_t> next_ = 0;
+};
+} // namespace rest_rpc