diff --git a/include/rest_rpc/io_context_pool.hpp b/include/rest_rpc/io_context_pool.hpp new file mode 100644 index 0000000..68df390 --- /dev/null +++ b/include/rest_rpc/io_context_pool.hpp @@ -0,0 +1,53 @@ +#pragma once +#include +#include + +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(); + 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 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 &get_io_context() { + size_t i = next_.fetch_add(1, std::memory_order::relaxed); + return io_contexts_[i % io_contexts_.size()]; + } + +private: + std::vector> io_contexts_; + std::vector works_; + std::once_flag run_flag_; + std::once_flag stop_flag_; + std::atomic next_ = 0; +}; +} // namespace rest_rpc \ No newline at end of file diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 71bdcd7..05db6f1 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -7,6 +7,7 @@ include_directories( ) add_executable(${project_name} + test_io_context_pool.cpp test_rest_rpc.cpp main.cpp ) diff --git a/tests/test_io_context_pool.cpp b/tests/test_io_context_pool.cpp new file mode 100644 index 0000000..f1c0795 --- /dev/null +++ b/tests/test_io_context_pool.cpp @@ -0,0 +1,76 @@ +#include "../include/rest_rpc/io_context_pool.hpp" +#include "doctest/doctest.h" +#include + +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(2); + std::promise 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); +} \ No newline at end of file