From 6da68dc0426afcc74ec0c382a8673c521780c374 Mon Sep 17 00:00:00 2001
From: Brian Picciano <mediocregopher@gmail.com>
Date: Sun, 18 Jun 2023 14:28:46 +0200
Subject: [PATCH] stop domain manager's inner tasks on shutdown

---
 src/domain/manager.rs | 12 +++++-------
 src/main.rs           | 15 +++++++++++++++
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/src/domain/manager.rs b/src/domain/manager.rs
index bf8c09b..00afda2 100644
--- a/src/domain/manager.rs
+++ b/src/domain/manager.rs
@@ -139,7 +139,7 @@ pub trait Manager: Sync + Send {
     fn all_domains(&self) -> Result<Vec<domain::Name>, unexpected::Error>;
 }
 
-struct ManagerImpl {
+pub struct ManagerImpl {
     origin_store: sync::Arc<dyn origin::store::Store>,
     domain_config_store: sync::Arc<dyn config::Store>,
     domain_checker: checker::DNSChecker,
@@ -170,7 +170,7 @@ pub fn new(
     domain_config_store: sync::Arc<dyn config::Store>,
     domain_checker: checker::DNSChecker,
     acme_manager: Option<sync::Arc<dyn acme::manager::Manager>>,
-) -> sync::Arc<dyn Manager> {
+) -> ManagerImpl {
     let canceller = CancellationToken::new();
 
     let origin_sync_handler = {
@@ -188,22 +188,20 @@ pub fn new(
         })
     };
 
-    sync::Arc::new(ManagerImpl {
+    ManagerImpl {
         origin_store,
         domain_config_store,
         domain_checker,
         acme_manager,
         canceller,
         origin_sync_handler,
-    })
+    }
 }
 
 impl ManagerImpl {
-    pub async fn stop(self) {
+    pub fn stop(self) -> tokio::task::JoinHandle<()> {
         self.canceller.cancel();
         self.origin_sync_handler
-            .await
-            .expect("origin_sync_handler errored");
     }
 }
 
diff --git a/src/main.rs b/src/main.rs
index 4341b3d..1360c7a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -12,6 +12,8 @@ use std::net::SocketAddr;
 use std::str::FromStr;
 use std::{future, path, sync};
 
+use domiply::domain::manager::Manager;
+
 #[derive(Parser, Debug)]
 #[command(version)]
 #[command(about = "A domiply to another dimension")]
@@ -152,6 +154,8 @@ async fn main() {
         https_params.as_ref().map(|p| p.domain_acme_manager.clone()),
     );
 
+    let domain_manager = sync::Arc::new(domain_manager);
+
     let service = domiply::service::new(
         domain_manager.clone(),
         config.domain_checker_target_a,
@@ -317,5 +321,16 @@ async fn main() {
 
     while wait_group.next().await.is_some() {}
 
+    // TODO this is currently required so that we can be sure domain_manager is no longer used by
+    // anything else, and the into_inner below works. It would be great if service could accept a
+    // ref to domain_manager instead, and then maybe this wouldn't be needed?
+    drop(service);
+
+    sync::Arc::into_inner(domain_manager)
+        .unwrap()
+        .stop()
+        .await
+        .expect("domain manager failed to shutdown cleanly");
+
     log::info!("Graceful shutdown complete");
 }