Clean shutdown the event logger for usage tracking (#24947)

GitOrigin-RevId: a86742f2a39d415bf12d08c35cafa093e14308ac

Author: nipunn1313

Cargo.lock

@@ -6,6 +6,7 @@ edition = "2021"
 license = "LicenseRef-FSL-1.1-Apache-2.0"
 
 [dependencies]
+anyhow = { workspace = true }
 async-trait = { workspace = true }
 proptest = { workspace = true, optional = true }
 proptest-derive = { workspace = true, optional = true }

crates/events/Cargo.toml

@@ -134,6 +134,9 @@ pub trait UsageEventLogger: Send + Sync + std::fmt::Debug {
 
     /// Dump events into a buffer, waiting for the buffer to empty if it's full.
     async fn record_async(&self, events: Vec<UsageEvent>);
+
+    /// Cleanly shutdown, flushing events
+    async fn shutdown(&self) -> anyhow::Result<()>;
 }
 
 #[derive(Debug, Clone)]
@@ -144,4 +147,8 @@ impl UsageEventLogger for NoOpUsageEventLogger {
     fn record(&self, _events: Vec<UsageEvent>) {}
 
     async fn record_async(&self, _events: Vec<UsageEvent>) {}
+
+    async fn shutdown(&self) -> anyhow::Result<()> {
+        Ok(())
+    }
 }

crates/events/src/usage.rs

@@ -102,8 +102,7 @@ pub struct LocalAppState {
     pub application: Application<ProdRuntime>,
     // Number of sync protocol workers.
     pub live_ws_count: Arc<AtomicU64>,
-    pub shutdown_rx: async_broadcast::Receiver<()>,
-    pub shutdown_tx: ShutdownSignal,
+    pub zombify_rx: async_broadcast::Receiver<()>,
 }
 
 impl LocalAppState {
@@ -122,8 +121,7 @@ impl Clone for LocalAppState {
             instance_name: self.instance_name.clone(),
             application: self.application.clone(),
             live_ws_count: self.live_ws_count.clone(),
-            shutdown_rx: self.shutdown_rx.clone(),
-            shutdown_tx: self.shutdown_tx.clone(),
+            zombify_rx: self.zombify_rx.clone(),
         }
     }
 }
@@ -135,16 +133,16 @@ pub async fn make_app(
     runtime: ProdRuntime,
     config: LocalConfig,
     persistence: Arc<dyn Persistence>,
-    shutdown_rx: async_broadcast::Receiver<()>,
-    shutdown_tx: ShutdownSignal,
+    zombify_rx: async_broadcast::Receiver<()>,
+    preempt_tx: ShutdownSignal,
 ) -> anyhow::Result<LocalAppState> {
     let key_broker = config.key_broker()?;
     let searcher: Arc<dyn Searcher> = Arc::new(InProcessSearcher::new(runtime.clone()).await?);
     let database = Database::load(
         persistence.clone(),
         runtime.clone(),
         searcher.clone(),
-        shutdown_tx.clone(),
+        preempt_tx,
         virtual_system_mapping(),
         Arc::new(NoOpUsageEventLogger),
     )
@@ -257,8 +255,7 @@ pub async fn make_app(
         instance_name,
         application,
         live_ws_count: Arc::new(AtomicU64::new(0)),
-        shutdown_rx,
-        shutdown_tx,
+        zombify_rx,
     };
 
     Ok(app_state)

crates/local_backend/src/lib.rs

@@ -82,7 +82,7 @@ pub async fn stream_udf_execution(
     let entries_future = st
         .application
         .stream_udf_execution(identity, query_args.cursor);
-    let mut shutdown_rx = st.shutdown_rx.clone();
+    let mut zombify_rx = st.zombify_rx.clone();
     futures::select_biased! {
         entries_future_r = entries_future.fuse() => {
             let (log_entries, new_cursor) = entries_future_r?;
@@ -103,7 +103,7 @@ pub async fn stream_udf_execution(
             };
             Ok(Json(response))
         },
-        _ = shutdown_rx.recv().fuse() => {
+        _ = zombify_rx.recv().fuse() => {
             // Return an error so the client reconnects after we come back up.
             Err(anyhow::anyhow!(ErrorMetadata::operational_internal_server_error()).context("Shutting down long poll request").into())
         },
@@ -139,7 +139,7 @@ pub async fn stream_function_logs(
     let entries_future = st
         .application
         .stream_function_logs(identity, query_args.cursor);
-    let mut shutdown_rx = st.shutdown_rx.clone();
+    let mut zombify_rx = st.zombify_rx.clone();
     let request_id = match (query_args.session_id, query_args.client_request_counter) {
         (Some(session_id), Some(client_request_counter)) => Some(RequestId::new_for_ws_session(
             session_id.parse().context("Invalid session ID")?,
@@ -218,7 +218,7 @@ pub async fn stream_function_logs(
             };
             Ok(Json(response))
         },
-        _ = shutdown_rx.recv().fuse() => {
+        _ = zombify_rx.recv().fuse() => {
             // Return an error so the client reconnects after we come back up.
             Err(anyhow::anyhow!(ErrorMetadata::operational_internal_server_error()).context("Shutting down long poll request").into())
         },

crates/local_backend/src/logs.rs

@@ -17,6 +17,7 @@ parking_lot = { workspace = true, features = ["hardware-lock-elision"] }
 pb = { path = "../pb" }
 proptest = { workspace = true, optional = true }
 proptest-derive = { workspace = true, optional = true }
+tracing = { workspace = true }
 value = { path = "../value" }
 
 [dev-dependencies]