From 925549051db1594a3ba68407915516fba1439a6d Mon Sep 17 00:00:00 2001
From: likepeng <likepeng0418@163.com>
Date: Fri, 3 Jan 2025 16:13:19 +0800
Subject: [PATCH] add SetThreadAffinity method

---
 myframe/common.cpp         | 37 +++++++++++++++++++++++++++++++++++++
 myframe/common.h           |  4 ++++
 myframe/worker_context.cpp | 13 +++++++++++++
 myframe/worker_context.h   |  1 +
 4 files changed, 55 insertions(+)

diff --git a/myframe/common.cpp b/myframe/common.cpp
index 10b98bc..8a587ed 100644
--- a/myframe/common.cpp
+++ b/myframe/common.cpp
@@ -126,4 +126,41 @@ std::vector<std::string_view> Common::SplitMsgName(const std::string& name) {
   return tokens;
 }
 
+// 可以通过std::thread::hardware_concurrency()获得核心数
+int Common::SetThreadAffinity(std::thread* t, int cpu_core) {
+#if defined(MYFRAME_OS_WINDOWS)
+  auto hThread = t->native_handle();
+  DWORD_PTR mask = 1 << cpu_core;
+  if (SetThreadAffinityMask(hThread, mask) == 0) {
+    return -1;
+  }
+  return 0;
+#else
+  cpu_set_t cpuset;
+  CPU_ZERO(&cpuset);
+  CPU_SET(cpu_core, &cpuset);
+  auto handle = t->native_handle();
+  int result = pthread_setaffinity_np(handle, sizeof(cpu_set_t), &cpuset);
+  return result;
+#endif
+}
+
+int Common::SetSelfThreadAffinity(int cpu_core) {
+#if defined(MYFRAME_OS_WINDOWS)
+  auto hThread = GetCurrentThread();
+  DWORD_PTR mask = 1 << cpu_core;
+  if (SetThreadAffinityMask(hThread, mask) == 0) {
+    return -1;
+  }
+  return 0;
+#else
+  cpu_set_t cpuset;
+  CPU_ZERO(&cpuset);
+  CPU_SET(cpu_core, &cpuset);
+  auto handle = pthread_self();
+  int result = pthread_setaffinity_np(handle, sizeof(cpu_set_t), &cpuset);
+  return result;
+#endif
+}
+
 }  // namespace myframe
diff --git a/myframe/common.h b/myframe/common.h
index 7fd8b5a..5a8b9b6 100644
--- a/myframe/common.h
+++ b/myframe/common.h
@@ -10,6 +10,7 @@ Author: 李柯鹏 <likepeng0418@163.com>
 #include <memory>
 #include <string>
 #include <vector>
+#include <thread>
 #include <string_view>
 #if __has_include(<filesystem>)
 #include <filesystem>
@@ -33,6 +34,9 @@ class MYFRAME_EXPORT Common final {
   static bool IsAbsolutePath(const std::string& path);
 
   static std::vector<std::string_view> SplitMsgName(const std::string& name);
+
+  static int SetThreadAffinity(std::thread* t, int cpu_core);
+  static int SetSelfThreadAffinity(int cpu_core);
 };
 
 }  // namespace myframe
diff --git a/myframe/worker_context.cpp b/myframe/worker_context.cpp
index d3b9070..5fe2d1d 100644
--- a/myframe/worker_context.cpp
+++ b/myframe/worker_context.cpp
@@ -13,6 +13,7 @@ Author: 李柯鹏 <likepeng0418@163.com>
 #include "myframe/msg.h"
 #include "myframe/worker.h"
 #include "myframe/app.h"
+#include "myframe/common.h"
 
 namespace myframe {
 
@@ -60,6 +61,18 @@ void WorkerContext::Join() {
   }
 }
 
+bool WorkerContext::SetThreadAffinity(int cpu_core) {
+  if (runing_.load()) {
+    if (0 == Common::SetThreadAffinity(&th_, cpu_core)) {
+      return true;
+    }
+    LOG(WARNING) << GetName() << " bind cpu " << cpu_core << " failed";
+  } else {
+    LOG(WARNING) << GetName() << " not runing, skip SetThreadAffinity";
+  }
+  return false;
+}
+
 void WorkerContext::Initialize() {
   mailbox_.SetAddr(worker_->GetWorkerName());
   worker_->Init();
diff --git a/myframe/worker_context.h b/myframe/worker_context.h
index a18a489..0b13db6 100644
--- a/myframe/worker_context.h
+++ b/myframe/worker_context.h
@@ -40,6 +40,7 @@ class WorkerContext final : public Event {
   void Join();
   bool IsRuning() { return runing_.load(); }
   std::thread::id GetThreadId() { return th_.get_id(); }
+  bool SetThreadAffinity(int cpu_core);
 
   /// event 相关函数
   ev_handle_t GetHandle() const override;