diff --git a/src/jl_exported_funcs.inc b/src/jl_exported_funcs.inc index e2527e3d7aeab..5e70566ab310e 100644 --- a/src/jl_exported_funcs.inc +++ b/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) \ diff --git a/src/julia_threads.h b/src/julia_threads.h index 3a4e9a66cf5a7..60f156b8fab14 100644 --- a/src/julia_threads.h +++ b/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 diff --git a/src/threading.c b/src/threading.c index eb76ac579b538..d6e28e3d80828 100644 --- a/src/threading.c +++ b/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 diff --git a/test/threads.jl b/test/threads.jl index 172385d1c130e..24da56e2b89fd 100644 --- a/test/threads.jl +++ b/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