Make wait for server ack more resilient in Lwt client.

paurkedal · Mar 15, 2022 · 4a894fb · 4a894fb
1 parent a87325e
commit 4a894fb
Showing 1 changed file with 13 additions and 11 deletions.
diff --git a/lwt/websocket_lwt_unix.ml b/lwt/websocket_lwt_unix.ml
@@ -130,17 +130,19 @@ let with_connection ?(extra_headers = Cohttp.Header.init ())
   let close_and_wait_for_remote_ack close =
     write_frame close
     >>= fun () ->
-    Lwt.pick
-      [ Lwt_unix.timeout 2.;
-        ( read_frame ()
-        >>= function
-        | {opcode= Close; _} -> Lwt.return_unit
-        | x ->
-            Lwt_log.warning_f ~section
-              "Client initiated close: expected a close frame from server, got \
-               %s"
-              (Websocket.Frame.show x) ) ]
-    >>= fun () -> Lwt_io.close oc in
+    let rec wait_for_ack () =
+      read_frame ()
+      >>= function
+      | {opcode= Close; _} -> Lwt.return_unit
+      | x ->
+          Lwt_log.warning_f ~section
+            "Client initiated close: expected a close frame from server, got \
+             %s"
+            (Websocket.Frame.show x)
+          >>= wait_for_ack in
+    Lwt.finalize
+      (fun () -> Lwt.pick [ Lwt_unix.timeout 2.; wait_for_ack () ])
+      (fun () -> Lwt_io.close oc) in
   let write_frame = function
     | `Continue frame -> write_frame frame
     | `Stop None ->