Add support to parse multiple cores for setting CPU affinity

include/realtime_tools/realtime_helpers.hpp

@@ -32,6 +32,7 @@
 #include <string>
 #include <thread>
 #include <utility>
+#include <vector>
 
 #ifdef _WIN32
 #include <windows.h>
@@ -80,6 +81,19 @@ bool lock_memory(std::string & message);
 */
 std::pair<bool, std::string> lock_memory();
 
+/**
+ * Configure the caller thread affinity - Tell the scheduler to prefer a certain
+ * set of cores for the given thread handle.
+ * \note The threads created by the calling thread will inherit the affinity.
+ * \param[in] thread the thread handle of the thread
+ * \param[in] core the cpu numbers of the core. If an empty vector is passed,
+ * the affinity is reset to the default.
+ * \returns a pair of a boolean indicating whether the operation succeeded or not
+ * and a message describing the result of the operation
+*/
+std::pair<bool, std::string> set_thread_affinity(
+  NATIVE_THREAD_HANDLE thread, const std::vector<int> & cores);
+
 /**
  * Configure the caller thread affinity - Tell the scheduler to prefer a certain
  * core for the given thread handle.
@@ -115,6 +129,17 @@ std::pair<bool, std::string> set_thread_affinity(std::thread & thread, int core)
 */
 std::pair<bool, std::string> set_current_thread_affinity(int core);
 
