Use TaskSet to cleanly shut down the http service

TODO

@@ -1 +1,3 @@
-- clean up main a lot
+- make domain_manager implement rusttls cert resolver
+- Try to switch from Arc to Box where possible
+- maybe build TaskSet into some kind of defer-like replacement

src/main.rs

@@ -148,7 +148,7 @@ async fn main() {
     let domain_manager = sync::Arc::new(domain_manager);
 
     {
-        let http_service = domiply::service::http::new(
+        let (http_service, http_service_task_set) = domiply::service::http::new(
             domain_manager.clone(),
             config.domain_checker_target_a,
             config.passphrase,
@@ -162,7 +162,7 @@ async fn main() {
 
         canceller.cancelled().await;
 
-        sync::Arc::into_inner(http_service).unwrap().stop().await;
+        domiply::service::http::stop(http_service, http_service_task_set).await;
     }
 
     sync::Arc::into_inner(domain_manager)

src/service/http.rs

@@ -1,16 +1,15 @@
 mod tasks;
 mod tpl;
 
-use futures::stream::futures_unordered::FuturesUnordered;
 use hyper::{Body, Method, Request, Response};
 use serde::{Deserialize, Serialize};
-use tokio_util::sync::CancellationToken;
 
 use std::convert::Infallible;
 use std::str::FromStr;
 use std::{future, net, sync};
 
-use crate::{domain, origin, service};
+use crate::error::unexpected;
+use crate::{domain, origin, service, util};
 
 type SvcResponse = Result<Response<hyper::body::Body>, String>;
 
@@ -20,9 +19,6 @@ pub struct Service {
     passphrase: String,
     http_domain: domain::Name,
     handlebars: handlebars::Handlebars<'static>,
-
-    canceller: CancellationToken,
-    wait_group: FuturesUnordered<tokio::task::JoinHandle<()>>,
 }
 
 pub struct HTTPSParams {
@@ -37,50 +33,52 @@ pub fn new(
     http_listen_addr: net::SocketAddr,
     http_domain: domain::Name,
     https_params: Option<HTTPSParams>,
-) -> sync::Arc<Service> {
+) -> (sync::Arc<Service>, util::TaskSet<unexpected::Error>) {
     let service = sync::Arc::new(Service {
         domain_manager: domain_manager.clone(),
         target_a,
         passphrase,
         http_domain: http_domain.clone(),
         handlebars: tpl::get(),
-        canceller: CancellationToken::new(),
-        wait_group: FuturesUnordered::new(),
     });
 
-    service.wait_group.push(tasks::listen_http(
+    let task_set = util::TaskSet::new();
+
+    task_set.spawn(|canceller| {
+        tasks::listen_http(
             service.clone(),
-        service.canceller.clone(),
+            canceller,
             http_listen_addr,
             http_domain.clone(),
-    ));
+        )
+    });
 
     if let Some(https_params) = https_params {
-        service.wait_group.push(tasks::listen_https(
+        task_set.spawn(|canceller| {
+            tasks::listen_https(
                 service.clone(),
-            service.canceller.clone(),
+                canceller,
-            https_params.cert_resolver,
+                https_params.cert_resolver.clone(),
                 https_params.listen_addr,
                 http_domain.clone(),
-        ));
+            )
+        });
 
-        service.wait_group.push(tasks::cert_refresher(
+        task_set.spawn(|canceller| {
-            domain_manager,
+            tasks::cert_refresher(domain_manager.clone(), canceller, http_domain.clone())
-            service.canceller.clone(),
+        });
-            http_domain,
-        ))
     }
 
-    return service;
+    return (service, task_set);
 }
 
-impl Service {
+pub async fn stop(service: sync::Arc<Service>, task_set: util::TaskSet<unexpected::Error>) {
-    pub async fn stop(self) {
+    task_set
-        self.canceller.cancel();
+        .stop()
-        for f in self.wait_group {
+        .await
-            f.await.expect("task failed");
+        .iter()
-        }
+        .for_each(|e| log::error!("error while shutting down http service: {e}"));
-    }
+    sync::Arc::into_inner(service).expect("service didn't get cleaned up");
 }
 
 #[derive(Serialize)]

src/service/http/tasks.rs

@@ -1,3 +1,4 @@
+use crate::error::unexpected::{self, Mappable};
 use crate::{domain, service};
 
 use std::{convert, future, net, sync};
@@ -5,13 +6,12 @@ use std::{convert, future, net, sync};
 use futures::StreamExt;
 use tokio_util::sync::CancellationToken;
 
-pub fn listen_http(
+pub async fn listen_http(
     service: sync::Arc<service::http::Service>,
     canceller: CancellationToken,
     addr: net::SocketAddr,
     domain: domain::Name,
-) -> tokio::task::JoinHandle<()> {
+) -> Result<(), unexpected::Error> {
-    tokio::spawn(async move {
     let make_service = hyper::service::make_service_fn(move |_| {
         let service = service.clone();
 
@@ -32,20 +32,16 @@ pub fn listen_http(
         canceller.cancelled().await;
     });
 
-        if let Err(e) = graceful.await {
+    graceful.await.or_unexpected()
-            panic!("server error: {}", e);
-        };
-    })
 }
 
-pub fn listen_https(
+pub async fn listen_https(
     service: sync::Arc<service::http::Service>,
     canceller: CancellationToken,
     cert_resolver: sync::Arc<dyn rustls::server::ResolvesServerCert>,
     addr: net::SocketAddr,
     domain: domain::Name,
-) -> tokio::task::JoinHandle<()> {
+) -> Result<(), unexpected::Error> {
-    tokio::spawn(async move {
     let make_service = hyper::service::make_service_fn(move |_| {
         let service = service.clone();
 
@@ -70,8 +66,7 @@ pub fn listen_https(
     let addr_incoming = hyper::server::conn::AddrIncoming::bind(&addr)
         .expect("https listen socket creation failed");
 
-        let incoming =
+    let incoming = tls_listener::TlsListener::new(server_config, addr_incoming).filter(|conn| {
-            tls_listener::TlsListener::new(server_config, addr_incoming).filter(|conn| {
         if let Err(err) = conn {
             log::error!("Error accepting TLS connection: {:?}", err);
             future::ready(false)
@@ -90,18 +85,14 @@ pub fn listen_https(
         canceller.cancelled().await;
     });
 
-        if let Err(e) = graceful.await {
+    graceful.await.or_unexpected()
-            panic!("server error: {}", e);
-        };
-    })
 }
 
-pub fn cert_refresher(
+pub async fn cert_refresher(
     domain_manager: sync::Arc<dyn domain::manager::Manager>,
     canceller: CancellationToken,
     http_domain: domain::Name,
-) -> tokio::task::JoinHandle<()> {
+) -> Result<(), unexpected::Error> {
-    tokio::spawn(async move {
     use tokio::time;
 
     let mut interval = time::interval(time::Duration::from_secs(60 * 60));
@@ -109,7 +100,7 @@ pub fn cert_refresher(
     loop {
         tokio::select! {
         _ = interval.tick() => (),
-            _ = canceller.cancelled() => return,
+        _ = canceller.cancelled() => return Ok(()),
         }
 
         _ = domain_manager
@@ -138,5 +129,4 @@ pub fn cert_refresher(
                 });
         }
     }
-    })
 }

src/util.rs

@@ -1,4 +1,7 @@
-use std::{fs, io, path};
+use std::{error, fs, io, path};
+
+use futures::stream::futures_unordered::FuturesUnordered;
+use tokio_util::sync::CancellationToken;
 
 pub fn open_file(path: &path::Path) -> io::Result<Option<fs::File>> {
     match fs::File::open(path) {
@@ -9,3 +12,45 @@ pub fn open_file(path: &path::Path) -> io::Result<Option<fs::File>> {
         },
     }
 }
+
+pub struct TaskSet<E>
+where
+    E: error::Error + Send + 'static,
+{
+    canceller: CancellationToken,
+    wait_group: FuturesUnordered<tokio::task::JoinHandle<Result<(), E>>>,
+}
+
+impl<E> TaskSet<E>
+where
+    E: error::Error + Send + 'static,
+{
+    pub fn new() -> TaskSet<E> {
+        TaskSet {
+            canceller: CancellationToken::new(),
+            wait_group: FuturesUnordered::new(),
+        }
+    }
+
+    pub fn spawn<F, Fut>(&self, mut f: F)
+    where
+        Fut: futures::Future<Output = Result<(), E>> + Send + 'static,
+        F: FnMut(CancellationToken) -> Fut,
+    {
+        let canceller = self.canceller.clone();
+        let handle = tokio::spawn(f(canceller));
+        self.wait_group.push(handle);
+    }
+
+    pub async fn stop(self) -> Vec<E> {
+        self.canceller.cancel();
+
+        let mut res = Vec::new();
+        for f in self.wait_group {
+            if let Err(err) = f.await.expect("task failed") {
+                res.push(err);
+            }
+        }
+        res
+    }
+}