Make sure that index backfill errors get reported and retry with backoff (#40638)

1. Ensure that errors from index backfill tasks propagate and get handled;
2. get rid of `IndexBackfill` structure to avoid confusion, by splitting up `pending` into individual items;
3. replace the semaphore with a simpler check on the size of the IndexWorker's `in_progress` JoinMap.

GitOrigin-RevId: ebca825c78317899f2bd9179d40028ef87233f78

Author: goffrie

crates/database/src/database_index_workers/mod.rs

@@ -5,10 +5,6 @@ use std::{
         HashSet,
     },
     future::Future,
-    hash::{
-        Hash,
-        Hasher,
-    },
     sync::Arc,
 };
 
@@ -42,13 +38,7 @@ use common::{
 use futures::FutureExt;
 use hashlink::LinkedHashSet;
 use keybroker::Identity;
-use tokio::{
-    select,
-    sync::{
-        OwnedSemaphorePermit,
-        Semaphore,
-    },
-};
+use tokio::select;
 use tokio_util::task::JoinMap;
 use value::{
     DeveloperDocumentId,
@@ -77,15 +67,16 @@ use crate::{
 pub mod index_writer;
 
 pub struct IndexWorker<RT: Runtime> {
-    // Index IDs that are currently being backfilled.
+    /// Index IDs that are currently being backfilled.
     in_progress_index_ids: HashSet<IndexId, ahash::RandomState>,
-    /// Order-preserving HashMap that represents the order that pending
-    /// index backfills will be loaded in.
-    pending: LinkedHashSet<IndexBackfill, ahash::RandomState>,
-    // Index IDs that are pending backfill.
-    pending_index_ids: HashSet<IndexId, ahash::RandomState>,
-    concurrency_limiter: Arc<Semaphore>,
-    work_pool: JoinMap<IndexBackfill, anyhow::Result<()>>,
+    /// The index backfill tasks
+    in_progress: JoinMap<Vec<IndexId>, anyhow::Result<()>>,
+    /// Order-preserving HashSet that represents the order that pending index
+    /// backfills will be processed. This does not include indexes that are
+    /// `in_progress`.
+    pending: LinkedHashSet<(IndexId, TabletId), ahash::RandomState>,
+    /// Limit on the size of `in_progress`
+    max_concurrency: usize,
     database: Database<RT>,
     index_writer: IndexWriter<RT>,
     backoff: Backoff,
@@ -94,19 +85,6 @@ pub struct IndexWorker<RT: Runtime> {
     pub should_terminate: bool,
 }
 
-#[derive(Clone, Debug, Eq, PartialEq)]
-pub struct IndexBackfill {
-    pub tablet_id: TabletId,
-    pub index_ids: Vec<IndexId>,
-}
-
-impl Hash for IndexBackfill {
-    fn hash<H: Hasher>(&self, state: &mut H) {
-        self.tablet_id.hash(state);
-        self.index_ids.hash(state);
-    }
-}
-
 impl<RT: Runtime> IndexWorker<RT> {
     #[allow(clippy::new_ret_no_self)]
     pub fn new(
@@ -124,10 +102,9 @@ impl<RT: Runtime> IndexWorker<RT> {
         );
         let mut worker = IndexWorker {
             in_progress_index_ids: Default::default(),
+            in_progress: JoinMap::new(),
             pending: Default::default(),
-            pending_index_ids: Default::default(),
-            concurrency_limiter: Arc::new(Semaphore::new(*INDEX_BACKFILL_CONCURRENCY)),
-            work_pool: JoinMap::new(),
+            max_concurrency: *INDEX_BACKFILL_CONCURRENCY,
             database,
             index_writer,
             backoff: Backoff::new(*INDEX_WORKERS_INITIAL_BACKOFF, *INDEX_WORKERS_MAX_BACKOFF),
@@ -171,10 +148,9 @@ impl<RT: Runtime> IndexWorker<RT> {
         );
         let mut worker = IndexWorker {
             in_progress_index_ids: Default::default(),
+            in_progress: JoinMap::new(),
             pending: Default::default(),
-            pending_index_ids: Default::default(),
-            concurrency_limiter: Arc::new(Semaphore::new(10)),
-            work_pool: JoinMap::new(),
+            max_concurrency: 10,
             database,
             index_writer,
             backoff: Backoff::new(*INDEX_WORKERS_INITIAL_BACKOFF, *INDEX_WORKERS_MAX_BACKOFF),
@@ -194,31 +170,32 @@ impl<RT: Runtime> IndexWorker<RT> {
                     tracing::error!("IndexWorker loop failed: {e:?}");
                     continue;
                 }
-                if worker.work_pool.is_empty() && worker.pending.is_empty() {
+                if worker.in_progress.is_empty() && worker.pending.is_empty() {
                     return r;
                 }
             }
         }
     }
 
     async fn run(&mut self) -> anyhow::Result<()> {
-        let mut backfill_queue = Vec::new();
         // Get all the documents from the `_index` table.
         let mut tx = self.database.begin(Identity::system()).await?;
         // _index doesn't have `by_creation_time` index, and thus must use `by_id`.
         let index_documents = tx
             .query_system(TableNamespace::Global, &SystemIndex::<IndexTable>::by_id())?
             .all()
             .await?;
-        let mut to_backfill_by_tablet = BTreeMap::new();
         let mut num_to_backfill = 0;
         for index_metadata in &index_documents {
             if let IndexConfig::Database { on_disk_state, .. } = &index_metadata.config {
                 if matches!(on_disk_state, DatabaseIndexState::Backfilling(_)) {
-                    to_backfill_by_tablet
-                        .entry(*index_metadata.name.table())
-                        .or_insert_with(Vec::new)
-                        .push(index_metadata.id().internal_id());
+                    let index_id = index_metadata.id().internal_id();
+                    let tablet_id = *index_metadata.name.table();
+                    if !self.in_progress_index_ids.contains(&index_id)
+                        && !self.pending.contains(&(index_id, tablet_id))
+                    {
+                        self.pending.insert((index_id, tablet_id));
+                    }
                     num_to_backfill += 1;
                 }
             }
@@ -229,118 +206,73 @@ impl<RT: Runtime> IndexWorker<RT> {
             tx.begin_timestamp()
         );
 
-        for (tablet_id, index_ids) in to_backfill_by_tablet {
-            let mut to_backfill = Vec::new();
-            for index_id in index_ids {
-                if !self.in_progress_index_ids.contains(&index_id)
-                    && !self.pending_index_ids.contains(&index_id)
-                {
-                    to_backfill.push(index_id);
-                }
-            }
-            backfill_queue.push(IndexBackfill {
-                tablet_id,
-                index_ids: to_backfill,
-            });
-        }
-
         let token = tx.into_token()?;
         let subscription = self.database.subscribe(token).await?;
-        for index_backfill in backfill_queue.clone() {
-            self.pending.insert(index_backfill.clone());
-            self.pending_index_ids.extend(index_backfill.index_ids);
-        }
 
         #[cfg(any(test, feature = "testing"))]
-        if self.should_terminate
-            && backfill_queue.is_empty()
-            && self.work_pool.is_empty()
-            && self.pending.is_empty()
-        {
+        if self.should_terminate && self.in_progress.is_empty() && self.pending.is_empty() {
             return Ok(());
         }
 
+        // Start new work if allowed by the concurrency limit
+        while self.in_progress.len() < self.max_concurrency
+            && let Some((index_id, tablet_id)) = self.pending.pop_front()
+        {
+            self.queue_index_backfill(index_id, tablet_id);
+        }
         select! {
             biased;
             // Start by finding indexes that have finished backfilling
-            (index_backfill, res) = self.work_pool.join_next().map(|o| o.expect("Can't be none if work pool is nonempty")), if !self.work_pool.is_empty() => {
-                if let Err(e) = res {
-                    tracing::error!("IndexWorker loop failed: {e:?}");
-                    report_error(&mut anyhow::anyhow!(format!("{e:?}: IndexWorker loop failed on backfill {index_backfill:?}"))).await;
-                }
-
-                for index_id in &index_backfill.index_ids {
-                    self.in_progress_index_ids.remove(index_id);
+            res = self.in_progress.join_next(), if !self.in_progress.is_empty() => {
+                let (index_ids, res) = res.expect("join_next cannot return None if nonempty");
+                // First, make sure `in_progress_index_ids` is always consistent with `in_progress`
+                for &index_id in &index_ids {
+                    self.in_progress_index_ids.remove(&index_id);
                 }
+                // If backfill tasks are failing, return an error here so that we back off
+                let () = res??;
+                tracing::info!("Finished backfilling {index_ids:?}");
+                // Return so that we possibly queue up more work
             }
-            // Next, queue up new work
-            permit = self.concurrency_limiter.clone().acquire_owned(),
-                if !self.pending.is_empty() =>
-            {
-                let permit = permit.expect("Semaphore can't go away");
-                self.queue_index_backfill(permit);
-            }
-            // Finally, wait for invalidation
+            // Alternatively, wait for invalidation
             _ = subscription.wait_for_invalidation().fuse() => {
-                tracing::info!("IndexIndexWorker received invalidation, going to sleep");
                 self.backoff.reset();
             }
         }
 
         Ok(())
     }
 
-    /// Spawns a task to load the process the next index backfill. The `permit`
-    /// is held to ensure we limit our concurrency.
-    fn queue_index_backfill(&mut self, permit: OwnedSemaphorePermit) {
-        // Try to queue an index backfill
-        let Some(index_backfill) = self.pending.pop_front() else {
-            return;
-        };
-        for index_id in &index_backfill.index_ids {
-            self.pending_index_ids.remove(index_id);
-        }
+    /// Spawns a task to process the next index backfill.
+    fn queue_index_backfill(&mut self, index_id: IndexId, tablet_id: TabletId) {
+        let mut index_ids = vec![index_id];
+        // Since we're mainly limited by the speed of reading the table, let's
+        // grab all the other pending indexes for this table at once
+        self.pending.retain(|&(other_index_id, other_tablet_id)| {
+            if other_tablet_id == tablet_id {
+                index_ids.push(other_index_id);
+                false
+            } else {
+                true
+            }
+        });
+
+        // TODO: this allows more than one `backfill_tablet` task to run
+        // simultaneously on the same table; it could be better to cancel old
+        // tasks and restart them.
 
-        if self.work_pool.contains_key(&index_backfill) {
-            let IndexBackfill { tablet_id, .. } = index_backfill;
-            tracing::warn!("Skipping duplicate index backfill {tablet_id} (already in progress)");
-            return;
+        for &index_id in &index_ids {
+            self.in_progress_index_ids.insert(index_id);
         }
-        if index_backfill
-            .index_ids
-            .iter()
-            .any(|index_id| self.in_progress_index_ids.contains(index_id))
-        {
-            let IndexBackfill {
+        self.in_progress.spawn(
+            index_ids.clone(),
+            Self::backfill_tablet(
                 tablet_id,
                 index_ids,
-                ..
-            } = index_backfill;
-            tracing::warn!(
-                "Skipping backfill of indexes {index_ids:?} on tablet {tablet_id} until current \
-                 backfill completes"
-            );
-            return;
-        }
-
-        let IndexBackfill {
-            tablet_id,
-            index_ids,
-            ..
-        } = index_backfill.clone();
-
-        let database = self.database.clone();
-        let index_writer = self.index_writer.clone();
-
-        for index_id in &index_ids {
-            self.in_progress_index_ids.insert(*index_id);
-        }
-
-        self.work_pool.spawn(index_backfill, async move {
-            let _permit = permit;
-
-            Self::backfill_tablet(tablet_id, index_ids, database, index_writer).await
-        });
+                self.database.clone(),
+                self.index_writer.clone(),
+            ),
+        );
     }
 
     async fn backfill_tablet(