CTest: Add options to control test process affinity to CPUs

In commit v2.8.0~170 (ENH: Added ctest test options PROCESSORS and
RUN_SERIAL, 2009-09-07) CTest learned to track the number of processors
allocated to running tests in order to balance it against the desired
level of parallelism.  Extend this idea by introducing a new
`PROCESSOR_AFFINITY` test property to ask that CTest run a test
with the CPU affinity mask set.  This will allow a set of tests
that are running concurrently to use disjoint CPU resources.

11 files changed, 138 insertions, 6 deletions

diff --git a/Source/CMakeLists.txt b/Source/CMakeLists.txt
index a0010a2..e547356 100644
--- a/Source/CMakeLists.txt
+++ b/Source/CMakeLists.txt
@@ -131,6 +131,8 @@ set(SRCS
   LexerParser/cmListFileLexer.c
   LexerParser/cmListFileLexer.in.l
 
+  cmAffinity.cxx
+  cmAffinity.h
   cmArchiveWrite.cxx
   cmBase32.cxx
   cmCacheManager.cxx
diff --git a/Source/CTest/cmCTestMultiProcessHandler.cxx b/Source/CTest/cmCTestMultiProcessHandler.cxx
index 50c2d86..99cf551 100644
--- a/Source/CTest/cmCTestMultiProcessHandler.cxx
+++ b/Source/CTest/cmCTestMultiProcessHandler.cxx
@@ -2,6 +2,7 @@
    file Copyright.txt or https://cmake.org/licensing for details.  */
 #include "cmCTestMultiProcessHandler.h"
 
+#include "cmAffinity.h"
 #include "cmCTest.h"
 #include "cmCTestRunTest.h"
 #include "cmCTestScriptHandler.h"
@@ -53,6 +54,8 @@ cmCTestMultiProcessHandler::cmCTestMultiProcessHandler()
   this->TestLoad = 0;
   this->Completed = 0;
   this->RunningCount = 0;
+  this->ProcessorsAvailable = cmAffinity::GetProcessorsAvailable();
+  this->HaveAffinity = this->ProcessorsAvailable.size();
   this->StopTimePassed = false;
   this->HasCycles = false;
   this->SerialTestRunning = false;
@@ -127,6 +130,21 @@ bool cmCTestMultiProcessHandler::StartTestProcess(int test)
     return false;
   }
 
+  if (this->HaveAffinity && this->Properties[test]->WantAffinity) {
+    size_t needProcessors = this->GetProcessorsUsed(test);
+    if (needProcessors > this->ProcessorsAvailable.size()) {
+      return false;
+    }
+    std::vector<size_t> affinity;
+    affinity.reserve(needProcessors);
+    for (size_t i = 0; i < needProcessors; ++i) {
+      auto p = this->ProcessorsAvailable.begin();
+      affinity.push_back(*p);
+      this->ProcessorsAvailable.erase(p);
+    }
+    this->Properties[test]->Affinity = std::move(affinity);
+  }
+
   cmCTestOptionalLog(this->CTest, HANDLER_VERBOSE_OUTPUT,
                      "test " << test << "\n", this->Quiet);
   this->TestRunningMap[test] = true; // mark the test as running
@@ -200,6 +218,11 @@ inline size_t cmCTestMultiProcessHandler::GetProcessorsUsed(int test)
   if (processors > this->ParallelLevel) {
     processors = this->ParallelLevel;
   }
+  // Cap tests that want affinity to the maximum affinity available.
+  if (this->HaveAffinity && processors > this->HaveAffinity &&
+      this->Properties[test]->WantAffinity) {
+    processors = this->HaveAffinity;
+  }
   return processors;
 }
 
@@ -398,6 +421,11 @@ void cmCTestMultiProcessHandler::FinishTestProcess(cmCTestRunTest* runner,
   this->UnlockResources(test);
   this->RunningCount -= GetProcessorsUsed(test);
 
+  for (auto p : properties->Affinity) {
+    this->ProcessorsAvailable.insert(p);
+  }
+  properties->Affinity.clear();
+
   delete runner;
   if (started) {
     this->StartNextTests();
diff --git a/Source/CTest/cmCTestMultiProcessHandler.h b/Source/CTest/cmCTestMultiProcessHandler.h
index 7837ff9..19e1a35 100644
--- a/Source/CTest/cmCTestMultiProcessHandler.h
+++ b/Source/CTest/cmCTestMultiProcessHandler.h
@@ -119,6 +119,8 @@ protected:
   // Number of tests that are complete
   size_t Completed;
   size_t RunningCount;
+  std::set<size_t> ProcessorsAvailable;
+  size_t HaveAffinity;
   bool StopTimePassed;
   // list of test properties (indices concurrent to the test map)
   PropertiesMap Properties;
diff --git a/Source/CTest/cmCTestRunTest.cxx b/Source/CTest/cmCTestRunTest.cxx
index 30ad38c..1ae36f7 100644
--- a/Source/CTest/cmCTestRunTest.cxx
+++ b/Source/CTest/cmCTestRunTest.cxx
@@ -515,7 +515,8 @@ bool cmCTestRunTest::StartTest(size_t total)
   }
 
   return this->ForkProcess(timeout, this->TestProperties->ExplicitTimeout,
-                           &this->TestProperties->Environment);
+                           &this->TestProperties->Environment,
+                           &this->TestProperties->Affinity);
 }
 
 void cmCTestRunTest::ComputeArguments()
