windows: stop postgres if we started it

2 files changed, 102 insertions, 35 deletions

diff --git a/src/bot.rs b/src/bot.rs
index 836fb82..de31c73 100644
--- a/src/bot.rs
+++ b/src/bot.rs
@@ -396,9 +396,9 @@ fn after_handle(ctx: PoiseContext) -> BoxFuture<()> {
 
 pub async fn run() -> anyhow::Result<()> {
     #[cfg(all(windows, feature = "windows_autostart_postgres"))]
-    unsafe {
-        crate::windows_util::ensure_postgres_started()?;
-    };
+    let started_pg =
+        tokio::task::spawn_blocking(|| unsafe { crate::windows_util::ensure_postgres_started() })
+            .await??;
 
     let token = &CONFIG.discord.auth.token;
 
@@ -414,16 +414,26 @@ pub async fn run() -> anyhow::Result<()> {
 
     let shard_manager = client.shard_manager.clone();
 
-    tokio::spawn(async move {
-        tokio::signal::ctrl_c().await.unwrap();
-        warn!("got ^C");
-
-        shard_manager.shutdown_all().await;
-        info!("shutdown");
+    let run_handle = tokio::spawn(async move {
+        info!("connecting to discord");
+        client.start().await.expect("running discord client");
     });
 
-    info!("starting bot");
-    client.start().await?;
+    tokio::signal::ctrl_c().await?;
+    warn!("got ^C, gracefully halting discord");
+
+    shard_manager.shutdown_all().await;
+    run_handle.await?;
+    info!("discord shutdown");
+
+    #[cfg(all(windows, feature = "windows_autostart_postgres"))]
+    unsafe {
+        if started_pg {
+            log::info!("we started postgres, stopping it before shutdown");
+
+            tokio::task::spawn_blocking(|| crate::windows_util::shutdown_postgres()).await??;
+        }
+    }
 
     Ok(())
 }
diff --git a/src/windows_util.rs b/src/windows_util.rs
index 348ea91..3b1339e 100644
--- a/src/windows_util.rs
+++ b/src/windows_util.rs
@@ -1,5 +1,8 @@
 use windows::{
-    core::HSTRING,
+    core::{
+        HSTRING,
+        PWSTR,
+    },
     Win32::{
         Foundation,
         System::Services,
@@ -10,43 +13,26 @@ lazy_static::lazy_static! {
     static ref PSQL_REGEX: regex::Regex = regex::Regex::new(r#"^(?:postgresql|psql|postgres)-.*([\d\.]*)"#).unwrap();
 }
 
-pub unsafe fn ensure_postgres_started() -> windows::core::Result<()> {
+pub unsafe fn ensure_postgres_started() -> windows::core::Result<bool> {
     let sc_manager = Services::OpenSCManagerW(
         None,
         None,
         (Foundation::GENERIC_READ | Foundation::GENERIC_WRITE).0,
     )?;
 
-    let services = list_services(sc_manager)?;
-
-    let mut psql_services =
-        services.into_iter().filter(|status| PSQL_REGEX.is_match(&status.name)).collect::<Vec<_>>();
-
-    match psql_services.len() {
-        1 => {},
-
-        0 => {
-            log::warn!("unable to identify an installed postgres service");
-        },
-
-        other => {
-            log::warn!("unable to identify postgres service: {other} potential matches found");
-            return Ok(());
-        },
-    }
-
-    let svc = psql_services.remove(0);
-    log::debug!("identified postgres service: {}", svc.name);
+    let Some(svc) = get_psql_service(sc_manager)? else {
+        return Ok(false);
+    };
 
     match svc.status.dwCurrentState {
         Services::SERVICE_RUNNING => {
             log::info!("postgres was already running, done");
-            return Ok(());
+            return Ok(false);
         },
 
         Services::SERVICE_START_PENDING => {
             log::info!("postgres was already in startup, done");
-            return Ok(());
+            return Ok(false);
         },
 
         _ => {
@@ -64,9 +50,80 @@ pub unsafe fn ensure_postgres_started() -> windows::core::Result<()> {
     Services::CloseServiceHandle(service)?;
     Services::CloseServiceHandle(sc_manager)?;
 
+    Ok(true)
+}
+
+pub unsafe fn shutdown_postgres() -> windows::core::Result<()> {
+    let sc_manager = Services::OpenSCManagerW(
+        None,
+        None,
+        (Foundation::GENERIC_READ | Foundation::GENERIC_WRITE).0,
+    )?;
+
+    let Some(svc) = get_psql_service(sc_manager)? else {
+        log::warn!("wasn't able to find postgres service");
+        return Ok(());
+    };
+
+    if svc.status.dwCurrentState != Services::SERVICE_RUNNING {
+        log::warn!("postgres wasn't in running state, not attempting to shut it down");
+        return Ok(());
+    }
+
+    let service =
+        Services::OpenServiceW(sc_manager, &HSTRING::from(&svc.name), Services::SERVICE_STOP)?;
+
+    let mut params = Services::SERVICE_CONTROL_STATUS_REASON_PARAMSW {
+        dwReason:      Services::SERVICE_STOP_REASON_FLAG_PLANNED
+            | Services::SERVICE_STOP_REASON_MAJOR_NONE
+            | Services::SERVICE_STOP_REASON_MINOR_NONE,
+        pszComment:    PWSTR::null(),
+        ServiceStatus: Services::SERVICE_STATUS_PROCESS::default(),
+    };
+
+    Services::ControlServiceExW(
+        service,
+        Services::SERVICE_CONTROL_STOP,
+        Services::SERVICE_CONTROL_STATUS_REASON_INFO,
+        &mut params as *mut _ as _,
+    )?;
+
+    log::info!("stopped postgres service");
+
+    Services::CloseServiceHandle(service)?;
+    Services::CloseServiceHandle(sc_manager)?;
+
     Ok(())
 }
 
+pub unsafe fn get_psql_service(
+    sc_manager: Services::SC_HANDLE,
+) -> windows::core::Result<Option<ServiceStatus>> {
+    let services = list_services(sc_manager)?;
+
+    let mut psql_services =
+        services.into_iter().filter(|status| PSQL_REGEX.is_match(&status.name)).collect::<Vec<_>>();
+
+    match psql_services.len() {
+        1 => {},
+
+        0 => {
+            log::warn!("unable to identify an installed postgres service");
+            return Ok(None);
+        },
+
+        other => {
+            log::warn!("unable to identify postgres service: {other} potential matches found");
+            return Ok(None);
+        },
+    }
+
+    let svc = psql_services.remove(0);
+    log::debug!("identified postgres service: {}", svc.name);
+
+    Ok(Some(svc))
+}
+
 const RESULT_SIZE: usize = size_of::<Services::ENUM_SERVICE_STATUS_PROCESSW>();
 
 pub struct ServiceStatus {