Add graceful shutdown and memory tracing

1 files changed, 9 insertions, 3 deletions

diff --git a/src/cert_store.rs b/src/cert_store.rs
index d561605..c1381db 100644
--- a/src/cert_store.rs
+++ b/src/cert_store.rs
@@ -4,7 +4,7 @@ use std::time::{Duration, Instant};
 
 use anyhow::Result;
 use chrono::Utc;
-use futures::TryFutureExt;
+use futures::{FutureExt, TryFutureExt};
 use log::*;
 use tokio::select;
 use tokio::sync::{mpsc, watch};
@@ -16,7 +16,6 @@ use rustls::sign::CertifiedKey;
 
 use crate::cert::{Cert, CertSer};
 use crate::consul::*;
-use crate::exit_on_err;
 use crate::proxy_config::*;
 
 pub struct CertStore {
@@ -33,6 +32,7 @@ impl CertStore {
 		consul: Consul,
 		rx_proxy_config: watch::Receiver<Arc<ProxyConfig>>,
 		letsencrypt_email: String,
+		exit_on_err: impl Fn(anyhow::Error) + Send + 'static,
 	) -> Arc<Self> {
 		let (tx, rx) = mpsc::unbounded_channel();
 
@@ -45,7 +45,13 @@ impl CertStore {
 			tx_need_cert: tx,
 		});
 
-		tokio::spawn(cert_store.clone().certificate_loop(rx).map_err(exit_on_err));
+		tokio::spawn(
+			cert_store
+				.clone()
+				.certificate_loop(rx)
+				.map_err(exit_on_err)
+				.then(|_| async { info!("Certificate renewal task exited") }),
+		);
 
 		cert_store
 	}