From a30136af8a7e5046d4e2d3588ef480db8141dabe Mon Sep 17 00:00:00 2001
From: qicosmos <qicosmos@163.com>
Date: Fri, 23 Feb 2024 16:08:50 +0800
Subject: [PATCH] allow set cpu affinity (#531)

---
 example/benchmark.cpp                         |  2 +-
 include/cinatra/coro_http_server.hpp          |  6 ++++--
 .../cinatra/ylt/coro_io/io_context_pool.hpp   | 21 ++++++++++++++++++-
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/example/benchmark.cpp b/example/benchmark.cpp
index 53637d9b..45d55f53 100644
--- a/example/benchmark.cpp
+++ b/example/benchmark.cpp
@@ -4,7 +4,7 @@ using namespace cinatra;
 using namespace std::chrono_literals;
 
 int main() {
-  coro_http_server server(std::thread::hardware_concurrency(), 8090);
+  coro_http_server server(std::thread::hardware_concurrency(), 8090, true);
   server.set_http_handler<GET>(
       "/plaintext", [](coro_http_request& req, coro_http_response& resp) {
         resp.get_conn()->set_multi_buf(false);
diff --git a/include/cinatra/coro_http_server.hpp b/include/cinatra/coro_http_server.hpp
index 01d57e62..9f9ae0b9 100644
--- a/include/cinatra/coro_http_server.hpp
+++ b/include/cinatra/coro_http_server.hpp
@@ -30,8 +30,10 @@ class coro_http_server {
   coro_http_server(asio::io_context &ctx, unsigned short port)
       : out_ctx_(&ctx), port_(port), acceptor_(ctx), check_timer_(ctx) {}
 
-  coro_http_server(size_t thread_num, unsigned short port)
-      : pool_(std::make_unique<coro_io::io_context_pool>(thread_num)),
+  coro_http_server(size_t thread_num, unsigned short port,
+                   bool cpu_affinity = false)
+      : pool_(std::make_unique<coro_io::io_context_pool>(thread_num,
+                                                         cpu_affinity)),
         port_(port),
         acceptor_(pool_->get_executor()->get_asio_executor()),
         check_timer_(pool_->get_executor()->get_asio_executor()) {}
diff --git a/include/cinatra/ylt/coro_io/io_context_pool.hpp b/include/cinatra/ylt/coro_io/io_context_pool.hpp
index b3667e29..df3fad53 100644
--- a/include/cinatra/ylt/coro_io/io_context_pool.hpp
+++ b/include/cinatra/ylt/coro_io/io_context_pool.hpp
@@ -27,6 +27,10 @@
 #include <thread>
 #include <type_traits>
 #include <vector>
+#ifdef __linux__
+#include <pthread.h>
+#include <sched.h>
+#endif
 
 namespace coro_io {
 
@@ -107,7 +111,8 @@ get_current_executor() {
 class io_context_pool {
  public:
   using executor_type = asio::io_context::executor_type;
-  explicit io_context_pool(std::size_t pool_size) : next_io_context_(0) {
+  explicit io_context_pool(std::size_t pool_size, bool cpu_affinity = false)
+      : next_io_context_(0), cpu_affinity_(cpu_affinity) {
     if (pool_size == 0) {
       pool_size = 1;  // set default value as 1
     }
@@ -139,6 +144,19 @@ class io_context_pool {
             svr->run();
           },
           io_contexts_[i]));
+
+#ifdef __linux__
+      if (cpu_affinity_) {
+        cpu_set_t cpuset;
+        CPU_ZERO(&cpuset);
+        CPU_SET(i, &cpuset);
+        int rc = pthread_setaffinity_np(threads.back()->native_handle(),
+                                        sizeof(cpu_set_t), &cpuset);
+        if (rc != 0) {
+          std::cerr << "Error calling pthread_setaffinity_np: " << rc << "\n";
+        }
+      }
+#endif
     }
 
     for (std::size_t i = 0; i < threads.size(); ++i) {
@@ -197,6 +215,7 @@ class io_context_pool {
   std::promise<void> promise_;
   std::atomic<bool> has_run_or_stop_ = false;
   std::once_flag flag_;
+  bool cpu_affinity_ = false;
 };
 
 class multithread_context_pool {