fix: more side effects inside fiber

pereb4ik · Feb 12, 2022 · 893a15c · 893a15c
1 parent 762acb8
commit 893a15c
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 27 deletions.
diff --git a/jsonrpc-fiber/src/jsonrpc_fiber.ml b/jsonrpc-fiber/src/jsonrpc_fiber.ml
@@ -228,15 +228,16 @@ struct
           (Dyn.to_string (Dyn.list Exn_with_backtrace.to_dyn errors));
         loop ()
     in
-    t.running <- true;
-    let* () =
-      Fiber.fork_and_join_unit
-        (fun () ->
-          let* () = loop () in
-          Fiber.Pool.stop later)
-        (fun () -> Fiber.Pool.run later)
-    in
-    close t
+    Fiber.of_thunk (fun () ->
+        t.running <- true;
+        let* () =
+          Fiber.fork_and_join_unit
+            (fun () ->
+              let* () = loop () in
+              Fiber.Pool.stop later)
+            (fun () -> Fiber.Pool.run later)
+        in
+        close t)
 
   let on_notification_fail ctx =
     let state = Context.state ctx in
@@ -288,14 +289,15 @@ struct
     | Error `Stopped -> raise (Stopped req)
 
   let request t (req : Message.request) =
-    check_running t;
-    let* () =
-      let req = { req with Message.id = Some req.id } in
-      Chan.send t.chan [ Message req ]
-    in
-    let ivar = Fiber.Ivar.create () in
-    register_request_ivar t req.id ivar;
-    read_request_ivar req ivar
+    Fiber.of_thunk (fun () ->
+        check_running t;
+        let* () =
+          let req = { req with Message.id = Some req.id } in
+          Chan.send t.chan [ Message req ]
+        in
+        let ivar = Fiber.Ivar.create () in
+        register_request_ivar t req.id ivar;
+        read_request_ivar req ivar)
 
   module Batch = struct
     type response =

diff --git a/jsonrpc-fiber/test/jsonrpc_fiber_tests.ml b/jsonrpc-fiber/test/jsonrpc_fiber_tests.ml
@@ -171,18 +171,19 @@ let%expect_test "concurrent requests" =
   let waiter_in, waitee_out = pipe () in
   let waitee = waitee (waitee_in, waitee_out) in
   let waiter = waiter (waiter_in, waiter_out) in
-  let initial_request () =
-    let request =
-      Jsonrpc.Message.create ~id:(`String "initial") ~method_:"init" ()
+  let run () =
+    let initial_request () =
+      let request =
+        Jsonrpc.Message.create ~id:(`String "initial") ~method_:"init" ()
+      in
+      print_endline "initial: waitee requests from waiter";
+      let+ resp = Jrpc.request waitee request in
+      print_endline "initial request response:";
+      print (Response resp)
     in
-    print_endline "initial: waitee requests from waiter";
-    let+ resp = Jrpc.request waitee request in
-    print_endline "initial request response:";
-    print (Response resp)
+    Fiber.all_concurrently_unit
+      [ Jrpc.run waitee; initial_request (); Jrpc.run waiter ]
   in
-  let all = [ Jrpc.run waiter; Jrpc.run waitee ] in
-  let all = initial_request () :: all in
-  let run () = Fiber.parallel_iter all ~f:Fun.id in
   Fiber_test.test Dyn.opaque run;
   [%expect
     {|