From 0eddd8d4058a6420ea41b21f27dede4717e3d56f Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Wed, 14 Sep 2022 14:38:14 +0300
Subject: [PATCH] Add thread::set_affinity_mask/get_affinity_mask.

---
 .../td/utils/port/detail/ThreadPthread.cpp    |  8 +++++
 tdutils/td/utils/port/detail/ThreadPthread.h  |  5 ++++
 tdutils/td/utils/port/detail/ThreadStl.h      | 30 +++++++++++++++++++
 tdutils/test/port.cpp                         | 27 +++++++++++++++++
 4 files changed, 70 insertions(+)

diff --git a/tdutils/td/utils/port/detail/ThreadPthread.cpp b/tdutils/td/utils/port/detail/ThreadPthread.cpp
index 64ea20dc0..91bbf8ca7 100644
--- a/tdutils/td/utils/port/detail/ThreadPthread.cpp
+++ b/tdutils/td/utils/port/detail/ThreadPthread.cpp
@@ -94,6 +94,14 @@ int ThreadPthread::do_pthread_create(pthread_t *thread, const pthread_attr_t *at
   return pthread_create(thread, attr, start_routine, arg);
 }
 
+Status ThreadPthread::set_affinity_mask(id thread_id, uint64 mask) {
+  return Status::Error("Unsupported");
+}
+
+uint64 ThreadPthread::get_affinity_mask(id thread_id) {
+  return 0;
+}
+
 namespace this_thread_pthread {
 ThreadPthread::id get_id() {
   return pthread_self();
diff --git a/tdutils/td/utils/port/detail/ThreadPthread.h b/tdutils/td/utils/port/detail/ThreadPthread.h
index ff54633c6..a0f15d1a5 100644
--- a/tdutils/td/utils/port/detail/ThreadPthread.h
+++ b/tdutils/td/utils/port/detail/ThreadPthread.h
@@ -17,6 +17,7 @@
 #include "td/utils/port/detail/ThreadIdGuard.h"
 #include "td/utils/port/thread_local.h"
 #include "td/utils/Slice.h"
+#include "td/utils/Status.h"
 
 #include <tuple>
 #include <type_traits>
@@ -66,6 +67,10 @@ class ThreadPthread {
 
   static void send_real_time_signal(id thread_id, int real_time_signal_number);
 
+  static Status set_affinity_mask(id thread_id, uint64 mask);
+
+  static uint64 get_affinity_mask(id thread_id);
+
  private:
   MovableValue<bool> is_inited_;
   pthread_t thread_;
diff --git a/tdutils/td/utils/port/detail/ThreadStl.h b/tdutils/td/utils/port/detail/ThreadStl.h
index a15dd6bda..ce6ac27fe 100644
--- a/tdutils/td/utils/port/detail/ThreadStl.h
+++ b/tdutils/td/utils/port/detail/ThreadStl.h
@@ -15,6 +15,7 @@
 #include "td/utils/port/detail/ThreadIdGuard.h"
 #include "td/utils/port/thread_local.h"
 #include "td/utils/Slice.h"
+#include "td/utils/Status.h"
 
 #include <thread>
 #include <tuple>
@@ -75,6 +76,35 @@ class ThreadStl {
     // not supported
   }
 
+  static Status set_affinity_mask(id thread_id, uint64 mask) {
+#if TD_WINDOWS
+    if (static_cast<DWORD_PTR>(mask) != mask) {
+      return Status::Error("Invalid thread affinity mask specified");
+    }
+    if (SetThreadAffinityMask(thread_id, static_cast<DWORD_PTR>(mask))) {
+      return Status::OK();
+    }
+    return OS_ERROR("Failed to set thread affinity mask");
+#else
+    return Status::Error("Unsupported");
+#endif
+  }
+
+  static uint64 get_affinity_mask(id thread_id) {
+#if TD_WINDOWS
+    DWORD_PTR process_mask = 0;
+    DWORD_PTR system_mask = 0;
+    if (GetProcessAffinityMask(GetCurrentProcess(), &process_mask, &system_mask)) {
+      auto result = SetThreadAffinityMask(thread_id, process_mask);
+      if (result != 0 && result != process_mask) {
+        SetThreadAffinityMask(thread_id, result);
+      }
+      return result;
+    }
+#endif
+    return 0;
+  }
+
  private:
   std::thread thread_;
 
diff --git a/tdutils/test/port.cpp b/tdutils/test/port.cpp
index a7c673e18..79a41596c 100644
--- a/tdutils/test/port.cpp
+++ b/tdutils/test/port.cpp
@@ -284,3 +284,30 @@ TEST(Port, EventFdAndSignals) {
 }
 #endif
 #endif
+
+#if !TD_THREAD_UNSUPPORTED
+TEST(Port, ThreadAffinityMask) {
+  auto thread_id = td::this_thread::get_id();
+  auto old_mask = td::thread::get_affinity_mask(thread_id);
+  LOG(INFO) << "Initial thread affinity mask: " << old_mask;
+  for (size_t i = 0; i < 64; i++) {
+    auto mask = td::thread::get_affinity_mask(thread_id);
+    LOG(INFO) << mask;
+    auto result = td::thread::set_affinity_mask(thread_id, static_cast<td::uint64>(1) << i);
+    LOG(INFO) << i << ": " << result;
+    mask = td::thread::get_affinity_mask(thread_id);
+    LOG(INFO) << mask;
+
+    if (i <= 1) {
+      td::thread thread([] {
+        auto mask = td::thread::get_affinity_mask(td::this_thread::get_id());
+        LOG(INFO) << "New thread affinity mask: " << mask;
+      });
+    }
+  }
+  auto result = td::thread::set_affinity_mask(thread_id, old_mask);
+  LOG(INFO) << result;
+  old_mask = td::thread::get_affinity_mask(thread_id);
+  LOG(INFO) << old_mask;
+}
+#endif