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.

20 files changed, 204 insertions, 6 deletions

diff --git a/Help/manual/cmake-properties.7.rst b/Help/manual/cmake-properties.7.rst
index d3e58d0..9d4a7e8 100644
--- a/Help/manual/cmake-properties.7.rst
+++ b/Help/manual/cmake-properties.7.rst
@@ -348,6 +348,7 @@ Properties on Tests
    /prop_test/LABELS
    /prop_test/MEASUREMENT
    /prop_test/PASS_REGULAR_EXPRESSION
+   /prop_test/PROCESSOR_AFFINITY
    /prop_test/PROCESSORS
    /prop_test/REQUIRED_FILES
    /prop_test/RESOURCE_LOCK
diff --git a/Help/prop_test/PROCESSORS.rst b/Help/prop_test/PROCESSORS.rst
index a1211fb..a927c10 100644
--- a/Help/prop_test/PROCESSORS.rst
+++ b/Help/prop_test/PROCESSORS.rst
@@ -2,6 +2,7 @@ PROCESSORS
 ----------
 
 Set to specify how many process slots this test requires.
+If not set, the default is ``1`` processor.
 
 Denotes the number of processors that this test will require.  This is
 typically used for MPI tests, and should be used in conjunction with
@@ -11,3 +12,5 @@ This will also be used to display a weighted test timing result in label and
 subproject summaries in the command line output of :manual:`ctest(1)`. The wall
 clock time for the test run will be multiplied by this property to give a
 better idea of how much cpu resource CTest allocated for the test.
+
+See also the :prop_test:`PROCESSOR_AFFINITY` test property.
diff --git a/Help/prop_test/PROCESSOR_AFFINITY.rst b/Help/prop_test/PROCESSOR_AFFINITY.rst
new file mode 100644
index 0000000..38ec179
--- /dev/null
+++ b/Help/prop_test/PROCESSOR_AFFINITY.rst
@@ -0,0 +1,11 @@
+PROCESSOR_AFFINITY
+------------------
+
+Set to a true value to ask CTest to launch the test process with CPU affinity
+for a fixed set of processors.  If enabled and supported for the current
+platform, CTest will choose a set of processors to place in the CPU affinity
+mask when launching the test process.  The number of processors in the set is
+determined by the :prop_test:`PROCESSORS` test property or the number of
+processors available to CTest, whichever is smaller.  The set of processors
+chosen will be disjoint from the processors assigned to other concurrently
+running tests that also have the ``PROCESSOR_AFFINITY`` property enabled.
diff --git a/Help/release/dev/ctest-affinity.rst b/Help/release/dev/ctest-affinity.rst
new file mode 100644
index 0000000..f4f72a5
--- /dev/null
+++ b/Help/release/dev/ctest-affinity.rst
@@ -0,0 +1,6 @@
+ctest-affinity
+--------------
+
+* A :prop_test:`PROCESSOR_AFFINITY` test property was added to request
+  that CTest run a test with CPU affinity for a set of processors
+  disjoint from other concurrently running tests with the property set.
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();
+}
diff --git a/Tests/CMakeLib/CMakeLists.txt b/Tests/CMakeLib/CMakeLists.txt
index 06df53f..126076d 100644
--- a/Tests/CMakeLib/CMakeLists.txt
+++ b/Tests/CMakeLib/CMakeLists.txt
@@ -49,3 +49,6 @@ if(TEST_CompileCommandOutput)
 endif()
 
 add_subdirectory(PseudoMemcheck)
+
+add_executable(testAffinity testAffinity.cxx)
+target_link_libraries(testAffinity CMakeLib)
diff --git a/Tests/CMakeLib/testAffinity.cxx b/Tests/CMakeLib/testAffinity.cxx
new file mode 100644
index 0000000..4b82280
--- /dev/null
+++ b/Tests/CMakeLib/testAffinity.cxx
@@ -0,0 +1,18 @@
+/* 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 <cstddef>
+#include <iostream>
+#include <set>
+
+int main()
+{
+  std::set<size_t> cpus = cmAffinity::GetProcessorsAvailable();
+  if (!cpus.empty()) {
+    std::cout << "CPU affinity mask count is '" << cpus.size() << "'.\n";
+  } else {
+    std::cout << "CPU affinity not supported on this platform.\n";
+  }
+  return 0;
+}
diff --git a/Tests/RunCMake/CMakeLists.txt b/Tests/RunCMake/CMakeLists.txt
index 31069fa..beadb1d 100644
--- a/Tests/RunCMake/CMakeLists.txt
+++ b/Tests/RunCMake/CMakeLists.txt
@@ -339,6 +339,9 @@ add_RunCMake_test(CPackConfig)
 add_RunCMake_test(CPackInstallProperties)
 add_RunCMake_test(ExternalProject)
 add_RunCMake_test(FetchContent)
+if(NOT CMake_TEST_EXTERNAL_CMAKE)
+  set(CTestCommandLine_ARGS -DTEST_AFFINITY=$<TARGET_FILE:testAffinity>)
+endif()
 add_RunCMake_test(CTestCommandLine)
 add_RunCMake_test(CacheNewline)
 # Only run this test on unix platforms that support
diff --git a/Tests/RunCMake/CTestCommandLine/RunCMakeTest.cmake b/Tests/RunCMake/CTestCommandLine/RunCMakeTest.cmake
index 0fafea5..3033c9c 100644
--- a/Tests/RunCMake/CTestCommandLine/RunCMakeTest.cmake
+++ b/Tests/RunCMake/CTestCommandLine/RunCMakeTest.cmake
@@ -141,3 +141,23 @@ function(run_TestOutputSize)
     )
 endfunction()
 run_TestOutputSize()
+
+function(run_TestAffinity)
+  set(RunCMake_TEST_BINARY_DIR ${RunCMake_BINARY_DIR}/TestAffinity)
+  set(RunCMake_TEST_NO_CLEAN 1)
+  file(REMOVE_RECURSE "${RunCMake_TEST_BINARY_DIR}")
+  file(MAKE_DIRECTORY "${RunCMake_TEST_BINARY_DIR}")
+  # Create a test with affinity enabled.  The default PROCESSORS
+  # value is 1, so our expected output checks that this is the
+  # number of processors in the mask.
+  file(WRITE "${RunCMake_TEST_BINARY_DIR}/CTestTestfile.cmake" "
+  add_test(Affinity \"${TEST_AFFINITY}\")
+  set_tests_properties(Affinity PROPERTIES PROCESSOR_AFFINITY ON)
+")
+  # Run ctest with a large parallel level so that the value is
+  # not responsible for capping the number of processors available.
+  run_cmake_command(TestAffinity ${CMAKE_CTEST_COMMAND} -V -j 64)
+endfunction()
+if(TEST_AFFINITY)
+  run_TestAffinity()
+endif()
diff --git a/Tests/RunCMake/CTestCommandLine/TestAffinity-stdout.txt b/Tests/RunCMake/CTestCommandLine/TestAffinity-stdout.txt
new file mode 100644
index 0000000..e23d30b
--- /dev/null
+++ b/Tests/RunCMake/CTestCommandLine/TestAffinity-stdout.txt
@@ -0,0 +1 @@
+1: CPU affinity (mask count is '1'|not supported on this platform)\.