Working with IO.CancelToken

Stream: lean4

[Eric Wieser]

Here's some code that ought to cancel a tactic after a certain duration:

import Lean

set_option linter.all false

open Lean Elab Tactic

/-! versions of try/catch that catch `interrupted` too  -/
section catch_interrupted
attribute [-instance]
  Lean.instMonadExceptOfExceptionCoreM Lean.Elab.Tactic.instMonadExceptExceptionTacticM
variable {α}

def Meta.tryCatchAll (m : MetaM α) (h : Exception → MetaM α) : MetaM α := tryCatch m h
def Term.tryCatchAll (m : TermElabM α) (h : Exception → TermElabM α) : TermElabM α := tryCatch m h
def Tactic.tryCatchAll (x : TacticM α) (h : Exception → TacticM α) : TacticM α := do
  let b ← saveState
  try x catch ex => b.restore; h ex

end catch_interrupted

elab "with_timeout " ms:num "=>" tac:tacticSeq : tactic => do
  let ms := ms.getNat.toUInt32
  let tk ← IO.CancelToken.new
  let t ← withTheReader Core.Context (fun s => { s with cancelTk? := some tk }) do
    let t ← IO.asTask do
      IO.println "Task is running"
      IO.sleep ms
      tk.set
      IO.println "Task is cancelled"
    Tactic.tryCatchAll
      (evalTactic tac)
      (fun e =>
        Lean.logInfo e.toMessageData)
      return t
  IO.cancel t

theorem foo : True := by
  with_timeout 500 =>
    sleep 1000
    run_tac Lean.logInfo "I wasn't cancelled!"
    trivial

theorem bar : True := by
  with_timeout 500 =>
    sleep 100
    run_tac Lean.logInfo "I wasn't cancelled!"
    trivial

#print foo
#print bar

This works when run with lake lean this_file.lean, but in vscode the server stops responding. Is this use of a custom cancel token legal?

[Sebastian Ullrich]

You're writing invalid data to the language server output. Try IO.eprintln

[Sebastian Ullrich]

(We redirect stdout writes during elaboration to diagnostics but it's not clear how to do that in arbitrary tasks. Perhaps we should at least redirect them to stderr.)

[Eric Wieser]

Can I use control to somehow stay in CoreM in the task?

[Sebastian Ullrich]

I'm in the process of developing API for this in context of parallel elaboration, but here is a sneak peek that may already be helpful

/-- Runs the given action in a separate task, discarding its final state. -/
def runAsync (act : TermElabM α) : TermElabM (Task (Except Exception α)) := do
  let mut coreSt ← getThe Core.State
  let metaSt ← getThe Meta.State
  let st ← get
  let coreCtx ← readThe Core.Context
  let metaCtx ← readThe Meta.Context
  let ctx ← read
  let heartbeats := (← IO.getNumHeartbeats) - coreCtx.initHeartbeats
  withCurrHeartbeats (do
      IO.addHeartbeats heartbeats.toUInt64
      act : TermElabM _)
    |>.run' ctx st
    |>.run' metaCtx metaSt
    |>.run' coreCtx coreSt
    |> EIO.asTask

[Eric Wieser]

I assume

def runAsync' (act : TermElabM α) : TermElabM (Task (Except Exception α)) := do
  controlAt (EIO Exception) fun run => EIO.asTask (run act)

doesn't behave correctly?

[Sebastian Ullrich]

I believe it doesn't because it would share the state IO.Ref across threads, sounds like a potential source of extremely fun bugs!

[Leni Aniva]

so does this work?