outbox module + fix quit handling yet again

lib/server/connection.ml

@@ -4,9 +4,8 @@ module User = Router.User
 type t = {
   addr : sockaddr;
   user : User.t;
+  outbox : Outbox.t;
   mutable pending_nick : string option;
-  outbox : Irc.Msg.t Lwt_stream.t;
-  push_outbox : (Irc.Msg.t option -> unit);
 }
 
 let make ~(router : Router.t) ~(addr : sockaddr) : t =
@@ -14,17 +13,15 @@ let make ~(router : Router.t) ~(addr : sockaddr) : t =
     | ADDR_INET (ia, _) -> Unix.string_of_inet_addr ia
     | ADDR_UNIX path -> path
   in
-  let user = User.make router ~hostname in
-  let outbox, push_outbox = Lwt_stream.create () in
-  { addr; user; pending_nick = None; outbox; push_outbox }
+  let outbox = Outbox.make () in
+  let user = User.make ~router ~hostname ~outbox in
+  { addr; user; outbox; pending_nick = None }
 
-let outbox t = Lwt_stream.choose [t.outbox; User.inbox t.user]
-let send t msg = try t.push_outbox (Some msg) with Lwt_stream.Closed -> ()
-let close t = try t.push_outbox None with Lwt_stream.Closed -> ()
+let outbox t = t.outbox
 
-let cleanup t =
-  close t;
-  User.cleanup t.user
+let shutdown t =
+  User.quit t.user;
+  Outbox.close t.outbox
 
 (* message handling *)
 
@@ -43,9 +40,12 @@ let attempt_to_register t =
     begin match User.set_nick t.user nick with
       | `nick_in_use -> `nicknameinuse nick
       | `nick_set ->
-        Logs.debug (fun m -> m "init mode: +%a" Irc.Mode.pp t.user.mode);
-        send t (Irc.Msg.make "MODE" [nick; Fmt.str "+%a" Irc.Mode.pp t.user.mode]
-                  ~prefix:(User.prefix t.user));
+        let send_mode () =
+          Outbox.send t.outbox
+            (Irc.Msg.make "MODE" [nick; Fmt.str "+%a" Irc.Mode.pp t.user.mode]
+               ~prefix:(User.prefix t.user));
+        in
+        Lwt.on_success (Lwt.pause ()) send_mode;
         `welcome
     end
   | _, _ -> `ok
@@ -80,7 +80,8 @@ let on_msg_privmsg t tgt msg _ =
   | Some dst -> Router.privmsg t.user (`user dst) msg; `ok
 
 let on_msg_join t tgt _ =
-  send t (Irc.Msg.make "JOIN" [tgt] ~prefix:(User.prefix t.user));
+  Outbox.send t.outbox
+    (Irc.Msg.make "JOIN" [tgt] ~prefix:(User.prefix t.user));
   `names ("@", tgt, ["@", User.nick t.user; "", "moe"; "", "barry"])
 
 let on_msg_privmsg t tgt msg = require_registered t (on_msg_privmsg t tgt msg)
@@ -113,7 +114,7 @@ let srv_motd_lines = [
 let rpl t cmd params =
   let prefix = Irc.Msg.Server_prefix srv_host in
   let target = User.nick t.user in
-  send t (Irc.Msg.make ~prefix cmd (target :: params))
+  Outbox.send t.outbox (Irc.Msg.make ~prefix cmd (target :: params))
 
 let rpl_tryagain t cmd = rpl t "263" [cmd; "Please wait a while and try again."]
 
@@ -172,7 +173,7 @@ let on_msg t (msg : Irc.Msg.t) : unit =
   in
   match result with
   | `ok -> ()
-  | `quit -> cleanup t
+  | `quit -> shutdown t
   | `welcome -> rpl_welcome t; rpl_motd t
   | `motd -> rpl_motd t
   | `names (cp, ch, us) -> rpl_names t cp ch us

lib/server/irc_server.ml

@@ -56,19 +56,20 @@ let writer (fd : fd) (obox : Irc.Msg.t Lwt_stream.t) : unit Lwt.t =
       | exn -> Lwt.fail exn)
 
 let handle_client (router : Router.t) (conn_fd : fd) (conn_addr : sockaddr) =
+  Logs.info (fun m -> m "new connection %a" pp_sockaddr conn_addr);
   let conn : Connection.t =
     Connection.make
       ~router
       ~addr:conn_addr
   in