@@ -591,7 +592,8 @@ void cmCTestRunTest::DartProcessing()
 }
 
 bool cmCTestRunTest::ForkProcess(cmDuration testTimeOut, bool explicitTimeout,
-                                 std::vector<std::string>* environment)
+                                 std::vector<std::string>* environment,
+                                 std::vector<size_t>* affinity)
 {
   this->TestProcess = cm::make_unique<cmProcess>(*this);
   this->TestProcess->SetId(this->Index);
@@ -637,7 +639,8 @@ bool cmCTestRunTest::ForkProcess(cmDuration testTimeOut, bool explicitTimeout,
     cmSystemTools::AppendEnv(*environment);
   }
 
-  return this->TestProcess->StartProcess(this->MultiTestHandler.Loop);
+  return this->TestProcess->StartProcess(this->MultiTestHandler.Loop,
+                                         affinity);
 }
 
 void cmCTestRunTest::WriteLogOutputTop(size_t completed, size_t total)
diff --git a/Source/CTest/cmCTestRunTest.h b/Source/CTest/cmCTestRunTest.h
index 4d57357..073af11 100644
--- a/Source/CTest/cmCTestRunTest.h
+++ b/Source/CTest/cmCTestRunTest.h
@@ -83,7 +83,8 @@ private:
   void DartProcessing();
   void ExeNotFound(std::string exe);
   bool ForkProcess(cmDuration testTimeOut, bool explicitTimeout,
-                   std::vector<std::string>* environment);
+                   std::vector<std::string>* environment,
+                   std::vector<size_t>* affinity);
   void WriteLogOutputTop(size_t completed, size_t total);
   // Run post processing of the process output for MemCheck
   void MemCheckPostProcess();
diff --git a/Source/CTest/cmCTestTestHandler.cxx b/Source/CTest/cmCTestTestHandler.cxx
index 84d8926..d1affd4 100644
--- a/Source/CTest/cmCTestTestHandler.cxx
+++ b/Source/CTest/cmCTestTestHandler.cxx
@@ -2165,6 +2165,9 @@ bool cmCTestTestHandler::SetTestsProperties(
               rt.Processors = 1;
             }
           }