+/**
+ * Configure the current thread affinity - Tell the scheduler to prefer a certain
+ * set of cores for the current thread.
+ * \note The threads created by the calling thread will inherit the affinity.
+ * \param[in] core the cpu numbers of the core. If an empty vector is passed,
+ * the affinity is reset to the default.
+ * \returns a pair of a boolean indicating whether the operation succeeded or not
+ * and a message describing the result of the operation
+*/
+std::pair<bool, std::string> set_current_thread_affinity(const std::vector<int> & cores);
+
 /**
  * Method to get the amount of available cpu cores
  * \ref https://man7.org/linux/man-pages/man3/sysconf.3.html

src/realtime_helpers.cpp

@@ -120,7 +120,8 @@ std::pair<bool, std::string> lock_memory()
 #endif
 }
 
-std::pair<bool, std::string> set_thread_affinity(NATIVE_THREAD_HANDLE thread, int core)
+std::pair<bool, std::string> set_thread_affinity(
+  NATIVE_THREAD_HANDLE thread, const std::vector<int> & cores)
 {
   std::string message;
 #ifdef _WIN32
@@ -156,34 +157,48 @@ std::pair<bool, std::string> set_thread_affinity(NATIVE_THREAD_HANDLE thread, in
   // Obtain available processors
   const auto number_of_cores = get_number_of_available_processors();
 
+  bool valid_cpu_set = true;
   // Reset affinity by setting it to all cores
-  if (core < 0) {
+  if (cores.empty()) {
     for (auto i = 0; i < number_of_cores; i++) {
       CPU_SET(i, &cpuset);
     }
-    // And actually tell the schedular to set the affinity of the thread of respective pid
-    const auto result = set_affinity_result_message(
-      pthread_setaffinity_np(thread, sizeof(cpu_set_t), &cpuset), message);
-    return std::make_pair(result, message);
+  } else {
+    for (const auto core : cores) {
+      if (core < 0 || core >= number_of_cores) {
+        valid_cpu_set = false;
+        break;
+      }
+      CPU_SET(core, &cpuset);
+    }
   }
 
-  if (core < number_of_cores) {
-    // Set the passed core to the cpu set
-    CPU_SET(core, &cpuset);
+  if (valid_cpu_set) {
     // And actually tell the schedular to set the affinity of the thread of respective pid
     const auto result = set_affinity_result_message(
       pthread_setaffinity_np(thread, sizeof(cpu_set_t), &cpuset), message);
     return std::make_pair(result, message);
   }
+  // create a string from the core numbers
+  std::string core_numbers;
+  for (const auto core : cores) {
+    core_numbers += std::to_string(core) + " ";
+  }
   // Invalid core number passed
-  message = "Invalid core number : '" + std::to_string(core) + "' passed! The system has " +
+  message = "Invalid core numbers : ['" + core_numbers + "'] passed! The system has " +
             std::to_string(number_of_cores) +
             " cores. Parsed core number should be between 0 and " +
             std::to_string(number_of_cores - 1);
   return std::make_pair(false, message);
 #endif
 }
 
+std::pair<bool, std::string> set_thread_affinity(NATIVE_THREAD_HANDLE thread, int core)
+{
+  const std::vector<int> affinity_cores = core < 0 ? std::vector<int>() : std::vector<int>{core};
+  return set_thread_affinity(thread, affinity_cores);
+}
+
 std::pair<bool, std::string> set_thread_affinity(std::thread & thread, int core)
 {
   if (!thread.joinable()) {
@@ -202,6 +217,15 @@ std::pair<bool, std::string> set_current_thread_affinity(int core)
 #endif
 }
 
+std::pair<bool, std::string> set_current_thread_affinity(const std::vector<int> & cores)
+{
+#ifdef _WIN32
+  return set_thread_affinity(GetCurrentThread(), cores);
+#else
+  return set_thread_affinity(pthread_self(), cores);
+#endif
+}
+
 int64_t get_number_of_available_processors()
 {
 #ifdef _WIN32

test/thread_priority_tests.cpp

@@ -56,6 +56,11 @@ TEST(thread_priority, set_thread_affinity_invalid_too_many_cores)
   const int count = static_cast<int>(realtime_tools::get_number_of_available_processors());
   // We should always have at least one core
   EXPECT_FALSE(realtime_tools::set_thread_affinity(t.native_handle(), count + 10).first);
+  EXPECT_FALSE(realtime_tools::set_thread_affinity(t.native_handle(), {0, 1, count + 2}).first);
+  EXPECT_FALSE(realtime_tools::set_thread_affinity(t.native_handle(), {0, -1, 2}).first);
+  EXPECT_FALSE(realtime_tools::set_thread_affinity(t.native_handle(), {0, 3, -1}).first);
+  EXPECT_FALSE(
+    realtime_tools::set_thread_affinity(t.native_handle(), std::vector<int>({-1})).first);
   t.join();
 }
 
@@ -64,6 +69,10 @@ TEST(thread_priority, set_current_thread_affinity_invalid_too_many_cores)
   const int count = static_cast<int>(realtime_tools::get_number_of_available_processors());
   // We should always have at least one core
   EXPECT_FALSE(realtime_tools::set_current_thread_affinity(count + 10).first);
+  EXPECT_FALSE(realtime_tools::set_current_thread_affinity({count + 10}).first);
+  EXPECT_FALSE(realtime_tools::set_current_thread_affinity({0, count + 10}).first);
+  EXPECT_FALSE(realtime_tools::set_current_thread_affinity({0, -1}).first);
+  EXPECT_FALSE(realtime_tools::set_current_thread_affinity(std::vector<int>({-1})).first);
 }
 
 TEST(thread_priority, set_thread_affinity_valid_reset)
@@ -72,17 +81,21 @@ TEST(thread_priority, set_thread_affinity_valid_reset)
   std::thread t([]() { std::this_thread::sleep_for(std::chrono::milliseconds(100)); });
   // Reset core affinity
   EXPECT_TRUE(realtime_tools::set_thread_affinity(t.native_handle(), -1).first);
+  EXPECT_TRUE(realtime_tools::set_thread_affinity(t.native_handle(), {}).first);
   t.join();
 }
 
 TEST(thread_priority, set_current_thread_affinity_valid_reset)
 {
   // Reset core affinity
   EXPECT_TRUE(realtime_tools::set_current_thread_affinity(-1).first);
+  EXPECT_TRUE(realtime_tools::set_current_thread_affinity({}).first);
 }
 
 TEST(thread_priority, set_current_thread_affinity_valid)
 {
   // We should always have at least one core
   EXPECT_TRUE(realtime_tools::set_current_thread_affinity(0).first);
+  const int count = static_cast<int>(realtime_tools::get_number_of_available_processors());
+  EXPECT_TRUE(realtime_tools::set_current_thread_affinity({0, count - 1}).first);
 }