From e2967fabd9beef8204523b12b01ded5cc5224f22 Mon Sep 17 00:00:00 2001
From: tali <iitalics@gmail.com>
Date: Thu, 1 Feb 2024 14:01:42 -0500
Subject: [PATCH] implement sd-notify protocol

---
 bin/main.ml          | 19 +++++++++++++++++++
 lib/server/server.ml | 44 ++++++++++++++++++++++++++++++--------------
 talircd.service      |  2 ++
 3 files changed, 51 insertions(+), 14 deletions(-)

diff --git a/bin/main.ml b/bin/main.ml
index 3f1d70a..60b4386 100644
--- a/bin/main.ml
+++ b/bin/main.ml
@@ -23,6 +23,22 @@ let () =
       ~timestamp:(not no_timestamp)
       ~namespace:(not no_namespace)
 
+let sd_notify_sock =
+  Option.map
+    (fun path ->
+       let sock_fd = Unix.socket PF_UNIX SOCK_DGRAM 0 ~cloexec:true in
+       let dest = Unix.ADDR_UNIX path in
+       sock_fd, dest)
+    (Sys.getenv_opt "NOTIFY_SOCKET")
+
+let sd_notify msg =
+  Option.iter
+    (fun (sock_fd, dest) ->
+       let dgram = Bytes.of_string msg in
+       Unix.sendto sock_fd dgram 0 (Bytes.length dgram) [] dest |> ignore)
+    sd_notify_sock
+
+
 (* TODO: s-exp/json/toml config format *)
 
 let port =
@@ -50,6 +66,9 @@ let config : Server.config = {
   ping_interval = 60;
   whowas_history_len = 1000;
   motd_file;
+  notify = function
+    | `ready -> sd_notify "READY=1"
+    | `stopping -> sd_notify "STOPPING=1"
 }
 
 let () =
diff --git a/lib/server/server.ml b/lib/server/server.ml
index 987bdbb..a8b781f 100644
--- a/lib/server/server.ml
+++ b/lib/server/server.ml
@@ -4,9 +4,23 @@ open Lwt.Infix
 
 include (val Logging.sublogs logger "Server")
 
+type config = {
+  port : int;
+  listen_backlog : int;
+  ping_interval : int;
+  whowas_history_len : int;
+  hostname : string;
+  motd_file : string;
+  notify : [`ready | `stopping] -> unit;
+}
+
 type ping_wheel = Connection.t Wheel.t
 
-let listener ~(port : int) ~(listen_backlog : int) : (fd * sockaddr) Lwt_stream.t =
+let listener
+      ~(port : int)
+      ~(listen_backlog : int)
+      ~(on_ready : unit -> unit)
+  : (fd * sockaddr) Lwt_stream.t =
   let sock : fd Lwt.t =
     let fd = Lwt_unix.socket PF_INET SOCK_STREAM 0 in
     Lwt_unix.setsockopt fd SO_KEEPALIVE false;
@@ -14,6 +28,7 @@ let listener ~(port : int) ~(listen_backlog : int) : (fd * sockaddr) Lwt_stream.
     let srv_adr = Unix.ADDR_INET (Unix.inet_addr_any, port) in
     let* () = Lwt_unix.bind fd srv_adr in
     Lwt_unix.listen fd listen_backlog;
+    on_ready ();
     info (fun m -> m "listening on %a" pp_sockaddr srv_adr);
     Lwt.return fd
   in
@@ -92,18 +107,16 @@ let handle_client
       (fun e -> error (fun m -> m "%a:@ %a" pp_sockaddr conn_addr Fmt.exn e));
   end
 
-type config = {
-  port : int;
-  listen_backlog : int;
-  ping_interval : int;
-  whowas_history_len : int;
-  hostname : string;
-  motd_file : string;
-}
-
-let run { port; listen_backlog; ping_interval; whowas_history_len;
-          hostname; motd_file } : unit Lwt.t
-  =
+let run {
+      port;
+      listen_backlog;
+      ping_interval;
+      whowas_history_len;
+      hostname;
+      motd_file;
+      notify;
+    }
+  : unit Lwt.t =
   debug (fun m -> m "ping interval:@ %ds" ping_interval);
   debug (fun m -> m "whowas history:@ %d" whowas_history_len);
 
@@ -164,7 +177,10 @@ let run { port; listen_backlog; ping_interval; whowas_history_len;
       on_con
       (listener
          ~port
-         ~listen_backlog)
+         ~listen_backlog
+         ~on_ready:(fun () -> notify `ready))
   in
 
+  (* TODO: graceful cleanup on ctrl-c *)
+
   listener_promise <&> pinger_promise
diff --git a/talircd.service b/talircd.service
index cac8450..c50bc31 100644
--- a/talircd.service
+++ b/talircd.service
@@ -6,4 +6,6 @@ Environment=IRC_HOSTNAME=irc.tali.software
 Environment=IRC_MOTD=/usr/local/share/talircd/motd
 #Environment=IRC_PORT=6667
 #Environment=LOG_LEVEL=DEBUG
+
+Type=notify
 ExecStart=/usr/local/bin/talircd