Fix reader getting stuck (never waking up) when using

CHANGES.md

@@ -1,34 +1,34 @@
 Unreleased
 --------------
 
-- httpaf-lwt-unix: replace the `dune` file (previously written in OCaml) with a
-  `(select)` form to avoid depending on `ocamlfind`
-  ([#18](https://github.com/anmonteiro/httpaf/pull/18))
-- httpaf-mirage: depend on `mirage-conduit` instead of `conduit-mirage`,
-  effectively placing a lower bound of OCaml 4.07 on httpaf-mirage
-  ([#16](https://github.com/anmonteiro/httpaf/pull/16))
-- httpaf-lwt, httpaf-lwt-unix, httpaf-mirage: deduplicate interface code via a
-  common `Httpaf_lwt_intf` interface
-  ([#13](https://github.com/anmonteiro/httpaf/pull/13))
-- httpaf, httpaf-lwt, httpaf-lwt-unix, httpaf-async, httpaf-mirage: add support
-  for persistent connections and pipelining in the client implementations
-  ([#5](https://github.com/anmonteiro/httpaf/pull/5))
-- httpaf, httpaf-lwt, httpaf-lwt-unix, httpaf-async, httpaf-mirage: add support
-  for switching protocols, e.g. to a WebSocket connection, via a new function
-  `Reqd.respond_with_upgrade`
-  ([#8](https://github.com/anmonteiro/httpaf/pull/8))
 - httpaf-lwt-unix: add support for HTTPS in the Lwt runtime, either via OpenSSL
   or `ocaml-tls` ([#2](https://github.com/anmonteiro/httpaf/pull/2))
 - httpaf, httpaf-lwt, httpaf-lwt-unix, httpaf-mirage: Add a Mirage adapter
   ([#3](https://github.com/anmonteiro/httpaf/pull/3))
 - Add an `esy.json` file ([#6](https://github.com/anmonteiro/httpaf/pull/6))
 - httpaf: Catch parsing errors and hand them to `error_handler` in
   `Server_connection` ([#4](https://github.com/anmonteiro/httpaf/pull/4))
+- httpaf, httpaf-lwt, httpaf-lwt-unix, httpaf-async, httpaf-mirage: add support
+  for persistent connections and pipelining in the client implementations
+  ([#5](https://github.com/anmonteiro/httpaf/pull/5))
+- httpaf, httpaf-lwt, httpaf-lwt-unix, httpaf-async, httpaf-mirage: add support
+  for switching protocols, e.g. to a WebSocket connection, via a new function
+  `Reqd.respond_with_upgrade`
+  ([#8](https://github.com/anmonteiro/httpaf/pull/8))
+- httpaf-lwt, httpaf-lwt-unix, httpaf-mirage: deduplicate interface code via a
+  common `Httpaf_lwt_intf` interface
+  ([#13](https://github.com/anmonteiro/httpaf/pull/13))
+- httpaf-mirage: depend on `mirage-conduit` instead of `conduit-mirage`,
+  effectively placing a lower bound of OCaml 4.07 on httpaf-mirage
+  ([#16](https://github.com/anmonteiro/httpaf/pull/16))
+- httpaf-lwt-unix: replace the `dune` file (previously written in OCaml) with a
+  `(select)` form to avoid depending on `ocamlfind`
+  ([#18](https://github.com/anmonteiro/httpaf/pull/18))
+- httpaf: Shutdown the writer after closing a non chunk-encoded request body on
+  the client ([#23](https://github.com/anmonteiro/httpaf/pull/23))
 - httpaf-mirage: Adapt to Mirage 3.7 interfaces. `httpaf_mirage` now requires
   `conduit-mirage` >= 2.0.2 and `mirage-flow` >= 2.0.0
   ([#24](https://github.com/anmonteiro/httpaf/pull/24))
-- httpaf: Shutdown the writer after closing a non chunk-encoded request body on
-  the client ([#23](https://github.com/anmonteiro/httpaf/pull/23))
 - httpaf, httpaf-lwt, httpaf-lwt-unix, httpaf-async, httpaf-mirage: after
   switching protocols, close the connection / file descriptors when the upgrade
   handler's returned promise resolves
@@ -48,6 +48,11 @@ Unreleased
   advertised in the interface file. They are now served to the respective
   handlers in the original transmission order
   ([#32](https://github.com/anmonteiro/httpaf/pull/32))
+- httpaf: Fix persistent connections getting stuck (reader never waking up)
+  when using `~flush_headers_immediately:true` in combination with an empty
+  response body ([#34](https://github.com/anmonteiro/httpaf/pull/34)). This is
+  a fix for an issue opened in the upstream repo:
+  [inhabitedtype/httpaf#162](https://github.com/inhabitedtype/httpaf/issues/162)
 
 httpaf (upstream) 0.6.5
 --------------

lib/body.ml

@@ -157,6 +157,20 @@ let when_ready_to_write t callback =
 
 let transfer_to_writer_with_encoding t ~encoding writer =
   let faraday = t.faraday in
+  begin match t.write_final_if_chunked, encoding with
+  | true, (`Fixed _ | `Close_delimited) ->
+    (* Play nicely with [has_pending_output] in the case of a fixed or
+       close-delimited encoding.
+
+       Immediately set `t.write_final_if_chunked` to `false` because we may
+       not have another opportunity to do so before advancing the request
+       queue. *)
+    t.write_final_if_chunked <- false;
+  | false, (`Fixed _ | `Close_delimited)
+  | _, `Chunked ->
+    (* Handled explicitly later when closing the writer. *)
+    ()
+  end;
   begin match Faraday.operation faraday with
   | `Yield -> ()
   | `Close ->
@@ -178,9 +192,6 @@ let transfer_to_writer_with_encoding t ~encoding writer =
       buffered := !buffered + lengthv;
       begin match encoding with
       | `Fixed _ | `Close_delimited ->
-        (* Play nicely with [has_pending_output] in the case of a fixed or
-           close-delimited encoding. *)
-        t.write_final_if_chunked <- false;
         Serialize.Writer.schedule_fixed writer iovecs
       | `Chunked ->
         Serialize.Writer.schedule_chunk writer iovecs

lib_test/test_httpaf.ml

@@ -836,6 +836,46 @@ Accept-Language: en-US,en;q=0.5\r\n\r\n";
     writer_closed t;
 	;;
 
+ let test_immediate_flush_empty_body () =
+    let reader_woken_up = ref false in
+    let response = Response.create `OK in
+    let request_handler reqd =
+      let resp_body = Reqd.respond_with_streaming
+        ~flush_headers_immediately:true reqd response
+      in
+      Body.close_writer resp_body;
+    in
+    let t = create ~error_handler request_handler in
+    reader_ready t;
+    yield_writer t (fun () ->
+      write_response t ~body:"" response);
+    read_request t (Request.create `GET "/");
+    yield_reader t (fun () -> reader_woken_up := true);
+    writer_yielded t;
+    Alcotest.(check bool) "Reader woken up"
+      true !reader_woken_up;
+	;;
+
+  let test_empty_body_no_immediate_flush () =
+    let reader_woken_up = ref false in
+    let response = Response.create `OK in
+    let request_handler reqd =
+      let resp_body = Reqd.respond_with_streaming
+        ~flush_headers_immediately:false reqd response
+      in
+      Body.close_writer resp_body;
+    in
+    let t = create ~error_handler request_handler in
+    reader_ready t;
+    writer_yielded t;
+    read_request t (Request.create `GET "/");
+    yield_reader t (fun () -> reader_woken_up := true);
+    write_response t ~body:"" response;
+    writer_yielded t;
+    Alcotest.(check bool) "Reader woken up"
+      true !reader_woken_up;
+	;;
+
   let tests =
     [ "initial reader state"  , `Quick, test_initial_reader_state
     ; "shutdown reader closed", `Quick, test_reader_is_closed_after_eof
@@ -859,6 +899,8 @@ Accept-Language: en-US,en;q=0.5\r\n\r\n";
     ; "respond with upgrade", `Quick, test_respond_with_upgrade
     ; "writer unexpected eof", `Quick, test_unexpected_eof
     ; "input shrunk", `Quick, test_input_shrunk
+    ; "`flush_headers_immediately` with empty body", `Quick, test_immediate_flush_empty_body
+    ; "empty body with no immediate flush", `Quick, test_empty_body_no_immediate_flush
     ]
 end