From 6117a71ec6c09d9ad9b29194704bb52b09bce399 Mon Sep 17 00:00:00 2001
From: David House <dhouse@janestreet.com>
Date: Tue, 27 Apr 2021 11:05:56 +0100
Subject: [PATCH] changes

---
 lib/client_connection.ml           |  6 +++++-
 lib_test/test_client_connection.ml | 23 ++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/lib/client_connection.ml b/lib/client_connection.ml
index e51c5533..bab4e6af 100644
--- a/lib/client_connection.ml
+++ b/lib/client_connection.ml
@@ -195,12 +195,16 @@ module Oneshot = struct
   let next_write_operation t =
     flush_request_body t;
     if Body.is_closed t.request_body
+    (* Even though we've just done [flush_request_body], it might still be the case that
+       [Body.has_pending_output] returns true, because it does so when we've written all
+       output except for the final chunk. *)
+    && not (Body.has_pending_output t.request_body)
     then Writer.close t.writer;
     Writer.next t.writer
   ;;
 
   let yield_writer t k =
-    if Body.is_closed t.request_body
+    if Body.is_closed t.request_body && not (Body.has_pending_output t.request_body)
     then begin
       Writer.close t.writer;
       k ()
diff --git a/lib_test/test_client_connection.ml b/lib_test/test_client_connection.ml
index 4075c590..838c2651 100644
--- a/lib_test/test_client_connection.ml
+++ b/lib_test/test_client_connection.ml
@@ -124,7 +124,27 @@ let test_get () =
   Body.close_writer body;
   write_request  t request';
   read_response  t response;
-  read_string    t "d\r\nHello, world!\r\n0\r\n\r\n";
+  read_string    t "d\r\nHello, world!\r\n0\r\n\r\n"
+;;
+
+let test_send_streaming_body () =
+  let request' = Request.create `GET "/" ~headers:Headers.encoding_chunked in
+  let response = Response.create `OK ~headers:Headers.encoding_chunked in
+  let body, t =
+    request
+      request'
+      ~response_handler:(default_response_handler response)
+      ~error_handler:no_error_handler
+  in
+  write_request  t request';
+  read_response  t response;
+  Body.write_string body "hello";
+  write_string t "5\r\nhello\r\n";
+  Body.write_string body "world";
+  Body.close_writer body;
+  write_string t "5\r\nworld\r\n";
+  write_string t "0\r\n\r\n";
+  writer_closed t
 ;;
 
 let test_response_eof () =
@@ -259,6 +279,7 @@ let test_failed_response_parse () =
 
 let tests =
   [ "GET"         , `Quick, test_get
+  ; "send streaming body", `Quick, test_send_streaming_body
   ; "Response EOF", `Quick, test_response_eof
   ; "Response header order preserved", `Quick, test_response_header_order
   ; "report_exn"  , `Quick, test_report_exn