refactor(agent): use exit flag everywhere

Author: Reisen

src/agent.rs

@@ -61,14 +61,6 @@ Metrics Server:
 Note that there is an Oracle and Exporter for each network, but only one Local Store and Global Store.
 
 ################################################################################################################################## */
-
-pub mod legacy_schedule;
-pub mod market_schedule;
-pub mod metrics;
-pub mod pythd;
-pub mod solana;
-pub mod state;
-pub mod store;
 use {
     self::{
         config::Config,
@@ -78,11 +70,33 @@ use {
     },
     anyhow::Result,
     futures_util::future::join_all,
+    lazy_static::lazy_static,
     slog::Logger,
     std::sync::Arc,
-    tokio::sync::broadcast,
+    tokio::sync::watch,
 };
 
+pub mod legacy_schedule;
+pub mod market_schedule;
+pub mod metrics;
+pub mod pythd;
+pub mod solana;
+pub mod state;
+pub mod store;
+
+lazy_static! {
+    /// A static exit flag to indicate to running threads that we're shutting down. This is used to
+    /// gracefully shut down the application.
+    ///
+    /// We make this global based on the fact the:
+    /// - The `Sender` side does not rely on any async runtime.
+    /// - Exit logic doesn't really require carefully threading this value through the app.
+    /// - The `Receiver` side of a watch channel performs the detection based on if the change
+    ///   happened after the subscribe, so it means all listeners should always be notified
+    ///   correctly.
+    pub static ref EXIT: watch::Sender<bool> = watch::channel(false).0;
+}
+
 pub struct Agent {
     config: Config,
 }
@@ -109,10 +123,6 @@ impl Agent {
         // job handles
         let mut jhs = vec![];
 
-        // Create the channels
-        // TODO: make all components listen to shutdown signal
-        let (shutdown_tx, _) = broadcast::channel(self.config.channel_capacities.shutdown);
-
         // Create the Pythd Adapter.
         let adapter =
             Arc::new(state::State::new(self.config.pythd_adapter.clone(), logger.clone()).await);
@@ -136,17 +146,13 @@ impl Agent {
         }
 
         // Create the Notifier task for the Pythd RPC.
-        jhs.push(tokio::spawn(notifier(
-            adapter.clone(),
-            shutdown_tx.subscribe(),
-        )));
+        jhs.push(tokio::spawn(notifier(adapter.clone())));
 
         // Spawn the Pythd API Server
         jhs.push(tokio::spawn(rpc::run(
             self.config.pythd_api_server.clone(),
             logger.clone(),
             adapter.clone(),
-            shutdown_tx.subscribe(),
         )));
 
         // Spawn the metrics server

src/agent/metrics.rs

@@ -1,44 +1,22 @@
 use {
-    super::state::{
-        local::PriceInfo,
-        State,
-    },
-    crate::agent::{
-        solana::oracle::PriceEntry,
-        store::PriceIdentifier,
-    },
+    super::state::{local::PriceInfo, State},
+    crate::agent::{solana::oracle::PriceEntry, store::PriceIdentifier},
     lazy_static::lazy_static,
     prometheus_client::{
-        encoding::{
-            text::encode,
-            EncodeLabelSet,
-        },
-        metrics::{
-            counter::Counter,
-            family::Family,
-            gauge::Gauge,
-        },
+        encoding::{text::encode, EncodeLabelSet},
+        metrics::{counter::Counter, family::Family, gauge::Gauge},
         registry::Registry,
     },
     serde::Deserialize,
     slog::Logger,
     solana_sdk::pubkey::Pubkey,
     std::{
         net::SocketAddr,
-        sync::{
-            atomic::AtomicU64,
-            Arc,
-        },
+        sync::{atomic::AtomicU64, Arc},
         time::Instant,
     },
     tokio::sync::Mutex,
-    warp::{
-        hyper::StatusCode,
-        reply,
-        Filter,
-        Rejection,
-        Reply,
-    },
+    warp::{hyper::StatusCode, reply, Filter, Rejection, Reply},
 };
 
 pub fn default_bind_address() -> SocketAddr {
@@ -68,8 +46,8 @@ lazy_static! {
 /// metrics.
 pub struct MetricsServer {
     pub start_time: Instant,
-    pub logger:     Logger,
-    pub adapter:    Arc<State>,
+    pub logger: Logger,
+    pub adapter: Arc<State>,
 }
 
 impl MetricsServer {
@@ -105,7 +83,11 @@ impl MetricsServer {
                 }
             });
 
-        warp::serve(metrics_route).bind(addr).await;
+        let (_, serve) = warp::serve(metrics_route).bind_with_graceful_shutdown(addr, async {
+            let _ = crate::agent::EXIT.subscribe().changed().await;
+        });
+
+        serve.await
     }
 }
 
@@ -169,12 +151,12 @@ pub struct PriceGlobalMetrics {
 
     /// f64 is used to get u64 support. Official docs:
     /// https://docs.rs/prometheus-client/latest/prometheus_client/metrics/gauge/struct.Gauge.html#using-atomicu64-as-storage-and-f64-on-the-interface
-    conf:      Family<PriceGlobalLabels, Gauge<f64, AtomicU64>>,
+    conf: Family<PriceGlobalLabels, Gauge<f64, AtomicU64>>,
     timestamp: Family<PriceGlobalLabels, Gauge>,
 
     /// Note: the exponent is not applied to this metric
-    prev_price:     Family<PriceGlobalLabels, Gauge>,
-    prev_conf:      Family<PriceGlobalLabels, Gauge<f64, AtomicU64>>,
+    prev_price: Family<PriceGlobalLabels, Gauge>,
+    prev_conf: Family<PriceGlobalLabels, Gauge<f64, AtomicU64>>,
     prev_timestamp: Family<PriceGlobalLabels, Gauge>,
 
     /// How many times this Price was updated in the global store
@@ -317,10 +299,10 @@ pub struct PriceLocalLabels {
 /// Metrics exposed to Prometheus by the local store for each price
 #[derive(Default)]
 pub struct PriceLocalMetrics {
-    price:     Family<PriceLocalLabels, Gauge>,
+    price: Family<PriceLocalLabels, Gauge>,
     /// f64 is used to get u64 support. Official docs:
     /// https://docs.rs/prometheus-client/latest/prometheus_client/metrics/gauge/struct.Gauge.html#using-atomicu64-as-storage-and-f64-on-the-interface
-    conf:      Family<PriceLocalLabels, Gauge<f64, AtomicU64>>,
+    conf: Family<PriceLocalLabels, Gauge<f64, AtomicU64>>,
     timestamp: Family<PriceLocalLabels, Gauge>,
 
     /// How many times this price was updated in the local store

src/agent/pythd/api/rpc.rs

@@ -50,10 +50,7 @@ use {
         net::SocketAddr,
         sync::Arc,
     },
-    tokio::sync::{
-        broadcast,
-        mpsc,
-    },
+    tokio::sync::mpsc,
     warp::{
         ws::{
             Message,
@@ -430,29 +427,20 @@ impl Default for Config {
     }
 }
 
-pub async fn run<S>(
-    config: Config,
-    logger: Logger,
-    adapter: Arc<S>,
-    shutdown_rx: broadcast::Receiver<()>,
-) where
+pub async fn run<S>(config: Config, logger: Logger, adapter: Arc<S>)
+where
     S: state::StateApi,
     S: Send,
     S: Sync,
     S: 'static,
 {
-    if let Err(err) = serve(config, &logger, adapter, shutdown_rx).await {
+    if let Err(err) = serve(config, &logger, adapter).await {
         error!(logger, "{}", err);
         debug!(logger, "error context"; "context" => format!("{:?}", err));
     }
 }
 
-async fn serve<S>(
-    config: Config,
-    logger: &Logger,
-    adapter: Arc<S>,
-    mut shutdown_rx: broadcast::Receiver<()>,
-) -> Result<()>
+async fn serve<S>(config: Config, logger: &Logger, adapter: Arc<S>) -> Result<()>
 where
     S: state::StateApi,
     S: Send,
@@ -490,8 +478,8 @@ where
 
     let (_, serve) = warp::serve(index).bind_with_graceful_shutdown(
         config.listen_address.as_str().parse::<SocketAddr>()?,
-        async move {
-            let _ = shutdown_rx.recv().await;
+        async {
+            let _ = crate::agent::EXIT.subscribe().changed().await;
         },
     );
 

src/agent/state.rs

@@ -185,7 +185,7 @@ mod tests {
         let (shutdown_tx, _) = broadcast::channel(1);
 
         // Spawn Price Notifier
-        let jh = tokio::spawn(notifier(adapter.clone(), shutdown_tx.subscribe()));
+        let jh = tokio::spawn(notifier(adapter.clone()));
 
         TestAdapter {
             adapter,

src/agent/state/api.rs

@@ -45,10 +45,7 @@ use {
         PriceStatus,
     },
     std::sync::Arc,
-    tokio::sync::{
-        broadcast,
-        mpsc,
-    },
+    tokio::sync::mpsc,
 };
 
 // TODO: implement Display on PriceStatus and then just call PriceStatus::to_string
@@ -172,12 +169,13 @@ pub trait StateApi {
     ) -> Result<()>;
 }
 
-pub async fn notifier(adapter: Arc<State>, mut shutdown_rx: broadcast::Receiver<()>) {
+pub async fn notifier(adapter: Arc<State>) {
     let mut interval = tokio::time::interval(adapter.notify_price_sched_interval_duration);
+    let mut exit = crate::agent::EXIT.subscribe();
     loop {
         adapter.drop_closed_subscriptions().await;
         tokio::select! {
-            _ = shutdown_rx.recv() => {
+            _ = exit.changed() => {
                 info!(adapter.logger, "shutdown signal received");
                 return;
             }

src/agent/state/keypairs.rs

@@ -5,34 +5,17 @@
 use {
     super::State,
     crate::agent::solana::network::Network,
-    anyhow::{
-        Context,
-        Result,
-    },
+    anyhow::{Context, Result},
     serde::Deserialize,
     slog::Logger,
     solana_client::nonblocking::rpc_client::RpcClient,
-    solana_sdk::{
-        commitment_config::CommitmentConfig,
-        signature::Keypair,
-        signer::Signer,
-    },
-    std::{
-        net::SocketAddr,
-        sync::Arc,
-    },
-    tokio::{
-        sync::RwLock,
-        task::JoinHandle,
-    },
+    solana_sdk::{commitment_config::CommitmentConfig, signature::Keypair, signer::Signer},
+    std::{net::SocketAddr, sync::Arc},
+    tokio::{sync::RwLock, task::JoinHandle},
     warp::{
         hyper::StatusCode,
-        reply::{
-            self,
-            WithStatus,
-        },
-        Filter,
-        Rejection,
+        reply::{self, WithStatus},
+        Filter, Rejection,
     },
 };
 
@@ -49,24 +32,24 @@ pub fn default_bind_address() -> SocketAddr {
 #[derive(Clone, Debug, Deserialize)]
 #[serde(default)]
 pub struct Config {
-    primary_min_keypair_balance_sol:   u64,
+    primary_min_keypair_balance_sol: u64,
     secondary_min_keypair_balance_sol: u64,
-    bind_address:                      SocketAddr,
+    bind_address: SocketAddr,
 }
 
 impl Default for Config {
     fn default() -> Self {
         Self {
-            primary_min_keypair_balance_sol:   default_min_keypair_balance_sol(),
+            primary_min_keypair_balance_sol: default_min_keypair_balance_sol(),
             secondary_min_keypair_balance_sol: default_min_keypair_balance_sol(),
-            bind_address:                      default_bind_address(),
+            bind_address: default_bind_address(),
         }
     }
 }
 
 #[derive(Default)]
 pub struct KeypairState {
-    primary_current_keypair:   RwLock<Option<Keypair>>,
+    primary_current_keypair: RwLock<Option<Keypair>>,
     secondary_current_keypair: RwLock<Option<Keypair>>,
 }
 
@@ -193,9 +176,13 @@ where
             }
         });
 
-    let http_api_jh = tokio::spawn(
-        warp::serve(primary_upload_route.or(secondary_upload_route)).bind(config.bind_address),
-    );
+    let http_api_jh = {
+        let (_, serve) = warp::serve(primary_upload_route.or(secondary_upload_route))
+            .bind_with_graceful_shutdown(config.bind_address, async {
+                let _ = crate::agent::EXIT.subscribe().changed().await;
+            });
+        tokio::spawn(serve)
+    };
 
     // WARNING: All jobs spawned here must report their join handles in this vec
     vec![http_api_jh]