Merge pull request #49 from inhabitedtype/simplify-api

Zero-copy read API

Author: seliopou

async/httpaf_async.ml

@@ -1,6 +1,56 @@
 open Core
 open Async
 
+(** XXX(seliopou): Replace Angstrom.Buffered with a module like this, while
+    also supporting growing the buffer. Clients can use this to buffer and the
+    use the unbuffered interface for actually running the parser. *)
+module Buffer : sig
+  type t
+
+  val create   : int -> t
+
+  val get : t -> f:(Bigstring.t -> off:int -> len:int -> int) -> int
+  val put : t -> f:(Bigstring.t -> off:int -> len:int -> int) -> int
+end= struct
+  type t =
+    { buffer      : Bigstring.t
+    ; mutable off : int
+    ; mutable len : int }
+
+  let create size =
+    let buffer = Bigstring.create size in
+    { buffer; off = 0; len = 0 }
+  ;;
+
+  let compress t =
+    if t.len = 0
+    then begin
+      t.off <- 0;
+      t.len <- 0;
+    end else if t.off > 0
+    then begin
+      Bigstring.blit ~src:t.buffer ~src_pos:t.off ~dst:t.buffer ~dst_pos:0 ~len:t.len;
+      t.off <- 0;
+    end
+  ;;
+
+  let get t ~f =
+    let n = f t.buffer ~off:t.off ~len:t.len in
+    t.off <- t.off + n;
+    t.len <- t.len - n;
+    if t.len = 0
+    then t.off <- 0;
+    n
+  ;;
+
+  let put t ~f =
+    compress t;
+    let n = f t.buffer ~off:(t.off + t.len) ~len:(Bigstring.length t.buffer - t.len) in
+    t.len <- t.len + n;
+    n
+  ;;
+end
+
 let read fd buffer =
   let badfd fd = failwithf "read got back fd: %s" (Fd.to_string fd) () in
   let rec finish fd buffer result =
@@ -25,16 +75,18 @@ let read fd buffer =
       finish fd buffer
         (Fd.syscall fd ~nonblocking:true
           (fun file_descr ->
-            Unix.Syscall_result.Int.ok_or_unix_error_exn ~syscall_name:"read"
-              (Bigstring.read_assume_fd_is_nonblocking file_descr buffer)))
+            Buffer.put buffer ~f:(fun bigstring ~off ~len ->
+              Unix.Syscall_result.Int.ok_or_unix_error_exn ~syscall_name:"read"
+                (Bigstring.read_assume_fd_is_nonblocking file_descr bigstring ~pos:off ~len))))
     else
       Fd.syscall_in_thread fd ~name:"read"
-        (fun file_descr -> Bigstring.read file_descr buffer)
+        (fun file_descr ->
+          Buffer.put buffer ~f:(fun bigstring ~off ~len ->
+            Bigstring.read file_descr bigstring ~pos:off ~len))
       >>= fun result -> finish fd buffer result
   in
   go fd buffer
 
-
 open Httpaf
 
 module Server = struct
@@ -46,13 +98,23 @@ module Server = struct
       let error_handler   = error_handler client_addr in
       let conn = Server_connection.create ?config ~error_handler request_handler in
       let read_complete = Ivar.create () in
+      (* XXX(seliopou): Make this configurable *)
+      let buffer = Buffer.create 0x1000 in
       let rec reader_thread () =
         match Server_connection.next_read_operation conn with
-        | `Read buffer ->
+        | `Read ->
           (* Log.Global.printf "read(%d)%!" (Fd.to_int_exn fd); *)
-          read fd buffer >>> fun result ->
-            Server_connection.report_read_result conn result;
-            reader_thread ()
+          read fd buffer
+          >>> begin function
+            | `Eof  ->
+              Server_connection.shutdown_reader conn;
+              reader_thread ()
+            | `Ok _ ->
+              Buffer.get buffer ~f:(fun bigstring ~off ~len ->
+                Server_connection.read conn bigstring ~off ~len)
+              |> ignore;
+              reader_thread ()
+          end
         | `Yield  ->
           (* Log.Global.printf "read_yield(%d)%!" (Fd.to_int_exn fd); *)
           Server_connection.yield_reader conn reader_thread
@@ -100,13 +162,22 @@ module Client = struct
     let request_body, conn   =
       Client_connection.request request ~error_handler ~response_handler in
     let read_complete = Ivar.create () in
+    let buffer = Buffer.create 0x1000 in
     let rec reader_thread () =
       match Client_connection.next_read_operation conn with
-      | `Read buffer ->
+      | `Read ->
         (* Log.Global.printf "read(%d)%!" (Fd.to_int_exn fd); *)
-        read fd buffer >>> fun result ->
-          Client_connection.report_read_result conn result;
-          reader_thread ()
+        read fd buffer
+          >>> begin function
+            | `Eof  ->
+              Client_connection.shutdown_reader conn;
+              reader_thread ()
+            | `Ok _ ->
+              Buffer.get buffer ~f:(fun bigstring ~off ~len ->
+                Client_connection.read conn bigstring ~off ~len)
+              |> ignore;
+              reader_thread ()
+          end
       | `Close ->
         (* Log.Global.printf "read_close(%d)%!" (Fd.to_int_exn fd); *)
         Ivar.fill read_complete ();

benchmarks/wrk_async_benchmark.ml

@@ -17,13 +17,13 @@ let error_handler _ ?request error start_response =
   | #Status.standard as error ->
     Body.write_string response_body (Status.default_reason_phrase error)
   end;
-  Body.close response_body
+  Body.close_writer response_body
 ;;
 
 let request_handler _ reqd =
   let { Request.target } = Reqd.request reqd in
   let request_body       = Reqd.request_body reqd in
-  Body.close request_body;
+  Body.close_reader request_body;
   match target with
   | "/" -> Reqd.respond_with_bigstring reqd (Response.create ~headers `OK) text;
   | _   -> Reqd.respond_with_string    reqd (Response.create `Not_found) "Route not found"

lib/client_connection.ml

@@ -91,11 +91,8 @@ module Oneshot = struct
       Body.transfer_to_writer_with_encoding t.request_body ~encoding t.writer
   ;;
 
-  let shutdown t =
-    flush_request_body t;
+  let shutdown_reader t =
     Reader.close t.reader;
-    Writer.close t.writer;
-    Body.close_writer t.request_body;
     begin match !(t.state) with
     | Awaiting_response | Closed -> ()
     | Received_response(_, response_body) ->
@@ -104,6 +101,17 @@ module Oneshot = struct
     end;
   ;;
 
+  let shutdown_writer t =
+    flush_request_body t;
+    Writer.close t.writer;
+    Body.close_writer t.request_body;
+  ;;
+
+  let shutdown t =
+    shutdown_reader t;
+    shutdown_writer t;
+  ;;
+
   let set_error_and_handle t error =
     shutdown t;
     t.state := Closed;
@@ -144,12 +152,13 @@ module Oneshot = struct
     | `Error (`Invalid_response_body_length _ as error) ->
       set_error_and_handle t error;
       `Close
-    | (`Read _ | `Close) as operation -> operation
+    | (`Read | `Close) as operation -> operation
   ;;
 
-  let report_read_result t result =
-    Reader.report_result t.reader result;
+  let read t bs ~off ~len =
+    let consumed = Reader.read t.reader bs ~off ~len in
     flush_response_body t;
+    consumed
   ;;
 
   let next_write_operation t =

lib/httpaf.mli

@@ -695,8 +695,7 @@ end
 module Server_connection : sig
   module Config : sig
     type t =
-      { read_buffer_size          : int (** Default is [4096] *)
-      ; response_buffer_size      : int (** Default is [1024] *)
+      { response_buffer_size      : int (** Default is [1024] *)
       ; response_body_buffer_size : int (** Default is [4096] *)
       }
 
@@ -723,27 +722,29 @@ module Server_connection : sig
   (** [create ?config ?error_handler ~request_handler] creates a connection
       handler that will service individual requests with [request_handler]. *)
 
-  val next_read_operation : _ t -> [ `Read of Bigstring.t | `Yield | `Close ]
+  val next_read_operation : _ t -> [ `Read | `Yield | `Close ]
   (** [next_read_operation t] returns a value describing the next operation
       that the caller should conduct on behalf of the connection. *)
 
-  val report_read_result : _ t -> [`Ok of int | `Eof] -> unit
-  (** [report_read_result t result] reports the result of the latest read
-      attempt to the connection. {report_read_result} should be called after a
-      call to {next_read_operation} that returns a [`Read buffer] value.
-
-        {ul
-        {- [`Ok n] indicates that the caller successfully received [n] bytes of
-        input and wrote them into the the read buffer that the caller was
-        provided by {next_read_operation}. }
-        {- [`Eof] indicates that the input source will no longer provide any
-        bytes to the read processor. }} *)
+  val read : _ t -> Bigstring.t -> off:int -> len:int -> int
+  (** [read t bigstring ~off ~len] reads bytes of input from the provided range
+      of [bigstring] and returns the number of bytes consumed by the
+      connection.  {!read} should be called after {!next_read_operation}
+      returns a [`Read] value and additional input is available for the
+      connection to consume. *)
 
   val yield_reader : _ t -> (unit -> unit) -> unit
   (** [yield_reader t continue] registers with the connection to call
       [continue] when reading should resume. {!yield_reader} should be called
       after {next_read_operation} returns a [`Yield] value. *)
 
+  val shutdown_reader : _ t -> unit
+  (** [shutdown_reader t] shutds own the read processor for the connection. All
+      subsequent calls to {!next_read_operations} will return [`Close].
+      {!shutdown_reader} should be called after {!next_read_operation} returns
+      a [`Read] value and there is no further input available for the
+      connection to consume. *)
+
   val next_write_operation : _ t -> [
     | `Write of Bigstring.t IOVec.t list
     | `Yield
@@ -808,21 +809,23 @@ module Client_connection : sig
     -> response_handler:response_handler
     -> [`write] Body.t * t
 
-  val next_read_operation : t -> [ `Read of Bigstring.t | `Close ]
+  val next_read_operation : t -> [ `Read | `Close ]
   (** [next_read_operation t] returns a value describing the next operation
       that the caller should conduct on behalf of the connection. *)
 
-  val report_read_result : t -> [`Ok of int | `Eof] -> unit
-  (** [report_read_result t result] reports the result of the latest read
-      attempt to the connection. {report_read_result} should be called after a
-      call to {next_read_operation} that returns a [`Read buffer] value.
-
-        {ul
-        {- [`Ok n] indicates that the caller successfully received [n] bytes of
-        input and wrote them into the the read buffer that the caller was
-        provided by {next_read_operation}. }
-        {- [`Eof] indicates that the input source will no longer provide any
-        bytes to the read processor. }} *)
+  val read : t -> Bigstring.t -> off:int -> len:int -> int
+  (** [read t bigstring ~off ~len] reads bytes of input from the provided range
+      of [bigstring] and returns the number of bytes consumed by the
+      connection.  {!read} should be called after {!next_read_operation}
+      returns a [`Read] value and additional input is available for the
+      connection to consume. *)
+
+  val shutdown_reader : t -> unit
+  (** [shutdown_reader t] shutds own the read processor for the connection. All
+      subsequent calls to {!next_read_operations} will return [`Close].
+      {!shutdown_reader} should be called after {!next_read_operation} returns
+      a [`Read] value and there is no further input available for the
+      connection to consume. *)
 
   val next_write_operation : t -> [
     | `Write of Bigstring.t IOVec.t list