diff options
| -rw-r--r-- | src/bot.rs | 32 | ||||
| -rw-r--r-- | src/windows_util.rs | 105 |
2 files changed, 102 insertions, 35 deletions
@@ -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 { |