Add graceful shutdown and memory tracing

1 files changed, 18 insertions, 5 deletions

diff --git a/src/proxy_config.rs b/src/proxy_config.rs
index 4add98c..e380885 100644
--- a/src/proxy_config.rs
+++ b/src/proxy_config.rs
@@ -9,7 +9,7 @@ use futures::future::BoxFuture;
 use futures::stream::{FuturesUnordered, StreamExt};
 
 use log::*;
-use tokio::{sync::watch, time::sleep};
+use tokio::{select, sync::watch, time::sleep};
 
 use crate::consul::*;
 
@@ -231,7 +231,10 @@ struct NodeWatchState {
 	retries: u32,
 }
 
-pub fn spawn_proxy_config_task(consul: Consul) -> watch::Receiver<Arc<ProxyConfig>> {
+pub fn spawn_proxy_config_task(
+	consul: Consul,
+	mut must_exit: watch::Receiver<bool>,
+) -> watch::Receiver<Arc<ProxyConfig>> {
 	let (tx, rx) = watch::channel(Arc::new(ProxyConfig {
 		entries: Vec::new(),
 	}));
@@ -244,8 +247,13 @@ pub fn spawn_proxy_config_task(consul: Consul) -> watch::Receiver<Arc<ProxyConfi
 
 		let mut node_site = HashMap::new();
 
-		loop {
-			match consul.list_nodes().await {
+		while !*must_exit.borrow() {
+			let list_nodes = select! {
+				ln = consul.list_nodes() => ln,
+				_ = must_exit.changed() => continue,
+			};
+
+			match list_nodes {
 				Ok(consul_nodes) => {
 					info!("Watched consul nodes: {:?}", consul_nodes);
 					for consul_node in consul_nodes {
@@ -271,7 +279,12 @@ pub fn spawn_proxy_config_task(consul: Consul) -> watch::Receiver<Arc<ProxyConfi
 				}
 			}
 
-			let (node, res): (String, Result<_>) = match watches.next().await {
+			let next_watch = select! {
+				nw = watches.next() => nw,
+				_ = must_exit.changed() => continue,
+			};
+
+			let (node, res): (String, Result<_>) = match next_watch {
 				Some(v) => v,
 				None => {
 					warn!("No nodes currently watched in proxy_config.rs");