cancellation: Add graceful vs. forced

Author: jpsamaroo

src/cancellation.jl

@@ -1,16 +1,29 @@
 # DTask-level cancellation
 
-struct CancelToken
-    cancelled::Base.RefValue{Bool}
+mutable struct CancelToken
+    @atomic cancelled::Bool
+    @atomic graceful::Bool
     event::Base.Event
 end
-CancelToken() = CancelToken(Ref(false), Base.Event())
-function cancel!(token::CancelToken)
-    token.cancelled[] = true
+CancelToken() = CancelToken(false, false, Base.Event())
+function cancel!(token::CancelToken; graceful::Bool=true)
+    if !graceful
+        @atomic token.graceful = false
+    end
+    @atomic token.cancelled = true
     notify(token.event)
     return
 end
-is_cancelled(token::CancelToken) = token.cancelled[]
+function is_cancelled(token::CancelToken; must_force::Bool=false)
+    if token.cancelled[]
+        if must_force && token.graceful[]
+            # If we're only responding to forced cancellation, ignore graceful cancellations
+            return false
+        end
+        return true
+    end
+    return false
+end
 Base.wait(token::CancelToken) = wait(token.event)
 # TODO: Enable this for safety
 #Serialization.serialize(io::AbstractSerializer, ::CancelToken) =
@@ -34,13 +47,15 @@ end
 # Global-level cancellation
 
 """
-    cancel!(task::DTask; force::Bool=false, halt_sch::Bool=false)
+    cancel!(task::DTask; force::Bool=false, graceful::Bool=true, halt_sch::Bool=false)
 
 Cancels `task` at any point in its lifecycle, causing the scheduler to abandon
-it. If `force` is `true`, the task will be interrupted with an
-`InterruptException` (not recommended, this is unsafe). If `halt_sch` is
-`true`, the scheduler will be halted after the task is cancelled (it will
-restart automatically upon the next `@spawn`/`spawn` call).
+it.
+
+# Keyword arguments
+- `force`: If `true`, the task will be interrupted with an `InterruptException` (not recommended, this is unsafe).
+- `graceful`: If `true`, the task will be allowed to finish its current execution before being cancelled; otherwise, it will be cancelled as soon as possible.
+- `halt_sch`: If `true`, the scheduler will be halted after the task is cancelled (it will restart automatically upon the next `@spawn`/`spawn` call).
 
 As an example, the following code will cancel task `t` before it finishes
 executing:
@@ -56,24 +71,24 @@ tasks which are waiting to run. Using `cancel!` is generally a much safer
 alternative to Ctrl+C, as it cooperates with the scheduler and runtime and
 avoids unintended side effects.
 """
-function cancel!(task::DTask; force::Bool=false, halt_sch::Bool=false)
+function cancel!(task::DTask; force::Bool=false, graceful::Bool=true, halt_sch::Bool=false)
     tid = lock(Dagger.Sch.EAGER_ID_MAP) do id_map
         id_map[task.uid]
     end
-    cancel!(tid; force, halt_sch)
+    cancel!(tid; force, graceful, halt_sch)
 end
 function cancel!(tid::Union{Int,Nothing}=nothing;
-                 force::Bool=false, halt_sch::Bool=false)
+                 force::Bool=false, graceful::Bool=true, halt_sch::Bool=false)
     remotecall_fetch(1, tid, force, halt_sch) do tid, force, halt_sch
         state = Sch.EAGER_STATE[]
 
         # Check that the scheduler isn't stopping or has already stopped
         if !isnothing(state) && !state.halt.set
-            @lock state.lock _cancel!(state, tid, force, halt_sch)
+            @lock state.lock _cancel!(state, tid, force, graceful, halt_sch)
         end
     end
 end
-function _cancel!(state, tid, force, halt_sch)
+function _cancel!(state, tid, force, graceful, halt_sch)
     @assert islocked(state.lock)
 
     # Get the scheduler uid
@@ -128,7 +143,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])
+                                cancel!(istate.cancel_tokens[tid]; graceful)
                             end
                         end
                     end

src/task-tls.jl

@@ -52,26 +52,30 @@ task_processor() = get_tls().processor
 @deprecate(thunk_processor(), task_processor())
 
 """
-    task_cancelled() -> Bool
+    task_cancelled(; must_force::Bool=false) -> Bool
 
 Returns `true` if the current [`DTask`](@ref) has been cancelled, else `false`.
+If `must_force=true`, then only return `true` if the cancellation was forced.
 """
-task_cancelled() = is_cancelled(get_tls().cancel_token)
+task_cancelled(; must_force::Bool=false) =
+    is_cancelled(get_tls().cancel_token; must_force)
 
 """
-    task_may_cancel!()
+    task_may_cancel!(; must_force::Bool=false)
 
 Throws an `InterruptException` if the current [`DTask`](@ref) has been cancelled.
+If `must_force=true`, then only throw if the cancellation was forced.
 """
-function task_may_cancel!()
-    if task_cancelled()
+function task_may_cancel!(;must_force::Bool=false)
+    if task_cancelled(;must_force)
         throw(InterruptException())
     end
 end
 
 """
-    task_cancel!()
+    task_cancel!(; graceful::Bool=true)
 
-Cancels the current [`DTask`](@ref).
+Cancels the current [`DTask`](@ref). If `graceful=true`, then the task will be
+cancelled gracefully, otherwise it will be forced.
 """
-task_cancel!() = cancel!(get_tls().cancel_token)
+task_cancel!(; graceful::Bool=true) = cancel!(get_tls().cancel_token; graceful)