From 39880cc6ace0a737928a2d1bab942839fca41d09 Mon Sep 17 00:00:00 2001
From: Lambda <csgl@xiretza.xyz>
Date: Thu, 26 Sep 2024 16:27:58 -0700
Subject: [PATCH] Abstract over sd_notify

---
 src/cli/serve.rs | 12 ++++--------
 src/main.rs      | 27 ++++++++++++++++++++++++++-
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/src/cli/serve.rs b/src/cli/serve.rs
index 0e646dc0..38832e6b 100644
--- a/src/cli/serve.rs
+++ b/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")]
-    sd_notify::notify(true, &[sd_notify::NotifyState::Ready])
-        .expect("should be able to notify systemd");
+    set_application_state(ApplicationState::Ready);
 
     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")]
-    sd_notify::notify(true, &[sd_notify::NotifyState::Stopping])
-        .expect("should be able to notify systemd");
+    set_application_state(ApplicationState::Stopping);
 }
 
 async fn federation_disabled(_: Uri) -> impl IntoResponse {
diff --git a/src/main.rs b/src/main.rs
index 903906a6..e0365fc3 100644
--- a/src/main.rs
+++ b/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();