Single-threaded IO timeout

Stream: lean4

[Jannis Limperg]

I'm trying to build a timeout mechanism for IO computations, but it seems to work only with >= 3 threads. MWE:

-- Test.lean
import Std

open Std.Internal.IO.Async (sleep) in
def withTimeoutCore {α} (timeout : Std.Time.Millisecond.Offset)
    (cancelTk : IO.CancelToken) (x : IO α) : IO α := do
  let task ← x.asTask
  let cancelTask ← (← sleep timeout).mapIO fun _ => do
    if ! (← IO.checkCanceled) then
      cancelTk.set
    throw <| IO.userError "timed out"
  let result?? ← IO.waitAny [task, cancelTask]
  IO.cancel cancelTask
  EIO.ofExcept result??

#eval show IO Unit from do
  let cancelTk ← IO.CancelToken.new
  withTimeoutCore 1000 cancelTk do
    while true do
      IO.sleep 5

$ lake env lean -j1 Test.lean # runs forever
$ lake env lean -j2 Test.lean # runs forever
$ lake env lean -j3 Test.lean
Test.lean:15:0: error: timed out

I don't really understand Lean's concurrency features well, so I'd be very grateful if someone could explain what's going on here. :thank_you: cc @Henrik Böving @Sofia Rodrigues since much of the async code seems to be from you.

[Jannis Limperg]

(Tested on v4.20.0 in case that matters.)

[Jannis Limperg]

On v4.24.0-rc1, all three invocations seem to hit an unreachable!.

[Jannis Limperg]

This version works on v4.24.0-rc1 with >= 2 threads, but not with 1:

import Std

open Std.Internal.IO.Async (sleep) in
def withTimeoutCore {α} (timeout : Std.Time.Millisecond.Offset) (x : IO α) : IO α := do
  let task ← x.asTask
  let cancelTask ← (← sleep timeout).mapIO fun _ => do
    if ! (← IO.checkCanceled) then
      IO.cancel task
    throw <| IO.userError "timed out"
  let result?? ← IO.waitAny [task, cancelTask]
  IO.cancel cancelTask
  EIO.ofExcept result??

#eval show IO Unit from do
  withTimeoutCore 1000 do
    while ! (← IO.checkCanceled) do
      IO.sleep 5

[Sofia Rodrigues]

Currently, the Async module works similarly to the Node.js event loop and execution model. If you run an infinite blocking while loop and try to create a setInterval in JS, the setInterval callback won’t execute because Node.js is single-threaded. So, you should avoid doing blocking stuff like infinite loops. In our case, we can summon multiple blocking tasks at the same time, but the limit depends on how many cores you have (Because Lean summons a fixed amount of threads to deal with Tasks). I think that right now, we spawn one thread per core. That’s probably the issue with the first code. With one thread, the main thread is running, and the other one will not be summoned. With two threads, both Tasks will be executed, but the third one with the timeout will not be executed because two tasks are blocking both threads.

The IO.CancelToken isn’t something you use with raw tasks, I made it to work with the Async module, which summons multiple short lived tasks, so we can’t really control them directly with the current Task cancellation operations. Instead, you use IO.checkCancelled and IO.cancel, because a Task can’t be cancelled in the middle of execution unless you explicitly check for it.

This is the correct way to use these things:

import Std

open Std.Internal.IO.Async (sleep) in
def withTimeoutCore {α} (timeout : Std.Time.Millisecond.Offset) (x : IO α) : IO α := do
  let task ← x.asTask

  let cancelTask ← (← sleep timeout).mapIO fun _ => do
    IO.cancel task
    throw <| IO.userError "timed out"

  let result?? ← IO.waitAny [task, cancelTask]
  IO.cancel cancelTask

  EIO.ofExcept result??

def main : IO Unit := do
  withTimeoutCore 1000 do
    while true do

      if ← IO.checkCanceled then
        break

      IO.println "running"
      IO.sleep 100

[Sofia Rodrigues]

About the last code that u sent, it’s not working with a single thread for the same reason I mentioned earlier about how many threads we have available for Tasks.

[Jannis Limperg]

I see, thank you for the thorough explanation! So I guess the question is which operations can be used to yield control to other async tasks. Python makes this explicit with await; I'm not familiar with JS's system. Is there something similar in Lean, or a set of operations that I can expect to yield?

In any case, the real use case is cancelling a CoreM computation. (Hence the weird CancelToken in my first code.) I guess I shouldn't expect Lean functions to yield frequently, so a single-threaded timeout seems out of the question. I guess my best bet is to give Lean a bunch of threads and maybe disable async elaboration, to maximise the chance that the timeout thread actually gets to run? Or maybe a signal-based timeout, but afaict Lean doesn't let me set signal handlers atm.

[Sebastian Ullrich]

You'll want to use (prio := .dedicated) when starting any task that should not be enqueued in the thread pool, such as for your timeout task