From 87597285e66b1d3a708fa2551698197601bd71b1 Mon Sep 17 00:00:00 2001
From: Elliot Saba <staticfloat@gmail.com>
Date: Sat, 19 Mar 2022 03:32:40 +0000
Subject: [PATCH] More flexibly test affinity setting

When running on a machine with `cpusets` applied, we are unable to
assign CPU affinity to CPUs 1 and 2; we may be locked to CPUs 9-16, for
example.  So we must inspect what our current cpumask is, and from that
select CPUs that we can safely assign affinity to in our tests.
---
 test/threads.jl | 28 +++++++++++++++-------------
 1 file changed, 15 insertions(+), 13 deletions(-)

diff --git a/test/threads.jl b/test/threads.jl
index dde50590ae08b8..975b65876f42fa 100644
--- a/test/threads.jl
+++ b/test/threads.jl
@@ -96,9 +96,10 @@ end
 const AFFINITY_SUPPORTED = (Sys.islinux() || Sys.iswindows()) && !running_under_rr()
 
 if AFFINITY_SUPPORTED
-    if Sys.CPU_THREADS > 1
-        @test run_with_affinity([2]) == "2"
-        @test run_with_affinity([1, 2]) == "1,2"
+    allowed_cpus = findall(uv_thread_getaffinity())
+    if length(allowed_cpus) ≥ 2
+        @test run_with_affinity(allowed_cpus[1:1]) == "$(allowed_cpus[1])"
+        @test run_with_affinity(allowed_cpus[1:2]) == "$(allowed_cpus[1]),$(allowed_cpus[2])"
     end
 end
 
@@ -113,18 +114,19 @@ function get_nthreads(options = ``; cpus = nothing)
 end
 
 @testset "nthreads determined based on CPU affinity" begin
-    if AFFINITY_SUPPORTED && Sys.CPU_THREADS ≥ 2
+    allowed_cpus = findall(uv_thread_getaffinity())
+    if AFFINITY_SUPPORTED && length(allowed_cpus) ≥ 2
         @test get_nthreads() ≥ 2
-        @test get_nthreads(cpus = [1]) == 1
-        @test get_nthreads(cpus = [2]) == 1
-        @test get_nthreads(cpus = [1, 2]) == 2
-        @test get_nthreads(`-t1`, cpus = [1]) == 1
-        @test get_nthreads(`-t1`, cpus = [2]) == 1
-        @test get_nthreads(`-t1`, cpus = [1, 2]) == 1
+        @test get_nthreads(cpus = allowed_cpus[1:1]) == 1
+        @test get_nthreads(cpus = allowed_cpus[2:2]) == 1
+        @test get_nthreads(cpus = allowed_cpus[1:2]) == 2
+        @test get_nthreads(`-t1`, cpus = allowed_cpus[1:1]) == 1
+        @test get_nthreads(`-t1`, cpus = allowed_cpus[2:2]) == 1
+        @test get_nthreads(`-t1`, cpus = allowed_cpus[1:2]) == 1
 
-        if Sys.CPU_THREADS ≥ 3
-            @test get_nthreads(cpus = [1, 3]) == 2
-            @test get_nthreads(cpus = [2, 3]) == 2
+        if length(allowed_cpus) ≥ 3
+            @test get_nthreads(cpus = allowed_cpus[1:2:3]) == 2
+            @test get_nthreads(cpus = allowed_cpus[2:3])   == 2
         end
     end
 end