Flush response before shutting down

lib/server_connection.ml

@@ -139,16 +139,21 @@ let create ?(config=Config.default) ?(error_handler=default_error_handler) reque
   }
 
 let shutdown_reader t =
-  Reader.force_close t.reader;
   if is_active t
-  then Reqd.close_request_body (current_reqd_exn t)
-  else wakeup_reader t
+  then Reqd.close_request_body (current_reqd_exn t);
+  Reader.force_close t.reader;
+  wakeup_reader t
 
 let shutdown_writer t =
+  if is_active t then (
+    let reqd = current_reqd_exn t in
+    (* XXX(dpatti): I'm not sure I understand why we close the *request* body
+       here. Maybe we can write a test such that removing this line causes it to
+       fail? *)
+    Reqd.close_request_body reqd;
+    Reqd.flush_response_body reqd);
   Writer.close t.writer;
-  if is_active t
-  then Reqd.close_request_body (current_reqd_exn t)
-  else wakeup_writer t
+  wakeup_writer t
 
 let error_code t =
   if is_active t
@@ -157,9 +162,7 @@ let error_code t =
 
 let shutdown t =
   shutdown_reader t;
-  shutdown_writer t;
-  wakeup_reader t;
-  wakeup_writer t
+  shutdown_writer t
 
 let set_error_and_handle ?request t error =
   if is_active t then begin

lib_test/test_server_connection.ml

@@ -1013,6 +1013,26 @@ let test_shutdown_during_asynchronous_request () =
   writer_closed t
 ;;
 
+let test_flush_response_before_shutdown () =
+  let request = Request.create `GET "/" ~headers:(Headers.encoding_fixed 0) in
+  let response = Response.create `OK ~headers:Headers.encoding_chunked in
+  let continue = ref (fun () -> ()) in
+  let request_handler reqd =
+    let body = Reqd.respond_with_streaming ~flush_headers_immediately:true reqd response in
+    continue := (fun () ->
+      Body.Writer.write_string body "hello world";
+      Body.Writer.close body);
+  in
+  let t = create request_handler in
+  read_request t request;
+  write_response t response;
+  !continue ();
+  shutdown t;
+  raises_writer_closed (fun () ->
+    write_string t "b\r\nhello world\r\n";
+    connection_is_shutdown t);
+;;
+
 let test_schedule_read_with_data_available () =
   let response = Response.create `OK in
   let body = ref None in
@@ -1084,5 +1104,6 @@ let tests =
   ; "response finished before body read", `Quick, test_response_finished_before_body_read
   ; "shutdown in request handler", `Quick, test_shutdown_in_request_handler
   ; "shutdown during asynchronous request", `Quick, test_shutdown_during_asynchronous_request
+  ; "flush response before shutdown", `Quick, test_flush_response_before_shutdown
   ; "schedule read with data available", `Quick, test_schedule_read_with_data_available
   ]