src/lwt/moonpool_lwt.ml

@@ -175,13 +175,6 @@ 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 *)
@@ -193,8 +186,13 @@ let await_lwt (fut : _ Lwt.t) =
       register_trigger_on_lwt_termination fut tr;
       M.Trigger.await_exn tr;
       await_lwt_terminated fut
-  ) else
+  ) else (
-    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 lwt_of_fut (fut : 'a M.Fut.t) : 'a Lwt.t =
   if not (Main_state.on_lwt_thread ()) then
@@ -297,16 +295,9 @@ let spawn_lwt f : _ Lwt.t =
   lwt_fut
 
 let spawn_lwt_ignore f = ignore (spawn_lwt f : unit Lwt.t)
+let run_in_lwt_and_await (f : unit -> 'a) : 'a = await_lwt @@ spawn_lwt f
 let on_lwt_thread = Main_state.on_lwt_thread
 
-let run_in_lwt_and_await (f : unit -> 'a) : 'a =
-  let st = Main_state.get_st () in
-  if Scheduler_state.on_lwt_thread_ st then
-    (* run immediately *)
-    f ()
-  else
-    await_lwt_from_another_thread @@ spawn_lwt f
-
 let lwt_main (f : _ -> 'a) : 'a =
   let st = setup () in
   (* make sure to cleanup *)