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) =
@@ -128,7 +141,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=false)
                             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)