From b49b669ac127d06aea34ef32799dac44538dcae7 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Wed, 14 Sep 2022 15:33:16 +0300
Subject: [PATCH] Support affinity mask on Linux.

---
 .../td/utils/port/detail/ThreadPthread.cpp    | 34 +++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/tdutils/td/utils/port/detail/ThreadPthread.cpp b/tdutils/td/utils/port/detail/ThreadPthread.cpp
index 91bbf8ca7..a6ed9fa67 100644
--- a/tdutils/td/utils/port/detail/ThreadPthread.cpp
+++ b/tdutils/td/utils/port/detail/ThreadPthread.cpp
@@ -11,6 +11,7 @@ char disable_linker_warning_about_empty_file_thread_pthread_cpp TD_UNUSED;
 #if TD_THREAD_PTHREAD
 
 #include "td/utils/misc.h"
+#include "td/utils/port/detail/skip_eintr.h"
 
 #include <pthread.h>
 #include <sched.h>
@@ -95,11 +96,44 @@ int ThreadPthread::do_pthread_create(pthread_t *thread, const pthread_attr_t *at
 }
 
 Status ThreadPthread::set_affinity_mask(id thread_id, uint64 mask) {
+#if TD_LINUX
+  cpu_set_t cpuset;
+  CPU_ZERO(&cpuset);
+  for (int j = 0; j < 64 && j < CPU_SETSIZE; j++) {
+    if ((mask >> j) & 1) {
+      CPU_SET(j, &cpuset);
+    }
+  }
+
+  auto res = skip_eintr([&] { return pthread_setaffinity_np(thread_id, sizeof(cpu_set_t), &cpuset); });
+  if (res) {
+    return OS_ERROR("Failed to set thread affinity mask");
+  }
+  return Status::OK();
+#else
   return Status::Error("Unsupported");
+#endif
 }
 
 uint64 ThreadPthread::get_affinity_mask(id thread_id) {
+#if TD_LINUX
+  cpu_set_t cpuset;
+  CPU_ZERO(&cpuset);
+  auto res = skip_eintr([&] { return pthread_getaffinity_np(thread_id, sizeof(cpu_set_t), &cpuset); });
+  if (res) {
+    return 0;
+  }
+
+  uint64 mask = 0;
+  for (int j = 0; j < 64 && j < CPU_SETSIZE; j++) {
+    if (CPU_ISSET(j, &cpuset)) {
+      mask |= static_cast<uint64>(1) << j;
+    }
+  }
+  return mask;
+#else
   return 0;
+#endif
 }
 
 namespace this_thread_pthread {