How do command cancellation tokens work?

Stream: lean4

[𝚠𝚘𝚓𝚌𝚒𝚎𝚌𝚑 𝚗𝚊𝚠𝚛𝚘𝚌𝚔𝚒]

I am trying to figure out how to respond to cancellation in custom elaborators and getting increasingly more confused. In the following example, it is very easy to get into a state where both commands are commented out, but the computations never get cancelled and keep going forever. This is on 4.28.0-rc1.

To try it out, uncomment the commands and comment them back. Most of the time, a runaway computation of a continues. Note this probably has to be done locally since the playground does not show dbg_trace outputs.

Is it not a guarantee of Lean that the cancel token will, in fact, be set on commands that are no longer needed, making the implementation unsound? If so, what is the right way to ensure a dedicated thread gets cancelled?

(This thread is related.)

import Lean

open Lean Elab Command

def printPtrAddr {α : Type u} {m : Type → Type} [Monad m] (a : α) (f : USize → String) : m Unit :=
  withPtrAddr a (fun addr => dbg_trace (f addr); return ()) (fun _ _ => by simp [dbgTrace])

elab "#test" n:ident : command => do
  let some tk := (← readThe Elab.Command.Context).cancelTk?
    | throwError "no cancel token"
  discard <| BaseIO.asTask (prio := .dedicated) do
    printPtrAddr tk (fun addr => s!"[{n} @ {addr}] starting..")
    repeat do
      printPtrAddr tk (fun addr => s!"[{n} @ {addr}] tick")
      if ← tk.isSet then
        printPtrAddr tk (fun addr => s!"[{n} @ {addr}] interrupted")
        return
      IO.sleep 1000

-- #test a
-- #test b

[𝚠𝚘𝚓𝚌𝚒𝚎𝚌𝚑 𝚗𝚊𝚠𝚛𝚘𝚌𝚔𝚒]

Ok, I think the answer is: this elaborator launches the thread and finishes immediately, so from the perspective of the language processor there is no elaboration task to be cancelled.

[𝚠𝚘𝚓𝚌𝚒𝚎𝚌𝚑 𝚗𝚊𝚠𝚛𝚘𝚌𝚔𝚒]

Hm, using logSnapshotTask doesn't seem to fix it. Maybe I am misusing the API?

set_option stderrAsMessages false
def printPtrAddr {α : Type u} (a : α) (f : USize → String) : BaseIO Unit :=
  withPtrAddr a (fun addr => dbg_trace (f addr); return ()) (fun _ _ => by simp [dbgTrace])

elab "#test" n:ident : command => do
  let cancelTk ← IO.CancelToken.new
  let runBg ← wrapAsyncAsSnapshot (cancelTk? := cancelTk) fun () => do
    printPtrAddr cancelTk (fun addr => s!"[{n} @ {addr}] starting..")
    repeat do
      printPtrAddr cancelTk (fun addr => s!"[{n} @ {addr}] tick")
      if ← cancelTk.isSet then
        printPtrAddr cancelTk (fun addr => s!"[{n} @ {addr}] interrupted")
        return
      IO.sleep 1000

  let task ← BaseIO.asTask (prio := .dedicated) (runBg ())

  logSnapshotTask { stx? := none, task, cancelTk? := cancelTk }

[Sebastian Ullrich]

@[incremental] fixes it - without it, no snapshot tree will be created. But this is clearly surprising, it may be possible to adjust the API so that use of logSnapshotTask is sufficient as these ought to be orthogonal concerns

[𝚠𝚘𝚓𝚌𝚒𝚎𝚌𝚑 𝚗𝚊𝚠𝚛𝚘𝚌𝚔𝚒]

It does not seem to fix it, in the sense that @[incremental] elab "#test" n:ident : command => do still loops.

[Sebastian Ullrich]

Ah, I must have test only changes before the command, not within it. For actual incrementality we need to fill Command.Context.snap? directly and make sure to apply the attribute to the elaborator, not the parser definition. But as said ideally it should not actually be necessary to define an incremental elaborator here.

syntax "#test" ident : command
@[incremental]
elab_rules : command | `(#test $n:ident) => do
  let snap? := (← read).snap?
  if let some snap := snap?.bind (·.old?) then
    snap.val.cancelRec
  let cancelTk ← IO.CancelToken.new
  let runBg ← wrapAsyncAsSnapshot (cancelTk? := cancelTk) fun () => do
    printPtrAddr cancelTk (fun addr => s!"[{n} @ {addr}] starting..")
    repeat do
      printPtrAddr cancelTk (fun addr => s!"[{n} @ {addr}] tick")
      if ← cancelTk.isSet then
        printPtrAddr cancelTk (fun addr => s!"[{n} @ {addr}] interrupted")
        return
      IO.sleep 1000

  let task ← BaseIO.asTask (prio := .dedicated) (runBg ())

  if let some prom := snap?.map (·.new) then
    prom.resolve (.ofTyped { element.diagnostics := .empty, children := #[{ stx? := none, task, cancelTk? := cancelTk }] : Language.SnapshotTree })

[𝚠𝚘𝚓𝚌𝚒𝚎𝚌𝚑 𝚗𝚊𝚠𝚛𝚘𝚌𝚔𝚒]

Hi again! In light of lean4#12553 I just tried my most recently posted snippet again (intentionally without @[incremental] since that's the case the PR addresses) on 4.29.0-rc2, but it continues to loop. I am guessing the PR was not meant to address this case?

The code using snap? directly does work, but it does require access to the CommandElabM context. Is there a way to obtain the correct behavior for tasks launched in CoreM (and only using the CoreM context/state, as opposed to threading through the snap? data)?