diff --git a/src/Init/System/CancelToken.lean b/src/Init/System/CancelToken.lean index 78d19bd6d8..6dfe26590b 100644 --- a/src/Init/System/CancelToken.lean +++ b/src/Init/System/CancelToken.lean @@ -16,16 +16,24 @@ namespace IO /-- Cancellation token for cooperative task cancellation: request cancellation with -`CancelToken.set` and check for it with `CancelToken.isSet`. +`CancelToken.set` and check for it with `CancelToken.isSet`. To react to cancellation without +polling, register a callback with `CancelToken.onSet`. This is a more flexible alternative to `Task.cancel` as the token can be shared between multiple -tasks. The underlying `Promise` allows waiting for cancellation via `CancelToken.task` without -polling. +tasks. + +Note: the underlying `Promise` and the task it produces are kept private. If a future change +exposes the underlying task (or the promise) publicly, the `set` implementation must be +audited for races between resolving the promise and writing `setRef`: callers observing +`isSet = true` and then waiting on the task currently rely on the order chosen in `set`, and +that ordering interacts with the cheap `Bool` fast path read by the C++ runtime. -/ structure CancelToken where private promise : IO.Promise Unit - /-- Plain `Bool` flag set in lockstep with `promise`; read by the cheap `isSet` fast path - (e.g. on the hot `Core.checkInterrupted` C++ path) without walking the promise. -/ + /-- Plain `Bool` flag set after `promise.resolve` so that observing `isSet = true` implies the + promise has already fired (and thus any `onSet` callback has already run). Read by the cheap + `isSet` fast path (e.g. on the hot `Core.checkInterrupted` C++ path) without walking the + promise. -/ private setRef : IO.Ref Bool deriving Nonempty @@ -42,31 +50,26 @@ def new : BaseIO CancelToken := do let setRef ← IO.mkRef false return { promise, setRef } -/-- Activates a cancellation token. Idempotent. -/ +/-- +Activates a cancellation token. Idempotent. + +Resolves the underlying promise *before* setting the `Bool` flag, so that observing `isSet = true` +implies that any synchronously-chained `onSet` callbacks have already run. The reverse implication +does not hold: a callback running synchronously inside `set` may briefly observe `isSet = false`. +-/ def set (tk : CancelToken) : BaseIO Unit := do - tk.setRef.set true tk.promise.resolve () + tk.setRef.set true /-- Checks whether the cancellation token has been activated. -/ def isSet (tk : CancelToken) : BaseIO Bool := tk.setRef.get /-- -A task that completes when the cancellation token is set or dropped. Useful for waiting on -cancellation without polling — e.g. as one of the tasks given to `IO.waitAny`. - -Returns the underlying promise's `result?` task directly, so the same task object is returned -on every call and can safely have further dependencies attached (`Task.map`, `BaseIO.bindTask`). -The `Option` distinguishes a normal `set` (`some ()`) from the token being dropped without ever -being set (`none`). For attaching a callback, prefer `onSet`. --/ -def task (tk : CancelToken) : Task (Option Unit) := - tk.promise.result? - -/-- -Registers a callback to run when the cancellation token is set. The callback runs as a -synchronous task dependency, so it executes inline on the thread that calls `set`. If the -token is already set when `onSet` is called, the callback runs immediately. +Registers a callback to run when the cancellation token is set or dropped. The callback runs as a +synchronous task dependency, so it executes inline on the thread that calls `set` (or on the +finalizer thread if the token is dropped). If the token is already set when `onSet` is called, +the callback runs immediately. -/ def onSet (tk : CancelToken) (action : BaseIO Unit) : BaseIO Unit := BaseIO.chainTask tk.promise.result? (sync := true) fun _ => action diff --git a/tests/elab/cancel_token.lean b/tests/elab/cancel_token.lean index 447fa5c8ca..9b120908df 100644 --- a/tests/elab/cancel_token.lean +++ b/tests/elab/cancel_token.lean @@ -1,4 +1,4 @@ -/-! Tests for `IO.CancelToken.onSet` and `IO.CancelToken.task`. -/ +/-! Tests for `IO.CancelToken.onSet`. -/ /-- `onSet` registered before `set` fires synchronously when `set` is called. -/ def testOnSetBeforeSet : IO Unit := do @@ -30,36 +30,13 @@ def testOnSetMultiple : IO Unit := do if n != 5 then panic! s!"expected 5 callbacks, got {n}" IO.println "ok" -/-- `task` returns a task that resolves once the token is set, with `some ()`. -/ -def testTask : IO Unit := do - let tk ← IO.CancelToken.new - let t := tk.task - tk.set - match t.get with - | some () => IO.println "ok" - | none => panic! "task fired with none after set" - -/-- `task` resolves to `none` when the token is dropped without being set. -/ -def testTaskDropped : IO Unit := do - let t ← do - let tk ← IO.CancelToken.new - pure tk.task - -- `tk` is now out of scope; the underlying promise is dropped. - match t.get with - | none => IO.println "ok" - | some () => panic! "task fired with some on dropped token" - /-- info: ok ok ok -ok -ok -/ #guard_msgs in #eval do testOnSetBeforeSet testOnSetAfterSet testOnSetMultiple - testTask - testTaskDropped diff --git a/tests/elab/delabConst.lean b/tests/elab/delabConst.lean index 47d5bac2f5..7e7fc41052 100644 --- a/tests/elab/delabConst.lean +++ b/tests/elab/delabConst.lean @@ -224,8 +224,8 @@ The name `IO.CancelToken.promise✝` is a private imported name. /-- info: def IO.CancelToken.set : IO.CancelToken → BaseIO Unit := fun tk => do - ST.Ref.set (IO.CancelToken.setRef✝ tk) true IO.Promise.resolve () (IO.CancelToken.promise✝ tk) + ST.Ref.set (IO.CancelToken.setRef✝ tk) true -/ #guard_msgs in #print IO.CancelToken.set /-! @@ -234,7 +234,7 @@ Even if `IO` is opened, it won't print as `CancelToken.promise✝`, but the full /-- info: def IO.CancelToken.set : CancelToken → BaseIO Unit := fun tk => do - ST.Ref.set (IO.CancelToken.setRef✝ tk) true Promise.resolve () (IO.CancelToken.promise✝ tk) + ST.Ref.set (IO.CancelToken.setRef✝ tk) true -/ #guard_msgs in open IO in #print IO.CancelToken.set