cancellation: Add cancel token support

Author: jpsamaroo

src/Dagger.jl

@@ -57,16 +57,16 @@ include("processor.jl")
 include("threadproc.jl")
 include("context.jl")
 include("utils/processors.jl")
+include("dtask.jl")
+include("cancellation.jl")
 include("task-tls.jl")
 include("scopes.jl")
 include("utils/scopes.jl")
-include("dtask.jl")
 include("queue.jl")
 include("thunk.jl")
 include("submission.jl")
 include("chunks.jl")
 include("memory-spaces.jl")
-include("cancellation.jl")
 
 # Task scheduling
 include("compute.jl")

src/cancellation.jl

@@ -1,3 +1,38 @@
+# DTask-level cancellation
+
+struct CancelToken
+    cancelled::Base.RefValue{Bool}
+    event::Base.Event
+end
+CancelToken() = CancelToken(Ref(false), Base.Event())
+function cancel!(token::CancelToken)
+    token.cancelled[] = true
+    notify(token.event)
+    return
+end
+is_cancelled(token::CancelToken) = token.cancelled[]
+Base.wait(token::CancelToken) = wait(token.event)
+# TODO: Enable this for safety
+#Serialization.serialize(io::AbstractSerializer, ::CancelToken) =
+#    throw(ConcurrencyViolationError("Cannot serialize a CancelToken"))
+
+const DTASK_CANCEL_TOKEN = TaskLocalValue{Union{CancelToken,Nothing}}(()->nothing)
+
+function clone_cancel_token_remote(orig_token::CancelToken, wid::Integer)
+    remote_token = remotecall_fetch(wid) do
+        return poolset(CancelToken())
+    end
+    errormonitor_tracked("remote cancel_token communicator", Threads.@spawn begin
+        wait(orig_token)
+        @dagdebug nothing :cancel "Cancelling remote token on worker $wid"
+        MemPool.access_ref(remote_token) do remote_token
+            cancel!(remote_token)
+        end
+    end)
+end
+
+# Global-level cancellation
+
 """
     cancel!(task::DTask; force::Bool=false, halt_sch::Bool=false)
 
@@ -80,11 +115,11 @@ function _cancel!(state, tid, force, halt_sch)
                             Tf === typeof(Sch.eager_thunk) && continue
                             istaskdone(task) && continue
                             any_cancelled = true
-                            @dagdebug tid :cancel "Cancelling running task ($Tf)"
                             if force
                                 @dagdebug tid :cancel "Interrupting running task ($Tf)"
                                 Threads.@spawn Base.throwto(task, InterruptException())
                             else
+                                @dagdebug tid :cancel "Cancelling running task ($Tf)"
                                 # Tell the processor to just drop this task
                                 task_occupancy = task_spec[4]
                                 time_util = task_spec[2]
@@ -93,6 +128,7 @@ function _cancel!(state, tid, force, halt_sch)
                                 push!(istate.cancelled, tid)
                                 to_proc = istate.proc
                                 put!(istate.return_queue, (myid(), to_proc, tid, (InterruptException(), nothing)))
+                                cancel!(istate.cancel_tokens[tid])
                             end
                         end
                     end

src/sch/Sch.jl

@@ -1183,6 +1183,7 @@ struct ProcessorInternalState
     proc_occupancy::Base.RefValue{UInt32}
     time_pressure::Base.RefValue{UInt64}
     cancelled::Set{Int}
+    cancel_tokens::Dict{Int,Dagger.CancelToken}
     done::Base.RefValue{Bool}
 end
 struct ProcessorState
@@ -1332,7 +1333,14 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
 
             # Execute the task and return its result
             t = @task begin
+                # Set up cancellation
+                cancel_token = Dagger.CancelToken()
+                Dagger.DTASK_CANCEL_TOKEN[] = cancel_token
+                lock(istate.queue) do _
+                    istate.cancel_tokens[thunk_id] = cancel_token
+                end
                 was_cancelled = false
+
                 result = try
                     do_task(to_proc, task)
                 catch err
@@ -1349,6 +1357,7 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
                             # Task was cancelled, so occupancy and pressure are
                             # already reduced
                             pop!(istate.cancelled, thunk_id)
+                            delete!(istate.cancel_tokens, thunk_id)
                             was_cancelled = true
                         end
                     end
@@ -1366,6 +1375,9 @@ function start_processor_runner!(istate::ProcessorInternalState, uid::UInt64, re
                     else
                         rethrow(err)
                     end
+                finally
+                    # Ensure that any spawned tasks get cleaned up
+                    Dagger.cancel!(cancel_token)
                 end
             end
             lock(istate.queue) do _
@@ -1415,6 +1427,7 @@ function do_tasks(to_proc, return_queue, tasks)
                                             Dict{Int,Vector{Any}}(),
                                             Ref(UInt32(0)), Ref(UInt64(0)),
                                             Set{Int}(),
+                                            Dict{Int,Dagger.CancelToken}(),
                                             Ref(false))
             runner = start_processor_runner!(istate, uid, return_queue)
             @static if VERSION < v"1.9"
@@ -1656,6 +1669,7 @@ function do_task(to_proc, task_desc)
             sch_handle,
             processor=to_proc,
             task_spec=task_desc,
+            cancel_token=Dagger.DTASK_CANCEL_TOKEN[],
         ))
 
         res = Dagger.with_options(propagated) do

src/task-tls.jl

@@ -5,6 +5,7 @@ struct DTaskTLS
     sch_uid::UInt
     sch_handle::Any # FIXME: SchedulerHandle
     task_spec::Vector{Any} # FIXME: TaskSpec
+    cancel_token::CancelToken
 end
 
 const DTASK_TLS = TaskLocalValue{Union{DTaskTLS,Nothing}}(()->nothing)
@@ -22,7 +23,7 @@ get_tls() = DTASK_TLS[]::DTaskTLS
 Sets all Dagger TLS variables from `tls`, which may be a `DTaskTLS` or a `NamedTuple`.
 """
 function set_tls!(tls)
-    DTASK_TLS[] = DTaskTLS(tls.processor, tls.sch_uid, tls.sch_handle, tls.task_spec)
+    DTASK_TLS[] = DTaskTLS(tls.processor, tls.sch_uid, tls.sch_handle, tls.task_spec, tls.cancel_token)
 end
 
 """
@@ -40,3 +41,21 @@ Get the current processor executing the current [`DTask`](@ref).
 """
 task_processor() = get_tls().processor
 @deprecate thunk_processor() task_processor()
+
+"""
+    task_cancelled() -> Bool
+
+Returns `true` if the current [`DTask`](@ref) has been cancelled, else `false`.
+"""
+task_cancelled() = get_tls().cancel_token.cancelled[]
+
+"""
+    task_may_cancel!()
+
+Throws an `InterruptException` if the current [`DTask`](@ref) has been cancelled.
+"""
+function task_may_cancel!()
+    if task_cancelled()
+        throw(InterruptException())
+    end
+end

src/threadproc.jl

@@ -27,8 +27,9 @@ function execute!(proc::ThreadProc, @nospecialize(f), @nospecialize(args...); @n
         return result[]
     catch err
         if err isa InterruptException
+            # Direct interrupt hit us, propagate cancellation signal
+            # FIXME: We should tell the scheduler that the user hit Ctrl-C
             if !istaskdone(task)
-                # Propagate cancellation signal
                 Threads.@spawn Base.throwto(task, InterruptException())
             end
         end