Skip to content

Conversation

@algebraic-dev
Copy link
Member

@algebraic-dev algebraic-dev commented Apr 17, 2025

This PR adds a new monadic interface for Async operations.

This is the design for the Async monad that I liked the most. The idea was refined with the help of @tydeu. Before that, I had some prerequisites in mind:

  1. Good performance
  2. Explicit yield points, so we could avoid using bindTask for every lifted IO operation
  3. A way to avoid creating an infinite chain of Tasks during recursion

The 2 and 3 points are not covered in this PR, I wish I had a good solution but right now only a few sketches of this.

Explicit yield points

I thought this would be easy at first, but it actually turned out kinda tricky. I ended up creating the suspend syntax, which is just a small modification of the lift method (<- ...) syntax. It desugars to Suspend.suspend task fun _ => .... So something like:

do
  IO.println "a"
  IO.println "b"
  let result := suspend (client.recv? 1024)
  IO.println "c"
  IO.println "d"

Would become:

Bind.bind (IO.println "a") fun _ =>
Bind.bind (IO.println "b") fun _ =>
Suspend.suspend (client.recv? 1024) fun message =>
  Bind.bind (IO.println "c") fun _ =>
  IO.println "d"

This makes things a bit more efficient. When using bind, we would try to avoid creating a Task chain, and the suspend would be the only place we use Task.bind. But there's a problem if we use bind with something that needs suspend, it’ll block the whole task. Blocking is the only way to prevent task accumulation when using plain bind inside a structure like that:

inductive AsyncResult (ε σ α : Type u) where
    | ok    : α → σ → AsyncResult ε σ α
    | error : ε → σ → AsyncResult ε σ α
    | ofTask  : Task (EStateM.Result ε σ α) → σ →AsyncResult ε σ α

Because we simply need to remove the ofTask and transform it into an ok.

Infinite chain of Tasks

If you create an infinite recursive function using Task (which is super common in servers like HTTP ones), it can lead to a lot of memory usage. Because those tasks get chained forever and won't be freed until the function returns.

To get around that, I used CPS and instead of just calling Task.bind, I’d spawn a new task and return an "empty" one like:

fun k => Task.bind (...) fun value => do k value; pure emptyTask

This works great with a CPS-style monad, but it generates a huge IR by itself.

Just doing CPS alone was too much, though, because every lifted operation created a new continuation and a Task.bind. So, I used it with suspend and got a better performance, but the usage is not good with suspend.

The current monad

Right now, the monad I’m using is super simple. It doesn't solve the earlier problems, but the API is clean, and the generated IR is small enough. An example of how we should use it is:

-- A loop that repeatedly sends a message and waits for a reply.
partial def writeLoop (client : Socket.Client) (message : String) : Async (AsyncTask Unit) := async do
  IO.println s!"sending: {message}"
  await (← client.send (String.toUTF8 message))

  if let some mes ← await (← client.recv? 1024) then
    IO.println s!"received: {String.fromUTF8! mes}"
    -- use parallel to avoid building up an infinite task chain
    parallel (writeLoop client message)
  else
    IO.println "client disconnected from receiving"

-- Server’s main accept loop, keeps accepting and echoing for new clients.
partial def acceptLoop (server : Socket.Server) (promise : IO.Promise Unit) : Async (AsyncTask Unit) := async do
  let client ← await (← server.accept)
  await (← client.send (String.toUTF8 "tutturu "))

  -- allow multiple clients to connect at the same time
  parallel (writeLoop client "hi!!")

  -- and keep accepting more clients, parallel again to avoid building up an infinite task chain
  parallel (acceptLoop server promise)

-- A simple client that connects and sends a message.
def echoClient (addr : SocketAddress) (message : String) : Async (AsyncTask Unit) := async do
  let socket ← Client.mk
  await (← socket.connect addr)
  parallel (writeLoop socket message)

-- TCP setup: bind, listen, serve, and run a sample client.
partial def mainTCP : Async Unit := do
  let addr := SocketAddressV4.mk (.ofParts 127 0 0 1) 8080

  let server ← Server.mk
  server.bind addr
  server.listen 128

  -- promise exists since the server is (probably) never going to stop
  let promise ← IO.Promise.new
  let acceptAction ← acceptLoop server promise

  await (← echoClient addr "hi!")
  await acceptAction
  await promise

