Added shutdown signal for service

SpeedSX · Jul 6, 2020 · 67c9cc1 · 67c9cc1
1 parent 791234e
commit 67c9cc1
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 10 deletions.
diff --git a/src/service_main.rs b/src/service_main.rs
@@ -1,16 +1,21 @@
 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,
     configuration::SERVICE_CONFIG, 
     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<T>(shutdown_rx: Receiver<T>) 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
@@ -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 {