Ensure Buf_write still flushes if an exception is raised

lib_eio/buf_write.ml

@@ -489,7 +489,17 @@ let with_flow ?(initial_size=0x1000) flow fn =
   Switch.run @@ fun sw ->
   let t = create ~sw initial_size in
   Fiber.fork ~sw (fun () -> Flow.copy (as_flow t) flow);
-  Fun.protect ~finally:(fun () -> close t) (fun () -> fn t)
+  match fn t with
+  | x ->
+    close t;
+    x
+  | exception ex ->
+    close t;
+    (* Raising the exception will cancel the writer thread, so do a flush first.
+       We don't want to flush if cancelled, but in that case the switch will
+       end the writer thread itself (and [flush] will raise). *)
+    flush t;
+    raise ex
 
 let rec serialize t writev =
   match await_batch t with

tests/buf_write.md

@@ -239,9 +239,11 @@ Cancelled while waiting for the underlying flow to perform the write:
 # Eio_mock.Backend.run @@ fun () ->
   let flow = Eio_mock.Flow.make "flow" in
   Eio_mock.Flow.on_copy_bytes flow [`Run Fiber.await_cancel];
-  Write.with_flow flow @@ fun t ->
   Fiber.both
-    (fun () -> Write.string t "Hello"; traceln "Did write")
+    (fun () ->
+       Write.with_flow flow @@ fun t ->
+       Write.string t "Hello"; traceln "Did write"
+    )
     (fun () -> Fiber.yield (); failwith "Simulated error");;
 +Did write
 Exception: Failure "Simulated error".
@@ -321,10 +323,13 @@ We still flush the output on error:
 
 ```ocaml
 # Eio_mock.Backend.run @@ fun () ->
+  Eio_mock.Flow.on_copy_bytes flow [`Return 1; `Yield_then (`Return 1)];
   Write.with_flow flow @@ fun t ->
   Write.string t "foo";
   failwith "Simulated error";;
-+flow: wrote "foo"
++flow: wrote "f"
++flow: wrote "o"
++flow: wrote "o"
 Exception: Failure "Simulated error".
 ```