-- Entry point
def main : IO Unit := mainTCP.wait

@github-actions github-actions bot added the toolchain-available A toolchain is available for this PR, at leanprover/lean4-pr-releases:pr-release-NNNN label Apr 17, 2025
@leanprover-community-bot
Copy link
Collaborator

leanprover-community-bot commented Apr 17, 2025

Mathlib CI status (docs):

  • ❗ Batteries/Mathlib CI will not be attempted unless your PR branches off the nightly-with-mathlib branch. Try git rebase acfc9c50d5ad32adf8868ba79ce9deee470cea4d --onto 020b8834c3e06da1cd1682431b6fa8d52206c8ec. You can force Mathlib CI using the force-mathlib-ci label. (2025-04-17 20:30:41)
  • ❗ Batteries/Mathlib CI will not be attempted unless your PR branches off the nightly-with-mathlib branch. Try git rebase 43e8288e3f073737f4b9f2283f455f0635a0f72e --onto d26d7973ad39eab976ed351255ee30f038439944. You can force Mathlib CI using the force-mathlib-ci label. (2025-05-02 21:26:07)
  • ❗ Batteries/Mathlib CI will not be attempted unless your PR branches off the nightly-with-mathlib branch. Try git rebase 7994e55d8067de8ec4b52c105df7c71341e9c190 --onto ca9b3eb75f3f794d3d526a7c37da389175b317a2. You can force Mathlib CI using the force-mathlib-ci label. (2025-05-16 13:24:08)
  • ❗ Batteries/Mathlib CI will not be attempted unless your PR branches off the nightly-with-mathlib branch. Try git rebase 7994e55d8067de8ec4b52c105df7c71341e9c190 --onto efe2ab4c04e81fe2a3edcc0d861449490b4431b2. You can force Mathlib CI using the force-mathlib-ci label. (2025-05-19 15:14:46)
  • ❗ Batteries/Mathlib CI will not be attempted unless your PR branches off the nightly-with-mathlib branch. Try git rebase 87cc330489f7eb778d6b670634fef41993d1557c --onto 3bf95e9b585ceca4e88baeda8b92a44f485cffdf. You can force Mathlib CI using the force-mathlib-ci label. (2025-05-21 09:01:10)
  • ❗ Batteries/Mathlib CI will not be attempted unless your PR branches off the nightly-with-mathlib branch. Try git rebase 87cc330489f7eb778d6b670634fef41993d1557c --onto 3af9ab64ed79a667fb8989f7a0c258b018aba371. You can force Mathlib CI using the force-mathlib-ci label. (2025-05-29 00:47:15)
  • ❗ Batteries/Mathlib CI will not be attempted unless your PR branches off the nightly-with-mathlib branch. Try git rebase 87cc330489f7eb778d6b670634fef41993d1557c --onto c12159b51982221bdd66f0c5997f85e1f9d91772. You can force Mathlib CI using the force-mathlib-ci label. (2025-06-04 22:16:44)
  • ❗ Batteries/Mathlib CI will not be attempted unless your PR branches off the nightly-with-mathlib branch. Try git rebase e7132326236c4cd5554b9de3436bbd8984aa03d7 --onto faffe863343211d03d07d4ef19d90a2aa9d6fc8c. You can force Mathlib CI using the force-mathlib-ci label. (2025-06-13 21:39:44)
  • ❗ Batteries/Mathlib CI will not be attempted unless your PR branches off the nightly-with-mathlib branch. Try git rebase e7132326236c4cd5554b9de3436bbd8984aa03d7 --onto 3b2990b3816b7e7ec289ad498496a61d18bd85f4. You can force Mathlib CI using the force-mathlib-ci label. (2025-06-18 05:58:38)
  • ❗ Batteries/Mathlib CI will not be attempted unless your PR branches off the nightly-with-mathlib branch. Try git rebase e7132326236c4cd5554b9de3436bbd8984aa03d7 --onto e83b7681407043c75f5167c2e4da13e25078a0c0. You can force Mathlib CI using the force-mathlib-ci label. (2025-06-18 18:47:25)
  • ❗ Batteries/Mathlib CI will not be attempted unless your PR branches off the nightly-with-mathlib branch. Try git rebase e7132326236c4cd5554b9de3436bbd8984aa03d7 --onto db499e96aac8ad654c8ed5ab40c4e6885d38c9a1. You can force Mathlib CI using the force-mathlib-ci label. (2025-06-26 02:56:15)

