diff --git a/src/service_main.rs b/src/service_main.rs index 6d85242..92b1876 100644 --- a/src/service_main.rs +++ b/src/service_main.rs @@ -1,6 +1,7 @@ use std::{convert::Infallible, env}; use warp::Filter; -use tokio::runtime::Runtime; +use tokio::{runtime::Runtime, sync::oneshot::{self, Receiver}}; +use tiberius::Config; use crate::{ handlers, problem, @@ -8,9 +9,13 @@ use crate::{ DBPool, connection_manager::TiberiusConnectionManager }; -use tiberius::Config; pub fn run() { + let (_tx, rx) = oneshot::channel::<()>(); + run_with_graceful_shutdown(rx); +} + +pub fn run_with_graceful_shutdown(shutdown_rx: Receiver) where T: Send + 'static { if env::var_os("RUST_LOG").is_none() { // Set `RUST_LOG=todos=debug` to see debug logs, // this only shows access logs. @@ -20,7 +25,7 @@ pub fn run() { // Create the runtime let mut rt = Runtime::new().unwrap(); - + // Spawn the root task rt.block_on(async { let manager = TiberiusConnectionManager::new(Config::from_ado_string(SERVICE_CONFIG.get_connection_string()).unwrap()).unwrap(); @@ -33,11 +38,14 @@ pub fn run() { info!(target: "service", "Listening on {}", SERVICE_CONFIG.get_addr()); - warp::serve(api) + let (_addr, server) = warp::serve(api) //.unstable_pipeline() //.run(([127, 0, 0, 1], service_config.get_port())) - .run(SERVICE_CONFIG.get_addr()) - .await; + //.run(SERVICE_CONFIG.get_addr()) + .bind_with_graceful_shutdown(SERVICE_CONFIG.get_addr(), async move { + shutdown_rx.await.ok(); + }); + server.await; }); } diff --git a/src/windows_service_main.rs b/src/windows_service_main.rs index 9064c27..cdf7427 100644 --- a/src/windows_service_main.rs +++ b/src/windows_service_main.rs @@ -1,6 +1,5 @@ use std::{ ffi::OsString, - sync::mpsc, time::Duration }; @@ -14,6 +13,8 @@ use windows_service::{ service_dispatcher, Result, }; +use tokio::sync::oneshot; + use crate::service_main; // Generate the windows service boilerplate. @@ -38,7 +39,8 @@ const SERVICE_TYPE: ServiceType = ServiceType::OWN_PROCESS; pub fn run_service() -> Result<()> { // Create a channel to be able to poll a stop event from the service worker loop. - let (shutdown_tx, shutdown_rx) = mpsc::channel(); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let mut mut_tx = Some(shutdown_tx); // this is added to make event_handler satisfy FnOnce // Define system service event handler that will be receiving service events. let event_handler = move |control_event| -> ServiceControlHandlerResult { @@ -49,7 +51,7 @@ pub fn run_service() -> Result<()> { // Handle stop ServiceControl::Stop => { - shutdown_tx.send(()).unwrap(); + let _ = mut_tx.take().unwrap().send(()).ok(); ServiceControlHandlerResult::NoError } @@ -72,7 +74,7 @@ pub fn run_service() -> Result<()> { process_id: None, })?; - service_main::run(); + service_main::run_with_graceful_shutdown(shutdown_rx); // Tell the system that service has stopped. status_handle.set_service_status(ServiceStatus {