add SetThreadAffinity method

lkpworkspace · Jan 3, 2025 · 9255490 · 9255490
1 parent bf51d9e
commit 9255490
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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;