Copy link
Contributor

@hargoniX hargoniX left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Which types in this PR are the intended surface level API? Generally people seem to expect one type that can be awaited (that's presumably Async?) and one type they can use for concurrent units of work (that's presumably AsyncTask?). I guess we should probably make this sort of intent clear? Though it should hopefully become clear based on tests/examples that end up being written later on.

Based on this I do wonder how far we want to align the rest of the async API that we have built so far with this Async type. Do you think e.g. the functions we've built so far in TCP/UDP/Timers should be Async instead of AsyncTask?

Beyond that as discussed with the user feedback on discord renaming parallel is probably a good idea but that's already on the TODO list.

@algebraic-dev algebraic-dev marked this pull request as ready for review May 20, 2025 19:00
@algebraic-dev algebraic-dev changed the title feat: simple async monad feat: monadic interface for asynchronous operations in Std. May 20, 2025
@algebraic-dev algebraic-dev changed the title feat: monadic interface for asynchronous operations in Std. feat: monadic interface for asynchronous operations in Std May 20, 2025
@algebraic-dev algebraic-dev requested review from TwoFX and hargoniX May 22, 2025 19:37
Copy link
Contributor

@hargoniX hargoniX left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would also be great if there was some sort of documentation, either in the code or this PR body, that actually explains the intentions behind all of these classes etc. So far with mini-redis I've mostly been guessing based on the test files here and conversations with you.

I also wonder about the planned integration with the existing async IO things, are we going to migrate those to Async? Are they going to stay AsyncTask?

As a last technical remark: I've been wondering in mini-redis how to do the following: When I'm in a monad stack that contains Async at the bottom and I end up producing a raw Async, how can I run that as a background computation? Running async on it doesn't seem to work and neither does discard.

Style wise the file looks "quite Haskell" which is not something we usually do but I'm not generally opposed to it, maybe other people have strong opinions here.

instance [inst : MonadAsync t m] [Monad m] : MonadAsync t (ReaderT n m) where
async p := inst.async ∘ p

instance [inst : MonadAsync t m] [Monad m] : MonadAsync t (StateRefT' s n m) where
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As discussed in DMs I'm not 100% sure about this instance, if we spawn a couple of computations in a StateRefT with this they are all going to be referring to the same IO.Ref and thus race all the times no? Also why does StateRefT have an instance but not StateT and the others above that have instances?

Copy link
Member Author

@algebraic-dev algebraic-dev May 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure if we should provide StateRefT and StateT instances. These instances behave weirdly with Async, so I think we should only provide an instance for ReaderT, just to make it easier to have some kind of context with Mutexes.

@algebraic-dev algebraic-dev requested review from hargoniX and tydeu June 13, 2025 18:37
Copy link
Member

@tydeu tydeu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! I have a few minor suggestions and one major one, but I am otherwise quite happy with how this is turning out. anxiously await being able to use this API in Lake to improve the codebase's readability!

@algebraic-dev algebraic-dev requested a review from tydeu June 18, 2025 03:40
@algebraic-dev algebraic-dev added the release-ci Enable all CI checks for a PR, like is done for releases label Jun 18, 2025
Copy link
Member

@tydeu tydeu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks even better! 😄

algebraic-dev and others added 3 commits June 18, 2025 01:17
Co-authored-by: Mac Malone <[email protected]>
Co-authored-by: Mac Malone <[email protected]>
@algebraic-dev algebraic-dev requested a review from tydeu June 18, 2025 16:32
Copy link
Member

@TwoFX TwoFX left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great! Feel free to merge after adressing the style nit.

@algebraic-dev algebraic-dev enabled auto-merge June 26, 2025 00:48
@algebraic-dev algebraic-dev added this pull request to the merge queue Jun 26, 2025
Merged via the queue into leanprover:master with commit b15cfad Jun 26, 2025
19 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

changelog-library Library release-ci Enable all CI checks for a PR, like is done for releases toolchain-available A toolchain is available for this PR, at leanprover/lean4-pr-releases:pr-release-NNNN

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants