cohttp-eio: improve server API

CHANGES.md

@@ -1,4 +1,5 @@
 ## Unreleased
+- cohttp-eio: use Eio.Buf_write and improve server API (talex5 #887)
 - cohttp-eio: update to Eio 0.3 (talex5 #886)
 - cohttp-eio: convert to Eio.Buf_read (talex5 #882)
 - cohttp lwt client: Connection cache and explicit pipelining (madroach #853)

cohttp-eio.opam

@@ -21,12 +21,13 @@ bug-reports: "https://github.com/mirage/ocaml-cohttp/issues"
 depends: [
   "dune" {>= "2.9"}
   "base-domains"
-  "eio" {>= "0.3"}
+  "eio" {>= "0.4"}
   "eio_main" {with-test}
   "uri" {with-test}
   "cstruct"
   "bigstringaf"
   "fmt"
+  "mdx" {with-test}
   "eio_main" {with-test}
   "http" {= version}
   "odoc" {with-doc}

cohttp-eio/examples/server1.ml

@@ -30,7 +30,7 @@ let text =
 
 open Cohttp_eio
 
-let app (req, _reader) =
+let app (req, _reader, _client_addr) =
   match Http.Request.resource req with
   | "/" -> Server.text_response text
   | "/html" -> Server.html_response text
@@ -42,5 +42,4 @@ let () =
     [ ("-p", Arg.Set_int port, " Listening port number(8080 by default)") ]
     ignore "An HTTP/1.1 server";
 
-  Eio_main.run @@ fun env ->
-  Eio.Switch.run @@ fun sw -> Server.run ~port:!port env sw app
+  Eio_main.run @@ fun env -> Server.run ~port:!port env app

cohttp-eio/src/body.ml

@@ -1,7 +1,9 @@
+module Write = Eio.Buf_write
+
 type t =
   | Fixed of string
   | Chunked of chunk_writer
-  | Custom of (Eio.Flow.sink -> unit)
+  | Custom of (Write.t -> unit)
   | Empty
 
 and chunk_writer = {
@@ -228,10 +230,10 @@ let read_chunked reader headers f =
 let write_headers t headers =
   Http.Header.iter
     (fun k v ->
-      Writer.write_string t k;
-      Writer.write_string t ": ";
-      Writer.write_string t v;
-      Writer.write_string t "\r\n")
+      Write.string t k;
+      Write.string t ": ";
+      Write.string t v;
+      Write.string t "\r\n")
     headers
 
 (* https://datatracker.ietf.org/doc/html/rfc7230#section-4.1 *)
@@ -242,21 +244,21 @@ let write_chunked t chunk_writer =
         let v =
           match value with None -> "" | Some v -> Printf.sprintf "=%s" v
         in
-        Writer.write_string t (Printf.sprintf ";%s%s" name v))
+        Write.string t (Printf.sprintf ";%s%s" name v))
       exts
   in
   let write_body = function
     | Chunk { size; data; extensions = exts } ->
-        Writer.write_string t (Printf.sprintf "%X" size);
+        Write.string t (Printf.sprintf "%X" size);
         write_extensions exts;
-        Writer.write_string t "\r\n";
-        Writer.write_string t data;
-        Writer.write_string t "\r\n"
+        Write.string t "\r\n";
+        Write.string t data;
+        Write.string t "\r\n"
     | Last_chunk exts ->
-        Writer.write_string t "0";
+        Write.string t "0";
         write_extensions exts;
-        Writer.write_string t "\r\n"
+        Write.string t "\r\n"
   in
   chunk_writer.body_writer write_body;
   chunk_writer.trailer_writer (write_headers t);
-  Writer.write_string t "\r\n"
+  Write.string t "\r\n"

cohttp-eio/src/cohttp_eio.mli

@@ -2,7 +2,7 @@ module Body : sig
   type t =
     | Fixed of string
     | Chunked of chunk_writer
-    | Custom of (Eio.Flow.sink -> unit)
+    | Custom of (Eio.Buf_write.t -> unit)
     | Empty
 
   and chunk_writer = {
@@ -28,25 +28,27 @@ end
 
 (** [Server] is a HTTP 1.1 server. *)
 module Server : sig
-  type middleware = handler -> handler
-  and handler = request -> response
-  and request = Http.Request.t * Eio.Buf_read.t
-  and response = Http.Response.t * Body.t
+  type request = Http.Request.t * Eio.Buf_read.t * Eio.Net.Sockaddr.stream
+  (** The request headers, a reader for the socket, and the address of the
+      client. To read the request body, use {!read_fixed} or {!read_chunked}. *)
+
+  type response = Http.Response.t * Body.t
+  type handler = request -> response
 
   (** {1 Request} *)
 
-  val read_fixed : request -> string
-  (** [read_fixed (request,reader)] is bytes of length [n] if "Content-Length"
-      header is a valid integer value [n] in [request]. [reader] is updated to
-      reflect that [n] bytes was read.
+  val read_fixed : Http.Request.t -> Eio.Buf_read.t -> string
+  (** [read_fixed request body] reads a string of length [n] if "Content-Length"
+      header is a valid integer value [n] in [request].
 
       @raise Invalid_argument
         if ["Content-Length"] header is missing or is an invalid value in
         [headers] OR if the request http method is not one of [POST], [PUT] or
         [PATCH]. *)
 
-  val read_chunked : request -> (Body.chunk -> unit) -> Http.Header.t
-  (** [read_chunked request chunk_handler] is [updated_headers] if
+  val read_chunked :
+    Http.Request.t -> Eio.Buf_read.t -> (Body.chunk -> unit) -> Http.Header.t
+  (** [read_chunked request body chunk_handler] is [updated_headers] if
       "Transfer-Encoding" header value is "chunked" in [headers] and all chunks
       in [reader] are read successfully. [updated_headers] is the updated
       headers as specified by the chunked encoding algorithm in
@@ -81,10 +83,14 @@ module Server : sig
     ?socket_backlog:int ->
     ?domains:int ->
     port:int ->
-    Eio.Stdenv.t ->
-    Eio.Switch.t ->
+    < domain_mgr : Eio.Domain_manager.t ; net : Eio.Net.t ; .. > ->
     handler ->
-    unit
+    'a
+
+  val connection_handler :
+    handler -> #Eio.Net.stream_socket -> Eio.Net.Sockaddr.stream -> unit
+  (** [connection_handler request_handler] is a connection handler, suitable for
+      passing to {!Eio.Net.accept_fork}. *)
 
   (** {1 Basic Handlers} *)
 

cohttp-eio/src/server.ml

@@ -1,8 +1,9 @@
 open Eio.Std
+module Write = Eio.Buf_write
 
 type middleware = handler -> handler
 and handler = request -> response
-and request = Http.Request.t * Eio.Buf_read.t
+and request = Http.Request.t * Eio.Buf_read.t * Eio.Net.Sockaddr.stream
 and response = Http.Response.t * Body.t
 
 let domain_count =
@@ -12,8 +13,8 @@ let domain_count =
 
 (* Request *)
 
-let read_fixed ((request, reader) : request) =
-  match request.meth with
+let read_fixed request reader =
+  match Http.Request.meth request with
   | `POST | `PUT | `PATCH -> Body.read_fixed reader request.headers
   | _ ->
       let err =
@@ -23,8 +24,8 @@ let read_fixed ((request, reader) : request) =
       in
       raise @@ Invalid_argument err
 
-let read_chunked : request -> (Body.chunk -> unit) -> Http.Header.t =
- fun (request, reader) f -> Body.read_chunked reader request.headers f
+let read_chunked request reader f =
+  Body.read_chunked reader (Http.Request.headers request) f
 
 (* Responses *)
 
@@ -64,64 +65,61 @@ let internal_server_error_response =
 let bad_request_response =
   (Http.Response.make ~status:`Bad_request (), Body.Empty)
 
-let write_response (writer : Writer.t)
+let write_response (writer : Write.t)
     ((response, body) : Http.Response.t * Body.t) =
   let version = Http.Version.to_string response.version in
   let status = Http.Status.to_string response.status in
-  Writer.write_string writer version;
-  Writer.write_string writer " ";
-  Writer.write_string writer status;
-  Writer.write_string writer "\r\n";
+  Write.string writer version;
+  Write.string writer " ";
+  Write.string writer status;
+  Write.string writer "\r\n";
   Body.write_headers writer response.headers;
-  Writer.write_string writer "\r\n";
+  Write.string writer "\r\n";
   match body with
-  | Fixed s -> Writer.write_string writer s
+  | Fixed s -> Write.string writer s
   | Chunked chunk_writer -> Body.write_chunked writer chunk_writer
-  | Custom f ->
-      Writer.wakeup writer;
-      f (writer.sink :> Eio.Flow.sink)
+  | Custom f -> f writer
   | Empty -> ()
 
 (* main *)
 
-let rec handle_request reader writer flow handler =
+let rec handle_request client_addr reader writer flow handler =
   match Reader.http_request reader with
   | request ->
-      let response, body = handler (request, reader) in
+      let response, body = handler (request, reader, client_addr) in
       write_response writer (response, body);
-      (* A custom response needs to write the main response before calling
-         the custom function for the body. Response.write wakes the writer for
-         us if that is the case. *)
-      if not (is_custom body) then Writer.wakeup writer;
       if Http.Request.is_keep_alive request then
-        handle_request reader writer flow handler
+        handle_request client_addr reader writer flow handler
   | (exception End_of_file) | (exception Eio.Net.Connection_reset _) -> ()
-  | exception Failure _e ->
+  | exception (Failure _ as ex) ->
       write_response writer bad_request_response;
-      Writer.wakeup writer
-  | exception _ ->
+      raise ex
+  | exception ex ->
       write_response writer internal_server_error_response;
-      Writer.wakeup writer
+      raise ex
+
+let connection_handler (handler : handler) flow client_addr =
+  let reader =
+    Eio.Buf_read.of_flow ~initial_size:0x1000 ~max_size:max_int flow
+  in
+  Write.with_flow flow (fun writer ->
+      handle_request client_addr reader writer flow handler)
 
 let run_domain ssock handler =
   let on_error exn =
     Printf.fprintf stderr "Error handling connection: %s\n%!"
       (Printexc.to_string exn)
   in
+  let handler = connection_handler handler in
   Switch.run (fun sw ->
-      while true do
-        Eio.Net.accept_fork ~sw ssock ~on_error (fun flow _addr ->
-            Eio.Switch.run @@ fun sw ->
-            let reader =
-              Eio.Buf_read.of_flow ~initial_size:0x1000 ~max_size:max_int
-                (flow :> Eio.Flow.source)
-            in
-            let writer = Writer.create (flow :> Eio.Flow.sink) in
-            Eio.Fiber.fork ~sw (fun () -> Writer.run writer);
-            handle_request reader writer flow handler)
-      done)
-
-let run ?(socket_backlog = 128) ?(domains = domain_count) ~port env sw handler =
+      let rec loop () =
+        Eio.Net.accept_fork ~sw ssock ~on_error handler;
+        loop ()
+      in
+      loop ())
+
+let run ?(socket_backlog = 128) ?(domains = domain_count) ~port env handler =
+  Switch.run @@ fun sw ->
   let domain_mgr = Eio.Stdenv.domain_mgr env in
   let ssock =
     Eio.Net.listen (Eio.Stdenv.net env) ~sw ~reuse_addr:true ~reuse_port:true

cohttp-eio/tests/dune

@@ -13,6 +13,10 @@
  (modules crlf)
  (libraries unix))
 
+(mdx
+ (package cohttp-eio)
+ (packages cohttp-eio))
+
 (env
  (_
   (binaries