Abstract over sd_notify

src/cli/serve.rs

@@ -42,12 +42,12 @@ use crate::{
     },
     config::{self, Config, ListenComponent, ListenTransport},
     database::KeyValueDatabase,
-    error, observability, services,
+    error, observability, services, set_application_state,
     utils::{
         self,
         error::{Error, Result},
     },
-    Services,
+    ApplicationState, Services,
 };
 
 pub(crate) async fn run(args: ServeArgs) -> Result<(), error::ServeCommand> {
@@ -226,9 +226,7 @@ async fn run_server() -> Result<(), error::Serve> {
         }
     }
 
-    #[cfg(feature = "systemd")]
+    set_application_state(ApplicationState::Ready);
-    sd_notify::notify(true, &[sd_notify::NotifyState::Ready])
-        .expect("should be able to notify systemd");
 
     tokio::spawn(shutdown_signal(handles));
 
@@ -573,9 +571,7 @@ async fn shutdown_signal(handles: Vec<ServerHandle>) {
         handle.graceful_shutdown(Some(Duration::from_secs(30)));
     }
 
-    #[cfg(feature = "systemd")]
+    set_application_state(ApplicationState::Stopping);
-    sd_notify::notify(true, &[sd_notify::NotifyState::Stopping])
-        .expect("should be able to notify systemd");
 }
 
 async fn federation_disabled(_: Uri) -> impl IntoResponse {

src/main.rs

@@ -5,7 +5,7 @@
 use std::process::ExitCode;
 
 use clap::Parser;
-use tracing::error;
+use tracing::{error, info};
 
 mod api;
 mod cli;
@@ -41,6 +41,31 @@ fn version() -> String {
     }
 }
 
+#[derive(Debug, Clone, Copy)]
+enum ApplicationState {
+    Ready,
+    Stopping,
+}
+
+fn set_application_state(state: ApplicationState) {
+    info!(?state, "Application state changed");
+
+    #[cfg(feature = "systemd")]
+    {
+        use sd_notify::NotifyState;
+
+        fn notify(states: &[NotifyState<'_>]) {
+            sd_notify::notify(false, states)
+                .expect("should be able to notify systemd");
+        }
+
+        match state {
+            ApplicationState::Ready => notify(&[NotifyState::Ready]),
+            ApplicationState::Stopping => notify(&[NotifyState::Stopping]),
+        };
+    }
+}
+
 #[tokio::main]
 async fn main() -> ExitCode {
     let args = cli::Args::parse();