From 2dbbad4ef22dc8a062ea8b2b88108a16441306f8 Mon Sep 17 00:00:00 2001 From: Simon Cruanes Date: Wed, 24 Sep 2025 15:19:21 -0400 Subject: [PATCH] refactor moonpool_lwt --- src/lwt/moonpool_lwt.ml | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/lwt/moonpool_lwt.ml b/src/lwt/moonpool_lwt.ml index 200ab72e..89040d2e 100644 --- a/src/lwt/moonpool_lwt.ml +++ b/src/lwt/moonpool_lwt.ml @@ -175,6 +175,13 @@ module Main_state = struct Scheduler_state.add_action_from_another_thread_ (get_st ()) f end +let await_lwt_from_another_thread fut = + let tr = M.Trigger.create () in + Main_state.add_action_from_another_thread (fun () -> + register_trigger_on_lwt_termination fut tr); + M.Trigger.await_exn tr; + await_lwt_terminated fut + let await_lwt (fut : _ Lwt.t) = if Scheduler_state.on_lwt_thread_ (Main_state.get_st ()) then ( (* can directly access the future *) @@ -186,13 +193,8 @@ let await_lwt (fut : _ Lwt.t) = register_trigger_on_lwt_termination fut tr; M.Trigger.await_exn tr; await_lwt_terminated fut - ) else ( - let tr = M.Trigger.create () in - Main_state.add_action_from_another_thread (fun () -> - register_trigger_on_lwt_termination fut tr); - M.Trigger.await_exn tr; - await_lwt_terminated fut - ) + ) else + await_lwt_from_another_thread fut let lwt_of_fut (fut : 'a M.Fut.t) : 'a Lwt.t = if not (Main_state.on_lwt_thread ()) then @@ -303,7 +305,7 @@ let run_in_lwt_and_await (f : unit -> 'a) : 'a = (* run immediately *) f () else - await_lwt @@ spawn_lwt f + await_lwt_from_another_thread @@ spawn_lwt f let lwt_main (f : _ -> 'a) : 'a = let st = setup () in