refactor: keep IO.CancelToken task private, resolve promise before setting flag (#13569)

This PR addresses two review points on `IO.CancelToken`:

* `set` now resolves the underlying promise *before* writing the `Bool`
  fast-path flag, so observing `isSet = true` implies any synchronously
chained `onSet` callback has already run. The previous order (flag
first,
then resolve) was a subtle footgun: code seeing `isSet = true` could not
  rely on the cancellation task having fired.
* The underlying promise and the task it produces are kept private. The
prior `task : Task (Option Unit)` accessor is removed; consumers should
use `onSet` to react to cancellation. A comment on the structure records
  that re-exposing the task in the future requires re-auditing the order
  in `set` for races between the promise and the `Bool` flag.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>

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

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

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