1 files changed, 42 insertions, 5 deletions

diff --git a/Source/CTest/cmCTestMultiProcessHandler.cxx b/Source/CTest/cmCTestMultiProcessHandler.cxx
index 50c2d86..14b5caa 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"
@@ -19,6 +20,7 @@
 
 #include <algorithm>
 #include <chrono>
+#include <cstring>
 #include <iomanip>
 #include <list>
 #include <math.h>
@@ -53,6 +55,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 +131,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
@@ -151,13 +170,19 @@ bool cmCTestMultiProcessHandler::StartTestProcess(int test)
     }
   }
 
-  cmWorkingDirectory workdir(this->Properties[test]->Directory);
-
-  // Lock the resources we'll be using
+  // Always lock the resources we'll be using, even if we fail to set the
+  // working directory because FinishTestProcess() will try to unlock them
   this->LockResources(test);
 
-  if (testRun->StartTest(this->Total)) {
-    return true;
+  cmWorkingDirectory workdir(this->Properties[test]->Directory);
+  if (workdir.Failed()) {
+    testRun->StartFailure("Failed to change working directory to " +
+                          this->Properties[test]->Directory + " : " +
+                          std::strerror(workdir.GetLastResult()));
+  } else {
+    if (testRun->StartTest(this->Total)) {
+      return true;
+    }
   }
 
   this->FinishTestProcess(testRun, false);
@@ -200,6 +225,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 +428,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();
@@ -666,6 +701,8 @@ void cmCTestMultiProcessHandler::PrintTestList()
     count++;
     cmCTestTestHandler::cmCTestTestProperties& p = *it.second;
 
+    // Don't worry if this fails, we are only showing the test list, not
+    // running the tests
     cmWorkingDirectory workdir(p.Directory);
 
     cmCTestRunTest testRun(*this);