[Bugfix] Propagate lease lost errors from TableSummaryWorker (#38952)

Propagate the lease lost error and properly shutdown when we hit the error in `TableSummaryWorker`. We were seeing instances repeatedly hit this error and not shutdown.

GitOrigin-RevId: 0043d518a03616a7d08d6d072eac358a550692df

Author: jordanhunt22

crates/application/src/lib.rs

@@ -108,6 +108,7 @@ use common::{
         DatabaseSchema,
         TableDefinition,
     },
+    shutdown::ShutdownSignal,
     types::{
         env_var_limit_met,
         env_var_name_not_unique,
@@ -684,6 +685,7 @@ impl<RT: Runtime> Application<RT> {
         cache: QueryCache,
         fetch_client: Arc<dyn FetchClient>,
         local_log_sink: Option<String>,
+        lease_lost_shutdown: ShutdownSignal,
     ) -> anyhow::Result<Self> {
         let module_cache =
             ModuleCache::new(runtime.clone(), application_storage.modules_storage.clone()).await;
@@ -717,8 +719,12 @@ impl<RT: Runtime> Application<RT> {
         let search_worker = Arc::new(Mutex::new(search_worker));
         let search_and_vector_bootstrap_worker =
             Arc::new(Mutex::new(database.start_search_and_vector_bootstrap()));
-        let table_summary_worker =
-            TableSummaryWorker::start(runtime.clone(), database.clone(), persistence.clone());
+        let table_summary_worker = TableSummaryWorker::start(
+            runtime.clone(),
+            database.clone(),
+            persistence.clone(),
+            lease_lost_shutdown,
+        );
         let schema_worker = Arc::new(Mutex::new(runtime.spawn(
             "schema_worker",
             SchemaWorker::start(runtime.clone(), database.clone()),

crates/application/src/table_summary_worker.rs

@@ -4,7 +4,11 @@ use std::{
 };
 
 use common::{
-    errors::report_error,
+    errors::{
+        lease_lost_error,
+        report_error,
+        LeaseLostError,
+    },
     knobs::{
         DATABASE_WORKERS_MAX_CHECKPOINT_AGE,
         DATABASE_WORKERS_MIN_COMMITS,
@@ -17,6 +21,7 @@ use common::{
         SpawnHandle,
         UnixTimestamp,
     },
+    shutdown::ShutdownSignal,
 };
 use database::{
     table_summary::write_snapshot,
@@ -64,6 +69,7 @@ impl<RT: Runtime> TableSummaryWorker<RT> {
         runtime: RT,
         database: Database<RT>,
         persistence: Arc<dyn Persistence>,
+        lease_lost_shutdown: ShutdownSignal,
     ) -> TableSummaryClient {
         let table_summary_worker = Self {
             runtime: runtime.clone(),
@@ -73,7 +79,7 @@ impl<RT: Runtime> TableSummaryWorker<RT> {
         let (cancel_sender, cancel_receiver) = oneshot::channel();
         let handle = runtime.spawn(
             "table_summary_worker",
-            table_summary_worker.go(cancel_receiver),
+            table_summary_worker.go(cancel_receiver, lease_lost_shutdown),
         );
         let inner = Inner {
             handle,
@@ -127,7 +133,7 @@ impl<RT: Runtime> TableSummaryWorker<RT> {
         Ok(())
     }
 
-    async fn go(self, cancel_receiver: oneshot::Receiver<()>) {
+    async fn go(self, cancel_receiver: oneshot::Receiver<()>, lease_lost_shutdown: ShutdownSignal) {
         tracing::info!("Starting background table summary worker");
         let mut timer = Some(table_summary_bootstrap_timer());
         let cancel_fut = cancel_receiver.fuse();
@@ -166,6 +172,11 @@ impl<RT: Runtime> TableSummaryWorker<RT> {
             }
             if let Err(mut err) = result {
                 report_error(&mut err).await;
+                if let Some(LeaseLostError) = err.downcast_ref() {
+                    lease_lost_shutdown.signal(
+                        lease_lost_error().context("Failed to write table summary checkpoint"),
+                    );
+                }
             }
             let wait_fut = self.runtime.wait(Duration::from_secs(10)).fuse();
             pin_mut!(wait_fut);

crates/application/src/test_helpers.rs

@@ -273,6 +273,7 @@ impl<RT: Runtime> ApplicationTestExt<RT> for Application<RT> {
             QueryCache::new(*UDF_CACHE_MAX_SIZE),
             fetch_client,
             None, // local_log_sink
+            ShutdownSignal::panic(),
         )
         .await?;
 

crates/local_backend/src/lib.rs

@@ -149,7 +149,7 @@ pub async fn make_app(
         persistence.clone(),
         runtime.clone(),
         searcher.clone(),
-        preempt_tx,
+        preempt_tx.clone(),
         virtual_system_mapping().clone(),
         Arc::new(NoOpUsageEventLogger),
     )
@@ -231,6 +231,7 @@ pub async fn make_app(
         QueryCache::new(*UDF_CACHE_MAX_SIZE),
         fetch_client,
         config.local_log_sink.clone(),
+        preempt_tx.clone(),
     )
     .await?;