Merge pull request #21 from inhabitedtype/allocation

allocation: be better about allocating a new buffer for writes

Author: seliopou

lib/faraday.ml

@@ -206,9 +206,7 @@ let flush_buffer t =
   if len > 0 then begin
     let off = t.scheduled_pos in
     schedule_iovec t ~off ~len (`Bigstring t.buffer);
-    t.buffer <- Bigarray.(Array1.create char c_layout (Array1.dim t.buffer));
-    t.write_pos <- 0;
-    t.scheduled_pos <- 0
+    t.scheduled_pos <- t.write_pos
   end
 
 let flush t f =
@@ -267,13 +265,13 @@ let schedule_bigstring =
   fun t ?off ?len a -> schedule_gen t ~length ~to_buffer ?off ?len a
 
 let ensure_space t len =
-  if free_bytes_in_buffer t < len then (
+  if free_bytes_in_buffer t < len then begin
     flush_buffer t;
-    `Not_enough_space
-  )
-  else (
-    `Go_ahead
-  )
+    t.buffer <-
+      Bigarray.(Array1.create char c_layout (max (Array1.dim t.buffer) len));
+    t.write_pos <- 0;
+    t.scheduled_pos <- 0
+  end
 
 let write_gen t ~length ~blit ?(off=0) ?len a =
   writable t;
@@ -282,14 +280,9 @@ let write_gen t ~length ~blit ?(off=0) ?len a =
     | None     -> length a - off
     | Some len -> len
   in
-  match ensure_space t len with
-  | `Not_enough_space ->
-    let buffer = Bigarray.(Array1.create char c_layout len) in
-    blit a off buffer 0 len;
-    schedule_iovec t ~len (`Bigstring buffer)
-  | `Go_ahead ->
-    blit a off t.buffer t.write_pos len;
-    t.write_pos <- t.write_pos + len
+  ensure_space t len;
+  blit a off t.buffer t.write_pos len;
+  t.write_pos <- t.write_pos + len
 
 let write_string =
   let length   = String.length in
@@ -421,8 +414,10 @@ let rec shift_buffers t written =
       Buffers.enqueue_front (IOVec.shift iovec written) t.scheduled
   with Not_found ->
     assert (written = 0);
-    t.scheduled_pos <- 0;
-    t.write_pos <- 0
+    if t.scheduled_pos = t.write_pos then begin
+      t.scheduled_pos <- 0;
+      t.write_pos <- 0
+    end
 
 let rec shift_flushes t =
   try

lib_test/test_faraday.ml

@@ -58,6 +58,14 @@ let write_tiny_buf =
       check ~buf_size:1 ~iovecs:1 ~msg:"string"    [`Write_string    "test"] "test";
       check ~buf_size:1 ~iovecs:1 ~msg:"bytes"     [`Write_bytes     "test"] "test";
       check ~buf_size:1 ~iovecs:1 ~msg:"bigstring" [`Write_bigstring "test"] "test"
+  end 
+  ; "multiple writes with tiny buffer", `Quick,  begin fun () ->
+      check ~buf_size:1 ~iovecs:2 ~msg:"string" [`Write_string "test1"; `Write_string "test2"] "test1test2"
+  end
+  ; "too many writes with tiny buffer", `Quick,  begin fun () ->
+      check ~buf_size:1 ~iovecs:5 ~msg:"string"
+        [`Write_string "te"; `Write_string "st"; `Write_string "te"; `Write_string "st"; `Write_string "te" ] 
+        "testtestte"
   end ]
 
 let schedule =
@@ -101,6 +109,6 @@ let () =
   Alcotest.run "test suite"
     [ "empty output"              , empty
     ; "single write"              , write
-    ; "single write (tiny buffer)", write
+    ; "writes (tiny buffer)"      , write_tiny_buf
     ; "single schedule"           , schedule
     ; "interleaved calls"         , interleaved ]