Graceful shutdown

Cargo.toml

@@ -4,7 +4,8 @@ members = [
   "api",
   "ingestion",
   "integration_tests",
-  "rove_connector", "util",
+  "rove_connector",
+  "util",
 ]
 resolver = "2"
 
@@ -32,4 +33,5 @@ serde = { version = "1.0.217", features = ["derive"] }
 thiserror = "1.0.69"
 tokio = { version = "1.43.0", features = ["rt-multi-thread", "macros"] }
 tokio-postgres = { version = "0.7.12", features = ["with-chrono-0_4"] }
+tokio-util = { version = "0.7.13", features = ["rt"] }
 toml = "0.8.19"

api/Cargo.toml

@@ -8,7 +8,9 @@ axum.workspace = true
 bb8.workspace = true
 bb8-postgres.workspace = true
 chrono.workspace = true
+util = { path = "../util" }
 postgres-types.workspace = true
 serde.workspace = true
 tokio.workspace = true
 tokio-postgres.workspace = true
+tokio-util.workspace = true

api/src/lib.rs

@@ -13,6 +13,7 @@ use timeseries::{
 };
 use timeslice::{get_timeslice, Timeslice};
 use tokio_postgres::NoTls;
+use tokio_util::sync::CancellationToken;
 
 pub mod latest;
 pub mod timeseries;
@@ -119,7 +120,7 @@ async fn latest_handler(
     Ok(Json(LatestResp { data }))
 }
 
