fix(portal): close socket only after all elements have been processed

1 files changed, 6 insertions, 4 deletions

diff --git a/portal/tcp/portal.ml b/portal/tcp/portal.ml
index 76fbbc8..854f76c 100644
--- a/portal/tcp/portal.ml
+++ b/portal/tcp/portal.ml
@@ -91,12 +91,14 @@ let connect (domain : string) : t Lwt.t =
   in
   let xml_stream, xml_push = Lwt_stream.create () in
   let push msg =
-    let none () = xml_push None; close tcp_socket
+    let none () = xml_push None
     and some fragments () =
-      Markup.iter (fun f -> xml_push (Some f)) fragments |> Lwt.return
-    in Option.fold ~none ~some msg |> Lwt.async
+      Markup.iter (fun f -> xml_push (Some f)) fragments
+    in Option.fold ~none ~some msg ()
   and report loc err = raise (MalformedStanza (loc, err)) in
   let open Markup_lwt in
   let stream = tcp_stream |> lwt_stream |> strings_to_bytes |> parse_xml ~report |> signals
-  in Lwt.async (fun () -> lwt_stream xml_stream |> Markup_lwt.write_xml |> iter send);
+  in Lwt.async (fun () ->
+         let* _ = lwt_stream xml_stream |> Markup_lwt.write_xml |> iter send
+         in Lwt_unix.close tcp_socket);
      stream, push