Filter out non-local updates in shard positions update event handler

quickwit/quickwit-ingest/src/ingest_v2/ingester.rs

@@ -45,7 +45,7 @@ use quickwit_proto::types::{
 use serde_json::{Value as JsonValue, json};
 use tokio::sync::Semaphore;
 use tokio::time::{sleep, timeout};
-use tracing::{debug, error, info, instrument, warn};
+use tracing::{Span, debug, error, info, instrument, warn};
 
 use super::IngesterPool;
 use super::broadcast::{BroadcastIngesterCapacityScoreTask, BroadcastLocalShardsTask};
@@ -1258,31 +1258,117 @@ impl IngesterService for Ingester {
 
 #[async_trait]
 impl EventSubscriber<ShardPositionsUpdate> for WeakIngesterState {
-    #[instrument(name = "ingester.shard_positions_gossip", skip_all)]
+    #[instrument(name = "ingester.truncate_shards_gossip", skip_all)]
     async fn handle_event(&mut self, shard_positions_update: ShardPositionsUpdate) {
         let Some(state) = self.upgrade() else {
-            warn!("ingester state update failed");
+            debug!("ingester was dropped: exiting");
             return;
         };
-        let Ok(mut state_guard) = state.lock_fully("truncate_shards_gossip").await else {
-            error!("failed to lock the ingester state");
+        let local_updates = filter_local_shard_updates(&state, shard_positions_update).await;
+
+        if local_updates.is_empty() {
             return;
-        };
-        let index_uid = shard_positions_update.source_uid.index_uid;
-        let source_id = shard_positions_update.source_uid.source_id;
+        }
+        // We're in no rush to process the updates, so yield to avoid starving other tasks waiting
+        // for the lock.
+        tokio::task::yield_now().await;
+
+        apply_local_shard_updates(&state, local_updates).await;
+    }
+}
 
-        for (shard_id, shard_position) in shard_positions_update.updated_shard_positions {
-            let queue_id = queue_id(&index_uid, &source_id, &shard_id);
+/// The gossiped update is not scoped to this ingester: it carries the positions of every shard
+/// of the source, most of which are typically hosted by other ingesters. This function filters
+/// down to the updates that will actually mutate our local state, using the cheap partial lock
+/// (`inner` only), sparing the caller from taking the full lock, which also holds the WAL write
+/// lock contended by persist/fetch operations on unrelated shards. The per-entry conditions
+/// below mirror the no-op checks performed by `delete_shard` and `truncate_shard` so we can
+/// discard useless entries without ever taking the full lock.
+#[instrument(
+    name = "ingester.filter_local_shard_updates",
+    skip_all,
+    fields(num_global_updates, num_local_updates)
+)]
+async fn filter_local_shard_updates(
+    state: &IngesterState,
+    shard_positions_update: ShardPositionsUpdate,
+) -> Vec<(QueueId, Position)> {
+    let index_uid = shard_positions_update.source_uid.index_uid;
+    let source_id = shard_positions_update.source_uid.source_id;
+
+    let Ok(state_guard) = state.lock_partially("filter_local_shard_updates").await else {
+        debug!("ingester was dropped: exiting");
+        return Vec::new();
+    };
+    let num_global_updates = shard_positions_update.updated_shard_positions.len();
+
+    let local_updates: Vec<(QueueId, Position)> = shard_positions_update
+        .updated_shard_positions
+        .into_iter()
+        .map(|(shard_id, shard_position)| {
+            (queue_id(&index_uid, &source_id, &shard_id), shard_position)
+        })
+        .filter(|(queue_id, shard_position)| {
+            let Some(shard) = state_guard.shards.get(queue_id) else {
+                return false;
+            };
             if shard_position.is_eof() {
-                state_guard.delete_shard(&queue_id, "indexer gossip").await;
-            } else if !shard_position.is_beginning() {
-                state_guard
-                    .truncate_shard(&queue_id, shard_position, "indexer gossip")
-                    .await;
+                return true;
+            }
+            if shard_position.is_beginning() {
+                return false;
             }
+            shard.truncation_position_inclusive < *shard_position
+        })
+        .collect();
+
+    Span::current().record("num_global_updates", num_global_updates);
+    Span::current().record("num_local_updates", local_updates.len());
+
+    debug!(
+        "filtered out {} of {num_global_updates} shard position update(s)",
+        num_global_updates - local_updates.len(),
+    );
+    local_updates
+}
+
+#[instrument(
+    name = "ingester.apply_local_shard_updates",
+    skip_all,
+    fields(num_deleted_shards, num_truncated_shards)
+)]
+async fn apply_local_shard_updates(state: &IngesterState, local_updates: Vec<(QueueId, Position)>) {
+    let now = Instant::now();
+
+    let Ok(mut state_guard) = state.lock_fully("apply_local_shard_updates").await else {
+        debug!("ingester was dropped: exiting");
+        return;
+    };
+    let mut num_deleted_shards = 0;
+    let mut num_truncated_shards = 0;
+
+    for (queue_id, shard_position) in local_updates {
+        if shard_position.is_eof() {
+            state_guard.delete_shard(&queue_id, "indexer gossip").await;
+            num_deleted_shards += 1;
+        } else if !shard_position.is_beginning() {
+            state_guard
+                .truncate_shard(&queue_id, shard_position, "indexer gossip")
+                .await;
+            num_truncated_shards += 1;
         }
-        state_guard.check_decommissioning_status().await;
     }
+    state_guard.check_decommissioning_status().await;
+
+    Span::current().record("num_deleted_shards", num_deleted_shards);
+    Span::current().record("num_truncated_shards", num_truncated_shards);
+
+    info!(
+        "deleted {} and truncated {} shard(s) via gossip in {}",
+        num_deleted_shards,
+        num_truncated_shards,
+        now.elapsed().pretty_display()
+    );
 }
 
 struct PendingPersistSubrequest {
@@ -1308,6 +1394,7 @@ mod tests {
     use bytes::Bytes;
     use quickwit_cluster::{ChitchatTransport, create_cluster_for_test_with_id};
     use quickwit_common::shared_consts::INGESTER_PRIMARY_SHARDS_PREFIX;
+    use quickwit_common::test_utils::wait_until_predicate;
     use quickwit_common::tower::ConstantRate;
     use quickwit_config::service::QuickwitService;
     use quickwit_proto::control_plane::{AdviseResetShardsResponse, MockControlPlaneService};
@@ -1320,7 +1407,6 @@ mod tests {
         ShardState,
     };
     use quickwit_proto::types::{DocMappingUid, DocUid, ShardId, SourceUid, queue_id};
-    use tokio::task::yield_now;
     use tokio::time::timeout;
     use tonic::transport::{Endpoint, Server};
 
@@ -3935,8 +4021,23 @@ mod tests {
         // Verify idempotency.
         event_broker.publish(shard_position_update);
 
-        // Yield so that the event is processed.
-        yield_now().await;
+        // Wait for both events to be processed.
+        wait_until_predicate(
+            || async {
+                ingester
+                    .state
+                    .lock_fully("test")
+                    .await
+                    .unwrap()
+                    .shards
+                    .len()
+                    == 1
+            },
+            Duration::from_secs(5),
+            Duration::from_millis(100),
+        )
+        .await
+        .expect("shard `2` should be deleted");
 
         let state_guard = ingester.state.lock_fully("test").await.unwrap();
         assert_eq!(state_guard.shards.len(), 1);
@@ -3951,6 +4052,83 @@ mod tests {
         assert!(!state_guard.mrecordlog.queue_exists(&queue_id_02));
     }
 
+    #[tokio::test]
+    async fn test_filter_local_shard_updates() {
+        let (_ingester_ctx, ingester) = IngesterForTest::default().build().await;
+
+        let index_uid = IndexUid::for_test("test-index", 0);
+        let source_id = SourceId::from("test-source");
+
+        let doc_mapping_uid = DocMappingUid::random();
+        let doc_mapping_json = format!(
+            r#"{{
+                "doc_mapping_uid": "{doc_mapping_uid}"
+            }}"#
+        );
+        let shard_01 = Shard {
+            index_uid: Some(index_uid.clone()),
+            source_id: source_id.clone(),
+            shard_id: Some(ShardId::from(1)),
+            shard_state: ShardState::Open as i32,
+            doc_mapping_uid: Some(doc_mapping_uid),
+            ..Default::default()
+        };
+        let queue_id_01 = queue_id(&index_uid, &source_id, &ShardId::from(1));
+
+        let mut state_guard = ingester.state.lock_fully("test").await.unwrap();
+        let now = Instant::now();
+
+        ingester
+            .init_primary_shard(
+                &mut state_guard.inner,
+                &mut state_guard.mrecordlog,
+                shard_01,
+                &doc_mapping_json,
+                now,
+                true,
+            )
+            .await
+            .unwrap();
+
+        // Truncate the shard once so that its truncation position is already at offset 0. This
+        // lets us exercise the "stale update" case below.
+        state_guard
+            .truncate_shard(&queue_id_01, Position::offset(0u64), "test")
+            .await;
+
+        drop(state_guard);
+
+        let shard_position_update = ShardPositionsUpdate {
+            source_uid: SourceUid {
+                index_uid,
+                source_id,
+            },
+            updated_shard_positions: vec![
+                // Shard hosted by another ingester: filtered out.
+                (ShardId::from(2), Position::offset(5u64)),
+                // Local shard, but the position does not advance the truncation position:
+                // filtered out.
+                (ShardId::from(1), Position::offset(0u64)),
+                // Local shard, `Beginning` position: filtered out.
+                (ShardId::from(1), Position::Beginning),
+                // Local shard, advances the truncation position: kept.
+                (ShardId::from(1), Position::offset(1u64)),
+                // Local shard, EOF position: always kept, regardless of the current truncation
+                // position.
+                (ShardId::from(1), Position::eof(0u64)),
+            ],
+        };
+        let local_updates =
+            filter_local_shard_updates(&ingester.state, shard_position_update).await;
+        assert_eq!(
+            local_updates,
+            vec![
+                (queue_id_01.clone(), Position::offset(1u64)),
+                (queue_id_01, Position::eof(0u64)),
+            ]
+        );
+    }
+
     #[tokio::test]
     async fn test_ingester_closes_idle_shards() {
         // The `CloseIdleShardsTask` task is already unit tested, so this test ensures the task is