-pub async fn run(pool: PgConnectionPool) {
+pub async fn run(pool: PgConnectionPool, cancel_token: CancellationToken) {
     // build our application with routes
     let app = Router::new()
         .route(
@@ -135,5 +136,8 @@ pub async fn run(pool: PgConnectionPool) {
 
     // run it with hyper on localhost:3000
     let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
-    axum::serve(listener, app).await.unwrap();
+    axum::serve(listener, app)
+        .with_graceful_shutdown(::util::await_cancellation(cancel_token))
+        .await
+        .unwrap();
 }

api/src/main.rs

@@ -1,5 +1,6 @@
 use bb8_postgres::PostgresConnectionManager;
 use tokio_postgres::NoTls;
+use tokio_util::sync::CancellationToken;
 
 #[tokio::main]
 async fn main() {
@@ -19,5 +20,9 @@ async fn main() {
     let manager = PostgresConnectionManager::new_from_stringlike(connect_string, NoTls).unwrap();
     let pool = bb8::Pool::builder().build(manager).await.unwrap();
 
-    lard_api::run(pool).await;
+    // set up cancellation token and signal catcher for graceful shutdown
+    let cancel_token = CancellationToken::new();
+    tokio::spawn(util::signal_catcher(cancel_token.clone()));
+
+    tokio::spawn(lard_api::run(pool, cancel_token.clone()));
 }

ingestion/Cargo.toml

@@ -16,6 +16,7 @@ bb8-postgres.workspace = true
 bytes.workspace = true
 chrono.workspace = true
 chronoutil.workspace = true
+util = { path = "../util" }
 csv.workspace = true
 futures.workspace = true
 kafka.workspace = true
@@ -27,4 +28,5 @@ serde.workspace = true
 thiserror.workspace = true
 tokio.workspace = true
 tokio-postgres.workspace = true
+tokio-util.workspace = true
 toml.workspace = true

ingestion/src/kvkafka.rs

@@ -3,6 +3,7 @@ use kafka::consumer::{Consumer, FetchOffset, GroupOffsetStorage};
 use serde::{Deserialize, Deserializer};
 use thiserror::Error;
 use tokio::sync::mpsc;
+use tokio_util::sync::CancellationToken;
 
 use crate::PgConnectionPool;
 
@@ -127,11 +128,15 @@ pub struct Msg {
     kvdata: Kvdata,
 }
 
-pub async fn read_and_insert(pool: PgConnectionPool, group_string: String) {
+pub async fn read_and_insert(
+    pool: PgConnectionPool,
+    group_string: String,
+    cancel_token: CancellationToken,
+) {
     let (tx, mut rx) = mpsc::channel(10);
 
     tokio::spawn(async move {
-        read_kafka(group_string, tx).await;
+        read_kafka(group_string, tx, cancel_token).await;
     });
 
     let mut client = pool.get().await.expect("couldn't connect to database");
@@ -212,7 +217,7 @@ pub async fn parse_message(message: &[u8], tx: &mpsc::Sender<Msg>) -> Result<(),
     Ok(())
 }
 
-async fn read_kafka(group_name: String, tx: mpsc::Sender<Msg>) {
+async fn read_kafka(group_name: String, tx: mpsc::Sender<Msg>, cancel_token: CancellationToken) {
     // NOTE: reading from the 4 redundant kafka queues, but only reading the checked data (other topics exists)
     let mut consumer = Consumer::from_hosts(vec![
         "kafka2-a1.met.no:9092".to_owned(),
@@ -229,27 +234,35 @@ async fn read_kafka(group_name: String, tx: mpsc::Sender<Msg>) {
 
     // Consume the kafka queue infinitely
     loop {
-        // https://docs.rs/kafka/latest/src/kafka/consumer/mod.rs.html#155
-        // poll asks for next available chunk of data as a MessageSet
-        match consumer.poll() {
-            Ok(sets) => {
-                for msgset in sets.iter() {
-                    for msg in msgset.messages() {
-                        if let Err(e) = parse_message(msg.value, &tx).await {
-                            eprintln!("{}", e);
+        tokio::select! {
+            _ = cancel_token.cancelled() => {
+                eprintln!("cancellation token triggered");
+                break;
+            }
+            // https://docs.rs/kafka/latest/src/kafka/consumer/mod.rs.html#155
+            // poll asks for next available chunk of data as a MessageSet
+            poll_result = async { consumer.poll() } => {
+                match poll_result {
+                    Ok(sets) => {
+                        for msgset in sets.iter() {
+                            for msg in msgset.messages() {
+                                if let Err(e) = parse_message(msg.value, &tx).await {
+                                    eprintln!("{}", e);
+                                }
+                            }
+                            if let Err(e) = consumer.consume_messageset(msgset) {
+                                eprintln!("{}", e);
+                            }
                         }
+                        consumer
+                            .commit_consumed()
+                            .expect("could not commit offset in consumer"); // ensure we keep offset
                     }
-                    if let Err(e) = consumer.consume_messageset(msgset) {
-                        eprintln!("{}", e);
+                    Err(e) => {
+                        eprintln!("{}\nRetrying in 5 seconds...", Error::Kafka(e));
+                        tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
                     }
                 }
-                consumer
-                    .commit_consumed()
-                    .expect("could not commit offset in consumer"); // ensure we keep offset
-            }
-            Err(e) => {
-                eprintln!("{}\nRetrying in 5 seconds...", Error::Kafka(e));
-                tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
             }
         }
     }

ingestion/src/lib.rs

@@ -17,6 +17,7 @@ use std::{
 };
 use thiserror::Error;
 use tokio_postgres::NoTls;
+use tokio_util::sync::CancellationToken;
 
 #[cfg(feature = "kafka")]
 pub mod kvkafka;
@@ -430,6 +431,7 @@ pub async fn run(
     permit_tables: Arc<RwLock<(ParamPermitTable, StationPermitTable)>>,
     rove_connector: rove_connector::Connector,
     qc_pipelines: HashMap<(i32, RelativeDuration), rove::Pipeline>,
+    cancel_token: CancellationToken,
 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     // set up param conversion map
     let param_conversions = get_conversions(param_conversion_path)?;
@@ -452,7 +454,9 @@ pub async fn run(
 
     // run our app with hyper, listening globally on port 3001
     let listener = tokio::net::TcpListener::bind("0.0.0.0:3001").await?;
-    axum::serve(listener, app).await?;
+    axum::serve(listener, app)
+        .with_graceful_shutdown(util::await_cancellation(cancel_token))
+        .await?;
 
     Ok(())
 }

ingestion/src/main.rs

@@ -3,6 +3,7 @@ use lard_ingestion::qc_pipelines::load_pipelines;
 use rove_connector::Connector;
 use std::sync::{Arc, RwLock};
 use tokio_postgres::NoTls;
+use tokio_util::sync::CancellationToken;
 
 use lard_ingestion::{getenv, permissions};
 
@@ -47,7 +48,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
 
     let qc_pipelines = load_pipelines("qc_pipelines/fresh")?;
 
-    println!("Spawing task to fetch permissions from StInfoSys...");
+    println!("Spawning task to fetch permissions from StInfoSys...");
     // background task to refresh permit tables every 30 mins
     tokio::task::spawn(async move {
         let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(30 * 60));
@@ -66,25 +67,41 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
         }
     });
 
-    // Spawn kvkafka reader
-    #[cfg(feature = "kafka_prod")]
-    {
-        let kafka_group = args[1].to_string();
-        println!("Spawing kvkafka reader...");
-        tokio::spawn(lard_ingestion::kvkafka::read_and_insert(
-            db_pool.clone(),
-            kafka_group,
-        ));
-    }
+    // set up cancellation token and signal catcher for graceful shutdown
+    let cancel_token = CancellationToken::new();
+    tokio::spawn(util::signal_catcher(cancel_token.clone()));
 
     // Set up and run our server + database
     println!("Ingestion server started!");
-    lard_ingestion::run(
+    let ingestor = tokio::spawn(lard_ingestion::run(
         db_pool,
         PARAMCONV,
         permit_tables,
         rove_connector,
         qc_pipelines,
-    )
-    .await
+        cancel_token,
+    ));
+
+    #[cfg(feature = "kafka_prod")]
+    // Spawn kvkafka reader
+    {
+        let kafka_group = args[1].to_string();
+        println!("Spawning kvkafka reader...");
+        let kvkafka_reader = tokio::spawn(lard_ingestion::kvkafka::read_and_insert(
+            db_pool.clone(),
+            kafka_group,
+            cancel_token.clone(),
+        ));
+
+        let (ingestor_res, kvkafka_reader_res) = tokio::join!(ingestor, kvkafka_reader);
+        (_, _) = (ingestor_res, kvkafka_reader_res); // ignore for now
+    }
+
+    #[cfg(not(feature = "kafka_prod"))]
+    {
+        let ingestor_res = tokio::join!(ingestor);
+        _ = ingestor_res; // ignore for now
+    }
+
+    Ok(())
 }

integration_tests/Cargo.toml

@@ -11,10 +11,12 @@ lard_api = { path = "../api" }
 lard_ingestion = { path = "../ingestion", features = ["integration_tests", "kafka"] }
 chrono.workspace = true
 chronoutil.workspace = true
+util = { path = "../util" }
 tokio.workspace = true
 tokio-postgres.workspace = true
+tokio-util.workspace = true
 bb8.workspace = true
-bb8-postgres.workspace = true 
+bb8-postgres.workspace = true
 rove.workspace = true
 rove_connector = { path = "../rove_connector" }
 serde.workspace = true

integration_tests/tests/end_to_end.rs

@@ -12,6 +12,7 @@ use futures::{Future, FutureExt};
 use rove::data_switch::{DataConnector, SpaceSpec, TimeSpec, Timestamp};
 use tokio::sync::mpsc;
 use tokio_postgres::NoTls;
+use tokio_util::sync::CancellationToken;
 
 use lard_api::{timeseries::Timeseries, LatestResp, TimeseriesResp, TimesliceResp};
 use lard_ingestion::{
@@ -251,15 +252,21 @@ async fn e2e_test_wrapper<T: Future<Output = ()>>(test: T) {
     };
     let qc_pipelines = load_pipelines("mock_qc_pipelines/fresh").expect("failed to load pipelines");
 
+    // set up cancellation token and signal catcher to detect premature shutdown
+    let cancel_token = CancellationToken::new();
+    let sig_catcher = tokio::spawn(util::signal_catcher(cancel_token.clone()));
+
+    let cancel_token2 = cancel_token.clone();
     let api_server = tokio::spawn(async move {
         tokio::select! {
-            output = lard_api::run(api_pool) => output,
+            output = lard_api::run(api_pool, cancel_token2) => output,
             _ = init_shutdown_rx1.recv() => {
-                api_shutdown_tx.send(()).unwrap();
-                ()
+                api_shutdown_tx.send(()).unwrap()
             },
         }
     });
+
+    let cancel_token2 = cancel_token.clone();
     let ingestor = tokio::spawn(async move {
         tokio::select! {
             output = lard_ingestion::run(
@@ -268,6 +275,7 @@ async fn e2e_test_wrapper<T: Future<Output = ()>>(test: T) {
                 mock_permit_tables(),
                 rove_connector,
                 qc_pipelines,
+                cancel_token2,
             ) => output,
             _ = init_shutdown_rx2.recv() => {
                 ingestor_shutdown_tx.send(()).unwrap();
@@ -279,6 +287,7 @@ async fn e2e_test_wrapper<T: Future<Output = ()>>(test: T) {
     tokio::select! {
         _ = api_server => panic!("API server task terminated first"),
         _ = ingestor => panic!("Ingestor server task terminated first"),
+        _ = sig_catcher => panic!("Signal catcher caught a shutdown signal"),
         // Clean up database even if test panics, to avoid test poisoning
         test_result = AssertUnwindSafe(test).catch_unwind() => {
             // For debugging a specific test, it might be useful to skip the cleanup process