Merge pull request #539 from talex5/fix-flush

Buf_write: fix flush returning too early

Author: talex5

lib_eio/buf_write.ml

@@ -469,27 +469,21 @@ let rec await_batch t =
       );
     await_batch t
 
-let read_into t buf =
-  let iovecs = await_batch t in
-  let n, _iovecs = Cstruct.fillv ~src:iovecs ~dst:buf in
-  shift t n;
-  n
-
-let read_source_buffer t fn =
-  let iovecs = await_batch t in
-  shift t (fn iovecs)
-
-let as_flow t =
-  object
-    inherit Flow.source
-    method! read_methods = [Flow.Read_source_buffer (read_source_buffer t)]
-    method read_into = read_into t
-  end
+(* We have to do our own copy, because we can't [shift] until the write is complete. *)
+let copy t flow =
+  let rec aux () =
+    let iovecs = await_batch t in
+    Flow.write flow iovecs;             (* todo: add a Flow.single_write and use that. *)
+    shift t (Cstruct.lenv iovecs);
+    aux ()
+  in
+  try aux ()
+  with End_of_file -> ()
 
 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);
+  Fiber.fork ~sw (fun () -> copy t flow);
   match fn t with
   | x ->
     close t;

tests/buf_write.md

@@ -182,7 +182,10 @@ Eio_mock.Flow.on_copy_bytes flow [
 - : unit = ()
 ```
 
-Multiple flushes:
+Multiple flushes.
+Note: ideally the flushes here would complete as soon as enough data has been flushed,
+but currently Eio.Flow.sink doesn't allow short writes and so Buf_write has to wait for
+the whole batch to be flushed.
 
 ```ocaml
 # Eio_mock.Backend.run @@ fun () ->
@@ -201,15 +204,43 @@ Multiple flushes:
   traceln "Done";;
 +flow: wrote (rsb) ["a"]
 +flow: wrote (rsb) ["b"; "c"]
-+1st flush
 +flow: wrote (rsb) ["d"; "e"]
-+2nd flush
 +flow: wrote (rsb) ["f"]
++1st flush
++2nd flush
 +3rd flush
 +Done
 - : unit = ()
 ```
 
+Check flush waits for the write to succeed:
+
+```ocaml
+let slow_writer = object
+  inherit Eio.Flow.sink
+  method copy src =
+    let buf = Cstruct.create 10 in
+    try
+      while true do
+        let len = Eio.Flow.single_read src buf in
+        Fiber.yield ();
+        traceln "Write %S" (Cstruct.to_string buf ~len)
+      done
+    with End_of_file -> ()
+end
+```
+
+```ocaml
+# Eio_mock.Backend.run @@ fun () ->
+  Write.with_flow slow_writer @@ fun t ->
+  Write.string t "test";
+  Write.flush t;
+  traceln "Flush complete"
++Write "test"
++Flush complete
+- : unit = ()
+```
+
 ## Scheduling
 
 ```ocaml