make Workqueue threadsafe (#30838)

base/boot.jl

@@ -380,8 +380,8 @@ eval(Core, :(LineInfoNode(mod::Module, method::Symbol, file::Symbol, line::Int,
 
 Module(name::Symbol=:anonymous, std_imports::Bool=true) = ccall(:jl_f_new_module, Ref{Module}, (Any, Bool), name, std_imports)
 
-function Task(@nospecialize(f), reserved_stack::Int=0)
- return ccall(:jl_new_task, Ref{Task}, (Any, Int), f, reserved_stack)
+function _Task(@nospecialize(f), reserved_stack::Int, completion_future)
+ return ccall(:jl_new_task, Ref{Task}, (Any, Any, Int), f, completion_future, reserved_stack)
 end
 
 # simple convert for use by constructors of types in Core

base/task.jl

@@ -2,6 +2,9 @@
 
 ## basic task functions and TLS
 
+const ThreadSynchronizer = GenericCondition{Threads.SpinLock}
+Core.Task(@nospecialize(f), reserved_stack::Int=0) = Core._Task(f, reserved_stack, ThreadSynchronizer())
+
 # Container for a captured exception and its backtrace. Can be serialized.
 struct CapturedException <: Exception
  ex::Any
@@ -135,6 +138,8 @@ istaskstarted(t::Task) = ccall(:jl_is_task_started, Cint, (Any,), t) != 0
 
 istaskfailed(t::Task) = (t.state == :failed)
 
+Threads.threadid(t::Task) = Int(ccall(:jl_get_task_tid, Int16, (Any,), t)+1)
+
 task_result(t::Task) = t.result
 
 task_local_storage() = get_task_tls(current_task())
@@ -181,13 +186,15 @@ end
 # NOTE: you can only wait for scheduled tasks
 function wait(t::Task)
  if !istaskdone(t)
- if t.donenotify === nothing
- t.donenotify = Condition()
+ lock(t.donenotify)
+ try
+ while !istaskdone(t)
+ wait(t.donenotify)
+ end
+ finally
+ unlock(t.donenotify)
  end
  end
- while !istaskdone(t)
- wait(t.donenotify)
- end
  if istaskfailed(t)
  throw(t.exception)
  end
@@ -273,7 +280,7 @@ end
 function register_taskdone_hook(t::Task, hook)
  tls = get_task_tls(t)
  push!(get!(tls, :TASKDONE_HOOKS, []), hook)
- t
+ return t
 end
 
 # runtime system hook called when a task finishes
@@ -286,9 +293,17 @@ function task_done_hook(t::Task)
  t.backtrace = catch_backtrace()
  end
 
- if isa(t.donenotify, Condition) && !isempty(t.donenotify.waitq)
- handled = true
- notify(t.donenotify, result, true, err)
+ donenotify = t.donenotify
+ if isa(donenotify, ThreadSynchronizer)
+ lock(donenotify)
+ try
+ if !isempty(donenotify.waitq)
+ handled = true
+ notify(donenotify, result, true, err)
+ end
+ finally
+ unlock(donenotify)
+ end
  end
 
  # Execute any other hooks registered in the TLS
@@ -298,8 +313,8 @@ function task_done_hook(t::Task)
  handled = true
  end
 
- if err && !handled
- if isa(result,InterruptException) && isdefined(Base,:active_repl_backend) &&
+ if err && !handled && Threads.threadid() == 1
+ if isa(result, InterruptException) && isdefined(Base, :active_repl_backend) &&
  active_repl_backend.backend_task.state == :runnable && isempty(Workqueue) &&
  active_repl_backend.in_eval
  throwto(active_repl_backend.backend_task, result) # this terminates the task
@@ -313,7 +328,8 @@ function task_done_hook(t::Task)
  # If an InterruptException happens while blocked in the event loop, try handing
  # the exception to the REPL task since the current task is done.
  # issue #19467
- if isa(e,InterruptException) && isdefined(Base,:active_repl_backend) &&
+ if Threads.threadid() == 1 &&
+ isa(e, InterruptException) && isdefined(Base, :active_repl_backend) &&
  active_repl_backend.backend_task.state == :runnable && isempty(Workqueue) &&
  active_repl_backend.in_eval
  throwto(active_repl_backend.backend_task, e)
@@ -360,12 +376,78 @@ end
 
 ## scheduler and work queue
 
-global const Workqueue = InvasiveLinkedList{Task}()
+struct InvasiveLinkedListSynchronized{T}
+ queue::InvasiveLinkedList{T}
+ lock::Threads.SpinLock
+ InvasiveLinkedListSynchronized{T}() where {T} = new(InvasiveLinkedList{T}(), Threads.SpinLock())
+end
+isempty(W::InvasiveLinkedListSynchronized) = isempty(W.queue)
+length(W::InvasiveLinkedListSynchronized) = length(W.queue)
+function push!(W::InvasiveLinkedListSynchronized{T}, t::T) where T
+ lock(W.lock)
+ try
+ push!(W.queue, t)
+ finally
+ unlock(W.lock)
+ end
+ return W
+end
+function pushfirst!(W::InvasiveLinkedListSynchronized{T}, t::T) where T
+ lock(W.lock)
+ try
+ pushfirst!(W.queue, t)
+ finally
+ unlock(W.lock)
+ end
+ return W
+end
+function pop!(W::InvasiveLinkedListSynchronized)
+ lock(W.lock)
+ try
+ return pop!(W.queue)
+ finally
+ unlock(W.lock)
+ end
+end
+function popfirst!(W::InvasiveLinkedListSynchronized)
+ lock(W.lock)
+ try
+ return popfirst!(W.queue)
+ finally
+ unlock(W.lock)
+ end
+end
+function list_deletefirst!(W::InvasiveLinkedListSynchronized{T}, t::T) where T
+ lock(W.lock)
+ try
+ list_deletefirst!(W.queue, t)
+ finally
+ unlock(W.lock)
+ end
+ return W
+end
+
+const StickyWorkqueue = InvasiveLinkedListSynchronized{Task}
+global const Workqueues = [StickyWorkqueue()]
+global const Workqueue = Workqueues[1] # default work queue is thread 1
+function __preinit_threads__()
+ if length(Workqueues) < Threads.nthreads()
+ resize!(Workqueues, Threads.nthreads())
+ for i = 2:length(Workqueues)
+ Workqueues[i] = StickyWorkqueue()
+ end
+ end
+ nothing
+end
 
 function enq_work(t::Task)
  (t.state == :runnable && t.queue === nothing) || error("schedule: Task not runnable")
- ccall(:uv_stop, Cvoid, (Ptr{Cvoid},), eventloop())
- push!(Workqueue, t)
+ tid = (t.sticky ? Threads.threadid(t) : 0)
+ if tid == 0
+ tid = Threads.threadid()
+ end
+ push!(Workqueues[tid], t)
+ tid == 1 && ccall(:uv_stop, Cvoid, (Ptr{Cvoid},), eventloop())
  return t
 end
 
@@ -418,11 +500,12 @@ end
 # fast version of `schedule(t, arg); wait()`
 function schedule_and_wait(t::Task, @nospecialize(arg)=nothing)
  (t.state == :runnable && t.queue === nothing) || error("schedule: Task not runnable")
- if isempty(Workqueue)
+ W = Workqueues[Threads.threadid()]
+ if isempty(W)
  return yieldto(t, arg)
  else
  t.result = arg
- push!(Workqueue, t)
+ push!(W, t)
  end
  return wait()
 end
@@ -487,23 +570,24 @@ end
 
 function ensure_rescheduled(othertask::Task)
  ct = current_task()
+ W = Workqueues[Threads.threadid()]
  if ct !== othertask && othertask.state == :runnable
  # we failed to yield to othertask
- # return it to the head of the queue to be scheduled later
- pushfirst!(Workqueue, othertask)
- end
- if ct.queue === Workqueue
- # if the current task was queued,
- # also need to return it to the runnable state
- # before throwing an error
- list_deletefirst!(Workqueue, ct)
+ # return it to the head of a queue to be retried later
+ tid = Threads.threadid(othertask)
+ Wother = tid == 0 ? W : Workqueues[tid]
+ pushfirst!(Wother, othertask)
  end
+ # if the current task was queued,
+ # also need to return it to the runnable state
+ # before throwing an error
+ list_deletefirst!(W, ct)
  nothing
 end
 
-function trypoptask()
- isempty(Workqueue) && return
- t = popfirst!(Workqueue)
+function trypoptask(W::StickyWorkqueue)
+ isempty(W) && return
+ t = popfirst!(W)
  if t.state != :runnable
  # assume this somehow got queued twice,
  # probably broken now, but try discarding this switch and keep going
@@ -516,25 +600,34 @@ function trypoptask()
  return t
 end
 
-@noinline function poptaskref()
+@noinline function poptaskref(W::StickyWorkqueue)
  local task
  while true
- task = trypoptask()
+ task = trypoptask(W)
  task === nothing || break
- if process_events(true) == 0
- task = trypoptask()
- task === nothing || break
- # if there are no active handles and no runnable tasks, just
- # wait for signals.
- pause()
+ if !Threads.in_threaded_loop[] && Threads.threadid() == 1
+ if process_events(true) == 0
+ task = trypoptask(W)
+ task === nothing || break
+ # if there are no active handles and no runnable tasks, just
+ # wait for signals.
+ pause()
+ end
+ else
+ if Threads.threadid() == 1
+ process_events(false)
+ end
+ ccall(:jl_gc_safepoint, Cvoid, ())
+ ccall(:jl_cpu_pause, Cvoid, ())
  end
  end
  return Ref(task)
 end
 
 
 function wait()
- reftask = poptaskref()
+ W = Workqueues[Threads.threadid()]
+ reftask = poptaskref(W)
  result = try_yieldto(ensure_rescheduled, reftask)
  process_events(false)
  # return when we come out of the queue

base/threadingconstructs.jl

@@ -68,16 +68,17 @@ function _threadsfor(iter,lbody)
  # Hack to make nested threaded loops kinda work
  if threadid() != 1 || in_threaded_loop[]
  # We are in a nested threaded loop
- threadsfor_fun(true)
+ Base.invokelatest(threadsfor_fun, true)
  else
  in_threaded_loop[] = true
  # the ccall is not expected to throw
- ccall(:jl_threading_run, Ref{Cvoid}, (Any,), threadsfor_fun)
+ ccall(:jl_threading_run, Cvoid, (Any,), threadsfor_fun)
  in_threaded_loop[] = false
  end
  nothing
  end
 end
+
 """
  Threads.@threads
 
@@ -96,7 +97,7 @@ macro threads(args...)
  throw(ArgumentError("need an expression argument to @threads"))
  end
  if ex.head === :for
- return _threadsfor(ex.args[1],ex.args[2])
+ return _threadsfor(ex.args[1], ex.args[2])
  else
  throw(ArgumentError("unrecognized argument to @threads"))
  end

src/gc-debug.c

@@ -578,11 +578,13 @@ static void gc_scrub_task(jl_task_t *ta)
 {
  int16_t tid = ta->tid;
  jl_ptls_t ptls = jl_get_ptls_states();
- jl_ptls_t ptls2 = jl_all_tls_states[tid];
+ jl_ptls_t ptls2 = NULL;
+ if (tid != -1)
+ ptls2 = jl_all_tls_states[tid];
 
  char *low;
  char *high;
- if (ta->copy_stack && ta == ptls2->current_task) {
+ if (ta->copy_stack && ptls2 && ta == ptls2->current_task) {
  low = (char*)ptls2->stackbase - ptls2->stacksize;
  high = (char*)ptls2->stackbase;
  }
@@ -593,7 +595,7 @@ static void gc_scrub_task(jl_task_t *ta)
  else
  return;
 
- if (ptls == ptls2 && ta == ptls2->current_task) {
+ if (ptls == ptls2 && ptls2 && ta == ptls2->current_task) {
  // scan up to current `sp` for current thread and task
  low = (char*)jl_get_frame_addr();
  }

src/gc.c

@@ -2331,7 +2331,9 @@ mark: {
  gc_scrub_record_task(ta);
  void *stkbuf = ta->stkbuf;
  int16_t tid = ta->tid;
- jl_ptls_t ptls2 = jl_all_tls_states[tid];
+ jl_ptls_t ptls2 = NULL;
+ if (tid != -1)
+ ptls2 = jl_all_tls_states[tid];
  if (gc_cblist_task_scanner) {
  export_gc_state(ptls, &sp);
  gc_invoke_callbacks(jl_gc_cb_task_scanner_t,
@@ -2347,7 +2349,7 @@ mark: {
  uintptr_t offset = 0;
  uintptr_t lb = 0;
  uintptr_t ub = (uintptr_t)-1;
- if (ta == ptls2->current_task) {
+ if (ptls2 && ta == ptls2->current_task) {
  s = ptls2->pgcstack;
  }
  else if (stkbuf) {

src/init.c

@@ -761,8 +761,6 @@ void _julia_init(JL_IMAGE_SEARCH rel)
 
  jl_init_codegen();
 
- jl_start_threads();
-
  jl_an_empty_vec_any = (jl_value_t*)jl_alloc_vec_any(0);
  jl_init_serializer();
  jl_init_intrinsic_properties();
@@ -818,7 +816,16 @@ void _julia_init(JL_IMAGE_SEARCH rel)
  // it does "using Base" if Base is available.
  if (jl_base_module != NULL) {
  jl_add_standard_imports(jl_main_module);
+ // Do initialization needed before starting child threads
+ jl_value_t *f = jl_get_global(jl_base_module, jl_symbol("__preinit_threads__"));
+ if (f) {
+ size_t last_age = ptls->world_age;
+ ptls->world_age = jl_get_world_counter();
+ jl_apply(&f, 1);
+ ptls->world_age = last_age;
+ }
  }
+ jl_start_threads();
 
  // This needs to be after jl_start_threads
  if (jl_options.handle_signals == JL_OPTIONS_HANDLE_SIGNALS_ON)