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

tests/CMakeLists.txt

@@ -7,6 +7,7 @@ include_directories(
         )
 
 add_executable(${project_name}
+        test_io_context_pool.cpp
 		test_rest_rpc.cpp
         main.cpp
         )

tests/test_io_context_pool.cpp

@@ -0,0 +1,76 @@
+#include "../include/rest_rpc/io_context_pool.hpp"
+#include "doctest/doctest.h"
+#include <iostream>
+
+using namespace rest_rpc;
+
+TEST_CASE("test context pool") {
+  io_context_pool pool(4);
+  bool quit = false;
+  std::thread thd([&pool, &quit] {
+    pool.run();
+    quit = true;
+  });
+
+  pool.stop();
+  thd.join();
+  CHECK(quit);
+}
+
+TEST_CASE("test context pool stop before run") {
+  io_context_pool pool(2);
+  bool quit = false;
+  pool.stop();
+  std::thread thd([&pool, &quit] {
+    pool.run();
+    quit = true;
+  });
+
+  pool.stop();
+  thd.join();
+  CHECK(quit);
+}
+
+TEST_CASE("test context pool multiple run") {
+  io_context_pool pool(2);
+  bool quit = false;
+  std::thread thd([&pool, &quit] {
+    pool.run();
+    quit = true;
+  });
+
+  std::thread thd1([&pool] { pool.run(); });
+
+  pool.stop();
+  thd.join();
+  thd1.join();
+  CHECK(quit);
+}
+
+TEST_CASE("test context pool automatic stop") {
+  auto pool = std::make_shared<io_context_pool>(2);
+  std::promise<void> p;
+  std::thread thd([&pool, &p] {
+    bool block = true;
+    p.set_value();
+    pool->run();
+    block = false;
+    CHECK(block);
+  });
+  p.get_future().wait();
+  std::this_thread::sleep_for(std::chrono::milliseconds(200));
+
+  thd.detach();
+}
+
+TEST_CASE("test get io context") {
+  io_context_pool pool(2);
+  CHECK(pool.size() == 2);
+  auto &ctx1 = pool.get_io_context();
+  auto &ctx2 = pool.get_io_context();
+  auto &ctx3 = pool.get_io_context();
+  auto ctx4 = pool.get_io_context();
+  CHECK(ctx1 != ctx2);
+  CHECK(ctx1 == ctx3);
+  CHECK(ctx2 == ctx4);
+}