-  let rd = Lwt_stream.iter (Connection.on_msg conn) (reader conn_fd) in
-  let wr = writer conn_fd (Connection.outbox conn) in
+  let reader = Lwt_stream.iter (Connection.on_msg conn) (reader conn_fd) in
+  let writer = writer conn_fd (Outbox.stream (Connection.outbox conn)) in
+  let shutdown () = Connection.shutdown conn in
+  Lwt.on_termination reader shutdown;
+  Lwt.on_termination writer shutdown;
   Lwt.finalize
+    (fun () -> writer)
     (fun () ->
-       Logs.info (fun m -> m "new connection %a" pp_sockaddr conn_addr);
-       Lwt.choose [rd; wr])
-    (fun () ->
-       Connection.cleanup conn;
        Lwt_unix.close conn_fd >|= fun () ->
        Logs.info (fun m -> m "connection closed %a" pp_sockaddr conn_addr))
 

lib/server/outbox.ml

@@ -0,0 +1,14 @@
+open! Import
+
+type t = {
+  stream : Irc.Msg.t Lwt_stream.t;
+  push : Irc.Msg.t option -> unit;
+}
+
+let make () =
+  let stream, push = Lwt_stream.create () in
+  { stream; push }
+
+let stream t = t.stream
+let send t msg = try t.push (Some msg) with Lwt_stream.Closed -> ()
+let close t = try t.push None with Lwt_stream.Closed -> ()

lib/server/router.ml

@@ -8,12 +8,11 @@ type t = {
 and user = {
   router : t;
   hostname : string;
+  outbox : Outbox.t;
   mutable key : string_ci;
   mutable nick : Irc.nick;
   mutable userinfo : Irc.userinfo option;
   mutable mode : Irc.Mode.t;
-  inbox : Irc.Msg.t Lwt_stream.t;
-  push_inbox : (Irc.Msg.t option -> unit);
 }
 
 let make () =
@@ -25,8 +24,7 @@ let find_user t nick =
 module User = struct
   type t = user
 
-  let make router ~hostname =
-    let inbox, push_inbox = Lwt_stream.create () in
+  let make ~router ~hostname ~outbox =
     {
       router;
       hostname;
@@ -34,17 +32,14 @@ module User = struct
       nick = "*";
       userinfo = None;
       mode = Irc.Mode.of_string "iw";
-      inbox; push_inbox;
+      outbox;
     }
 
+  let outbox t = t.outbox
   let nick t = t.nick
   let prefix t = Irc.Msg.User_prefix (t.nick, t.userinfo, Some t.hostname)
   let is_registered t = t.key <> empty_string_ci
 
-  let inbox t = t.inbox
-  let send t msg = try t.push_inbox (Some msg) with Lwt_stream.Closed -> ()
-  let close t = try t.push_inbox None with Lwt_stream.Closed -> ()
-
   let unregister t =
     Hashtbl.remove t.router.users t.key;
     t.key <- empty_string_ci
@@ -56,9 +51,10 @@ module User = struct
     else begin
       ((* TODO: relay NICK message *));
       if is_registered t then
-        send t (Irc.Msg.make "NICK" [new_nick]
-                  ~prefix:(prefix t)
-                  ~always_trailing:true);
+        Outbox.send t.outbox
+          (Irc.Msg.make "NICK" [new_nick]
+             ~prefix:(prefix t)
+             ~always_trailing:true);
       unregister t;
       Hashtbl.add t.router.users key t;
       t.key <- key;
@@ -66,16 +62,16 @@ module User = struct
       `nick_set
     end
 
-  let cleanup t =
+  let quit t =
     if is_registered t then begin
       (* TODO: quit reason *)
-      send t (Irc.Msg.make "QUIT" ["Closed"]
-                ~prefix:(prefix t)
-                ~always_trailing:true);
+      Outbox.send t.outbox
+        (Irc.Msg.make "QUIT" ["Closed"]
+           ~prefix:(prefix t)
+           ~always_trailing:true);
       (* TODO: relay QUIT message *)
       unregister t
-    end;
-    close t
+    end
 end
 
 let privmsg src dst txt =
@@ -83,4 +79,4 @@ let privmsg src dst txt =
   match dst with
   | `user dst ->
     let msg = Irc.Msg.make "PRIVMSG" [User.nick dst; txt] ~prefix in
-    User.send dst msg
+    Outbox.send (User.outbox dst) msg