Add jl_getaffinity and jl_setaffinity

src/jl_exported_funcs.inc

@@ -194,6 +194,7 @@
  XX(jl_generating_output) \
  XX(jl_generic_function_def) \
  XX(jl_gensym) \
+ XX(jl_getaffinity) \
  XX(jl_getallocationgranularity) \
  XX(jl_getnameinfo) \
  XX(jl_getpagesize) \
@@ -406,6 +407,7 @@
  XX(jl_safepoint_suspend_thread) \
  XX(jl_safepoint_resume_thread) \
  XX(jl_SC_CLK_TCK) \
+ XX(jl_setaffinity) \
  XX(jl_set_ARGS) \
  XX(jl_set_const) \
  XX(jl_set_errno) \

src/julia_threads.h

@@ -373,6 +373,9 @@ JL_DLLEXPORT int8_t jl_gc_is_in_finalizer(void) JL_NOTSAFEPOINT;
 
 JL_DLLEXPORT void jl_wakeup_thread(int16_t tid);
 
+JL_DLLEXPORT int jl_getaffinity(int16_t tid, char *mask, int cpumasksize);
+JL_DLLEXPORT int jl_setaffinity(int16_t tid, char *mask, int cpumasksize);
+
 #ifdef __cplusplus
 }
 #endif

src/threading.c

@@ -984,6 +984,52 @@ JL_DLLEXPORT int jl_alignment(size_t sz)
  return jl_gc_alignment(sz);
 }
 
+// Return values:
+// 0 == success
+// 1 == invalid thread id provided
+// 2 == ptls2 was NULL
+// <0 == uv_thread_getaffinity exit code
+JL_DLLEXPORT int jl_getaffinity(int16_t tid, char *mask, int cpumasksize) {
+ int nthreads = jl_atomic_load_acquire(&jl_n_threads);
+ if (tid < 0 || tid >= nthreads)
+ return 1;
+
+ // TODO: use correct lock. system_id is only legal if the thread is alive.
+ jl_ptls_t ptls2 = jl_atomic_load_relaxed(&jl_all_tls_states)[tid];
+ if (ptls2 == NULL)
+ return 2;
+ uv_thread_t uvtid = ptls2->system_id;
+
+ int ret_uv = uv_thread_getaffinity(&uvtid, mask, cpumasksize);
+ if (ret_uv != 0)
+ return ret_uv;
+
+ return 0; // success
+}
+
+// Return values:
+// 0 == success
+// 1 == invalid thread id provided
+// 2 == ptls2 was NULL
+// <0 == uv_thread_getaffinity exit code
+JL_DLLEXPORT int jl_setaffinity(int16_t tid, char *mask, int cpumasksize) {
+ int nthreads = jl_atomic_load_acquire(&jl_n_threads);
+ if (tid < 0 || tid >= nthreads)
+ return 1;
+
+ // TODO: use correct lock. system_id is only legal if the thread is alive.
+ jl_ptls_t ptls2 = jl_atomic_load_relaxed(&jl_all_tls_states)[tid];
+ if (ptls2 == NULL)
+ return 2;
+ uv_thread_t uvtid = ptls2->system_id;
+
+ int ret_uv = uv_thread_setaffinity(&uvtid, mask, NULL, cpumasksize);
+ if (ret_uv != 0)
+ return ret_uv;
+
+ return 0; // success
+}
+
 #ifdef __cplusplus
 }
 #endif

test/threads.jl

@@ -354,3 +354,15 @@ end
  @test istaskfailed(t)
  end
 end
+
+@testset "jl_*affinity" begin
+ cpumasksize = @ccall uv_cpumask_size()::Cint
+ if !Sys.iswindows() && cpumasksize > 0 # otherwise affinities are not supported on the platform (UV_ENOTSUP)
+ mask = zeros(Cchar, cpumasksize);
+ jl_getaffinity = (tid, mask, cpumasksize) -> ccall(:jl_getaffinity, Int32, (Int16, Ptr{Cchar}, Int32), tid, mask, cpumasksize)
+ jl_setaffinity = (tid, mask, cpumasksize) -> ccall(:jl_setaffinity, Int32, (Int16, Ptr{Cchar}, Int32), tid, mask, cpumasksize)
+ @test jl_getaffinity(1, mask, cpumasksize) == 0
+ fill!(mask, 1)
+ @test jl_setaffinity(1, mask, cpumasksize) == 0
+ end
+end