Bigstring body

aws-s3-lwt/io.ml

@@ -68,21 +68,36 @@ module Pipe = struct
       failwith (__LOC__ ^ ": Closed")
 
   let flush writer =
-    match Queue.length writer.queue = 0 && writer.closed with
+    match writer.closed with
     | true -> Lwt.return ()
     | false ->
     let waiter, wakeup = Lwt.wait () in
     Queue.add (Flush wakeup) writer.queue;
     if Queue.length writer.queue = 1 then Lwt_condition.signal writer.cond ();
     waiter
 
+  (* Only a reader reaching a marker resolves the flush that queued it, so a
+     pipe nobody reads again has to release them itself: a producer bounding
+     itself on flush waits forever otherwise. *)
+  let wake_flushers pipe =
+    let data = Queue.create () in
+    Queue.iter
+      (function
+        | Flush wakeup -> Lwt.wakeup_later wakeup ()
+        | Data _ as elem -> Queue.add elem data)
+      pipe.queue;
+    Queue.clear pipe.queue;
+    Queue.transfer data pipe.queue
+
   let close (writer : 'a writer) =
     writer.closed <- true;
+    wake_flushers writer;
     Lwt_condition.broadcast writer.cond ();
     on_close writer
 
   let close_reader (reader : 'a reader) =
     reader.closed <- true;
+    wake_flushers reader;
     Lwt_condition.broadcast reader.cond ();
     on_close reader
 

aws-s3.opam

@@ -15,6 +15,7 @@ depends: [
   "ocaml" {>= "4.08.0"}
   "dune" {>= "2.0.0"}
   "inifiles"
+  "bigstringaf" {>= "0.5.0"}
   "digestif" {>= "0.7"}
   "ptime"
   "uri"

aws-s3/authorization.ml

@@ -12,6 +12,9 @@ let log fmt = match debug with
 let hash_sha256 s =
   Digestif.SHA256.digest_string s
 
+let hash_sha256_bigstring s =
+  Digestif.SHA256.digest_bigstring s
+
 let hmac_sha256 ~key v =
   Digestif.SHA256.hmac_string ~key v
 

aws-s3/authorization.mli

@@ -1,5 +1,6 @@
 (**/**)
 val hash_sha256 : string -> Digestif.SHA256.t
+val hash_sha256_bigstring : Bigstringaf.t -> Digestif.SHA256.t
 val hmac_sha256 : key:string -> string -> Digestif.SHA256.t
 val to_hex : Digestif.SHA256.t -> string
 

aws-s3/aws.ml

@@ -96,13 +96,45 @@ module Make(Io : Types.Io) = struct
     in
     Pipe.create_reader ~f:(transfer initial_signature Digestif.SHA256.empty 0 [] None)
 
+  (* Sized to what the transport allocates per read anyway, so a bigstring body
+     reaches the socket without an object-sized string on the heap. *)
+  let body_slice_size = 65536
+
+  (* One slice ahead of the consumer, as {!chunk_writer} does: a pipe write does
+     not block, so without the flush the whole body would queue up as the
+     strings this exists to avoid. *)
+  let bigstring_reader body =
+    let rec send writer offset =
+      match Bigstringaf.length body - offset with
+      | _ when Pipe.is_closed writer -> return ()
+      | 0 -> return ()
+      | remain ->
+        let flushed = Pipe.flush writer in
+        let len = min body_slice_size remain in
+        Pipe.write writer (Bigstringaf.substring body ~off:offset ~len) >>= fun () ->
+        flushed >>= fun () ->
+        send writer (offset + len)
+    in
+    Pipe.create_reader ~f:(fun writer -> send writer 0)
+
+  let reader_of_body ~chunked = function
+    | Body.String body ->
+      let reader, writer = Pipe.create () in
+      Pipe.write writer body >>= fun () ->
+      Pipe.close writer;
+      return (Some reader)
+    | Body.Bigstring body -> return (Some (bigstring_reader body))
+    | Body.Empty -> return None
+    | Body.Chunked { pipe; chunk_size; _ } -> chunked ~pipe ~chunk_size
+
   let make_request ~(endpoint: Region.endpoint) ?connect_timeout_ms ?(expect=false) ~sink ?(body=Body.Empty) ?(credentials:Credentials.t option) ~headers ~meth ~path ~query () =
     let (date, time)  = Unix.gettimeofday () |> Time.iso8601_of_time in
 
     (* Create headers structure *)
     let content_length =
       match meth, body with
       | (`PUT | `POST), Body.String body -> Some (String.length body |> string_of_int)
+      | (`PUT | `POST), Body.Bigstring body -> Some (Bigstringaf.length body |> string_of_int)
       | (`PUT | `POST), Body.Chunked { length; chunk_size; _ } ->
         Some (get_chunked_length ~chunk_size length |> string_of_int )
       | (`PUT | `POST), Body.Empty -> Some "0"
@@ -111,6 +143,7 @@ module Make(Io : Types.Io) = struct
     let payload_sha = match body with
       | Body.Empty -> empty_sha
       | Body.String body -> Authorization.hash_sha256 body |> Authorization.to_hex
+      | Body.Bigstring body -> Authorization.hash_sha256_bigstring body |> Authorization.to_hex
       | Body.Chunked _ -> "STREAMING-AWS4-HMAC-SHA256-PAYLOAD"
     in
     let token = match credentials with
@@ -153,32 +186,16 @@ module Make(Io : Types.Io) = struct
               ~headers ~query:query ~scope ~signing_key ~payload_sha
           in
           let auth = (Authorization.make_auth_header ~credentials ~scope ~signed_headers ~signature) in
-          let body = match body with
-            | Body.String body ->
-              let reader, writer = Pipe.create () in
-              Pipe.write writer body >>= fun () ->
-              Pipe.close writer;
-              return (Some reader)
-            | Body.Empty -> return None
-            | Body.Chunked { pipe; chunk_size; _ } ->
-              let pipe =
+          let body =
+            reader_of_body body ~chunked:(fun ~pipe ~chunk_size ->
                 (* Get errors if the chunk_writer fails *)
-                chunk_writer ~signing_key ~scope
-                  ~initial_signature:signature ~date ~time ~chunk_size pipe
-              in
-              return (Some pipe)
+                return (Some (chunk_writer ~signing_key ~scope
+                                ~initial_signature:signature ~date ~time ~chunk_size pipe)))
           in
           Some auth, body
         | None ->
-          let body = match body with
-            | Body.String body ->
-              let reader, writer = Pipe.create () in
-              Pipe.write writer body >>= fun () ->
-              Pipe.close writer;
-              return (Some reader)
-            | Body.Empty -> return None
-            | Body.Chunked { pipe; _} ->
-              return (Some pipe)
+          let body =
+            reader_of_body body ~chunked:(fun ~pipe ~chunk_size:_ -> return (Some pipe))
           in
           None, body
       in

aws-s3/body.ml

@@ -1,11 +1,67 @@
 open StdLabels
 
+(* Bytes gathered off the OCaml heap. It doubles because a sink is never told
+   the content length it is about to receive. *)
+module Accumulator = struct
+  type t = { mutable buffer: Bigstringaf.t; mutable filled: int }
+
+  let initial_size = 65536
+
+  let create () = { buffer = Bigstringaf.empty; filled = 0 }
+
+  let reserve t needed =
+    match Bigstringaf.length t.buffer >= needed with
+    | true -> ()
+    | false ->
+      let rec double size = match size >= needed with
+        | true -> size
+        | false -> double (size * 2)
+      in
+      let size = max (Bigstringaf.length t.buffer) initial_size in
+      let grown = Bigstringaf.create (double size) in
+      Bigstringaf.blit t.buffer ~src_off:0 grown ~dst_off:0 ~len:t.filled;
+      t.buffer <- grown
+
+  let add_string t data =
+    let len = String.length data in
+    reserve t (t.filled + len);
+    Bigstringaf.blit_from_string data ~src_off:0 t.buffer ~dst_off:t.filled ~len;
+    t.filled <- t.filled + len
+
+  (* A [sub] would share the slack doubling left behind, and the caller holds
+     the result for as long as it holds the body. *)
+  let contents t =
+    match t.filled = Bigstringaf.length t.buffer with
+    | true -> t.buffer
+    | false -> Bigstringaf.copy t.buffer ~off:0 ~len:t.filled
+end
+
+let%test "accumulator answers exactly what was added, in order" =
+  let acc = Accumulator.create () in
+  let expected = Buffer.create 0 in
+  (* Pieces of varying length carrying their own index, so a lost, doubled or
+     misplaced one shows up rather than being covered by its neighbour. *)
+  for i = 1 to 5000 do
+    let piece = Printf.sprintf "%d:%s|" i (String.make (i mod 97) 'x') in
+    Buffer.add_string expected piece;
+    Accumulator.add_string acc piece
+  done;
+  let result = Accumulator.contents acc in
+  let expected = Buffer.contents expected in
+  String.length expected > Accumulator.initial_size
+  && Bigstringaf.length result = String.length expected
+  && Bigstringaf.to_string result = expected
+
+let%test "an accumulator nothing was added to is empty" =
+  Bigstringaf.length (Accumulator.contents (Accumulator.create ())) = 0
+
 module Make(Io : Types.Io) = struct
   open Io
   open Deferred
 
   type t =
     | String of string
+    | Bigstring of Bigstringaf.t
     | Empty
     | Chunked of { pipe: string Pipe.reader; length: int; chunk_size: int }
 
@@ -27,6 +83,19 @@ module Make(Io : Types.Io) = struct
     in
     loop []
 
+  (* Each fragment is dropped as soon as it is gathered, where {!to_string}
+     holds every one of them until the end. *)
+  let to_bigstring body =
+    let acc = Accumulator.create () in
+    let rec loop () =
+      Pipe.read body >>= function
+      | Some data ->
+        Accumulator.add_string acc data;
+        loop ()
+      | None -> return (Accumulator.contents acc)
+    in
+    loop ()
+
   let read_string ?start ~length reader =
     let rec loop acc data remain =
       match data, remain with

aws-s3/body.mli

@@ -2,13 +2,19 @@ module Make(Io : Types.Io) : sig
   open Io
   type t =
     | String of string
+    | Bigstring of Bigstringaf.t
     | Empty
     | Chunked of { pipe : string Pipe.reader; length : int; chunk_size : int; }
   (**/**)
   val to_string :
     string Pipe.reader ->
     string Deferred.t
 
+  (** {!to_string} without putting the body on the OCaml heap. *)
+  val to_bigstring :
+    string Pipe.reader ->
+    Bigstringaf.t Deferred.t
+
   val read_string :
     ?start:string ->
     length:int ->

aws-s3/dune

@@ -2,7 +2,7 @@
  (name        aws_s3)
  (public_name aws-s3)
  (synopsis "Amazon S3 access library")
- (libraries ptime inifiles digestif.c
+ (libraries ptime inifiles bigstringaf digestif.c
             base64 uri yojson
             ppx_protocol_conv_json
             ppx_protocol_conv_xmlm str)

aws-s3/http.ml

@@ -164,5 +164,7 @@ module Make(Io : Types.Io) = struct
     Pipe.close writer;
     Pipe.close_reader reader;
     Pipe.close sink;
+    (* A body the request never got to send still has a producer behind it. *)
+    (match body with Some body -> Pipe.close_reader body | None -> ());
     return result
 end