+          if (key == "PROCESSOR_AFFINITY") {
+            rt.WantAffinity = cmSystemTools::IsOn(val.c_str());
+          }
           if (key == "SKIP_RETURN_CODE") {
             rt.SkipReturnCode = atoi(val.c_str());
             if (rt.SkipReturnCode < 0 || rt.SkipReturnCode > 255) {
@@ -2336,6 +2339,7 @@ bool cmCTestTestHandler::AddTest(const std::vector<std::string>& args)
   test.ExplicitTimeout = false;
   test.Cost = 0;
   test.Processors = 1;
+  test.WantAffinity = false;
   test.SkipReturnCode = -1;
   test.PreviousRuns = 0;
   if (this->UseIncludeRegExpFlag &&
diff --git a/Source/CTest/cmCTestTestHandler.h b/Source/CTest/cmCTestTestHandler.h
index f4978b6..d2694a1 100644
--- a/Source/CTest/cmCTestTestHandler.h
+++ b/Source/CTest/cmCTestTestHandler.h
@@ -130,6 +130,8 @@ public:
     int Index;
     // Requested number of process slots
     int Processors;
+    bool WantAffinity;
+    std::vector<size_t> Affinity;
     // return code of test which will mark test as "not run"
     int SkipReturnCode;
     std::vector<std::string> Environment;
diff --git a/Source/CTest/cmProcess.cxx b/Source/CTest/cmProcess.cxx
index 09ed0a9..5c9b169 100644
--- a/Source/CTest/cmProcess.cxx
+++ b/Source/CTest/cmProcess.cxx
@@ -83,7 +83,7 @@ void cmProcess::SetCommandArguments(std::vector<std::string> const& args)
   this->Arguments = args;
 }
 
-bool cmProcess::StartProcess(uv_loop_t& loop)
+bool cmProcess::StartProcess(uv_loop_t& loop, std::vector<size_t>* affinity)
 {
   this->ProcessState = cmProcess::State::Error;
   if (this->Command.empty()) {
@@ -138,6 +138,22 @@ bool cmProcess::StartProcess(uv_loop_t& loop)
   options.stdio_count = 3; // in, out and err
   options.exit_cb = &cmProcess::OnExitCB;
   options.stdio = stdio;
+#if !defined(CMAKE_USE_SYSTEM_LIBUV)
+  std::vector<char> cpumask;
+  if (affinity && !affinity->empty()) {
+    cpumask.resize(static_cast<size_t>(uv_cpumask_size()), 0);
+    for (auto p : *affinity) {
+      cpumask[p] = 1;
+    }
+    options.cpumask = cpumask.data();
+    options.cpumask_size = cpumask.size();
+  } else {
+    options.cpumask = nullptr;
+    options.cpumask_size = 0;
+  }
+#else
+  static_cast<void>(affinity);
+#endif
 
   status =
     uv_read_start(pipe_reader, &cmProcess::OnAllocateCB, &cmProcess::OnReadCB);
diff --git a/Source/CTest/cmProcess.h b/Source/CTest/cmProcess.h
index 20e24b9..b2d87fa 100644
--- a/Source/CTest/cmProcess.h
+++ b/Source/CTest/cmProcess.h
@@ -36,7 +36,7 @@ public:
   void ChangeTimeout(cmDuration t);
   void ResetStartTime();
   // Return true if the process starts
-  bool StartProcess(uv_loop_t& loop);
+  bool StartProcess(uv_loop_t& loop, std::vector<size_t>* affinity);
 
   enum class State
   {
diff --git a/Source/cmAffinity.cxx b/Source/cmAffinity.cxx
new file mode 100644
index 0000000..bdf1f42
--- /dev/null
+++ b/Source/cmAffinity.cxx
@@ -0,0 +1,62 @@
+/* Distributed under the OSI-approved BSD 3-Clause License.  See accompanying
+   file Copyright.txt or https://cmake.org/licensing for details.  */
+#include "cmAffinity.h"
+
+#include "cm_uv.h"
+
+#ifndef CMAKE_USE_SYSTEM_LIBUV
+#ifdef _WIN32
+#define CM_HAVE_CPU_AFFINITY
+#include <windows.h>
+#elif defined(__linux__) || defined(__FreeBSD__)
+#define CM_HAVE_CPU_AFFINITY
+#include <pthread.h>
+#include <sched.h>
+#if defined(__FreeBSD__)
+#include <pthread_np.h>
+#include <sys/cpuset.h>
+#include <sys/param.h>
+#endif
+#if defined(__linux__)
+typedef cpu_set_t cm_cpuset_t;
+#else
+typedef cpuset_t cm_cpuset_t;
+#endif
+#endif
+#endif
+
+namespace cmAffinity {
+
+std::set<size_t> GetProcessorsAvailable()
+{
+  std::set<size_t> processorsAvailable;
+#ifdef CM_HAVE_CPU_AFFINITY
+  int cpumask_size = uv_cpumask_size();
+  if (cpumask_size > 0) {
+#ifdef _WIN32
+    DWORD_PTR procmask;
+    DWORD_PTR sysmask;
+    if (GetProcessAffinityMask(GetCurrentProcess(), &procmask, &sysmask) !=
+        0) {
+      for (int i = 0; i < cpumask_size; ++i) {
+        if (procmask & (((DWORD_PTR)1) << i)) {
+          processorsAvailable.insert(i);
+        }
+      }
+    }
+#else
+    cm_cpuset_t cpuset;
+    CPU_ZERO(&cpuset); // NOLINT(clang-tidy)
+    if (pthread_getaffinity_np(pthread_self(), sizeof(cpuset), &cpuset) == 0) {
+      for (int i = 0; i < cpumask_size; ++i) {
+        if (CPU_ISSET(i, &cpuset)) {
+          processorsAvailable.insert(i);
+        }
+      }
+    }
+#endif
+  }
+#endif
+  return processorsAvailable;
+}
+}
diff --git a/Source/cmAffinity.h b/Source/cmAffinity.h
new file mode 100644
index 0000000..3775bae
--- /dev/null
+++ b/Source/cmAffinity.h
@@ -0,0 +1,12 @@
+/* Distributed under the OSI-approved BSD 3-Clause License.  See accompanying
+   file Copyright.txt or https://cmake.org/licensing for details.  */
+#pragma once
+#include "cmConfigure.h" // IWYU pragma: keep
+
+#include <cstddef>
+#include <set>
+
+namespace cmAffinity {
+
+std::set<size_t> GetProcessorsAvailable();
+}