feat(portal_ws): send close stanza after pushing None

1 files changed, 11 insertions, 2 deletions

diff --git a/portal/ws/portal.ml b/portal/ws/portal.ml
index 8160e77..0d15e5a 100644
--- a/portal/ws/portal.ml
+++ b/portal/ws/portal.ml
@@ -138,9 +138,18 @@ let connect domain =
      WebSocket. *)
   let fragment_stream, fragment_push = Lwt_stream.create () in
   let stream = ws_stream |> lwt_stream |> strings_to_bytes |> parse_xml ~report |> signals
-  and push = function
+  in
+  let push = function
     | Some fs -> Markup.iter (fun f -> fragment_push (Some f)) fs
-    | None -> fragment_push None
+    | None -> begin
+        (* We need to send the [<close>] stanza in full to the WebSocket. *)
+        Markup.iter (fun f -> fragment_push (Some f)) close;
+        Lwt.async
+          (fun () ->
+            (* We drain completely the stream when closing, so the socket can close. *)
+            let+ () = Markup_lwt.drain stream
+            in fragment_push None)
+      end
   in
   (* Elements filters all elements based on the `Start_element. By simply making it
      return true every time, we get a stream of stream of elements, with each sub-stream