Add integration test

Author: rdettai

quickwit/Cargo.lock

@@ -264,20 +264,26 @@ pub enum FileSourceMessageType {
 pub struct FileSourceSqs {
     pub queue_url: String,
     pub message_type: FileSourceMessageType,
-    #[serde(default = "default_deduplication_window_duration_sec")]
-    pub deduplication_window_duration_sec: u32,
+    #[serde(default = "default_deduplication_window_duration_secs")]
+    pub deduplication_window_duration_secs: u32,
     #[serde(default = "default_deduplication_window_max_messages")]
     pub deduplication_window_max_messages: u32,
+    #[serde(default = "default_checkpoint_cleanup_interval_secs")]
+    pub checkpoint_cleanup_interval_secs: u32,
 }
 
-fn default_deduplication_window_duration_sec() -> u32 {
+fn default_deduplication_window_duration_secs() -> u32 {
     3600
 }
 
 fn default_deduplication_window_max_messages() -> u32 {
     100_000
 }
 
+fn default_checkpoint_cleanup_interval_secs() -> u32 {
+    60
+}
+
 #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, utoipa::ToSchema)]
 #[serde(tag = "type", rename_all = "snake_case")]
 pub enum FileSourceNotification {
@@ -903,8 +909,10 @@ mod tests {
                     queue_url: "https://sqs.us-east-1.amazonaws.com/123456789012/queue-name"
                         .to_string(),
                     message_type: FileSourceMessageType::S3Notification,
-                    deduplication_window_duration_sec: default_deduplication_window_duration_sec(),
+                    deduplication_window_duration_secs: default_deduplication_window_duration_secs(
+                    ),
                     deduplication_window_max_messages: default_deduplication_window_max_messages(),
+                    checkpoint_cleanup_interval_secs: default_checkpoint_cleanup_interval_secs()
                 })),
             );
             let file_params_reserialized = serde_json::to_value(&file_params_deserialized).unwrap();
@@ -914,8 +922,9 @@ mod tests {
                     "type": "sqs",
                     "queue_url": "https://sqs.us-east-1.amazonaws.com/123456789012/queue-name",
                     "message_type": "s3_notification",
-                    "deduplication_window_duration_sec": default_deduplication_window_duration_sec(),
-                    "deduplication_window_max_messages": default_deduplication_window_max_messages()
+                    "deduplication_window_duration_secs": default_deduplication_window_duration_secs(),
+                    "deduplication_window_max_messages": default_deduplication_window_max_messages(),
+                    "checkpoint_cleanup_interval_secs": default_checkpoint_cleanup_interval_secs(),
                 }]})
             );
         }

quickwit/quickwit-config/src/source_config/mod.rs

@@ -435,8 +435,9 @@ mod localstack_tests {
             FileSourceParams::Notifications(FileSourceNotification::Sqs(FileSourceSqs {
                 queue_url,
                 message_type: FileSourceMessageType::RawUri,
-                deduplication_window_duration_sec: 100,
+                deduplication_window_duration_secs: 100,
                 deduplication_window_max_messages: 100,
+                checkpoint_cleanup_interval_secs: 60,
             }));
         let source_config = SourceConfig::for_test(
             "test-file-source-sqs-notifications",

quickwit/quickwit-indexing/src/source/file_source.rs

@@ -19,7 +19,7 @@
 
 use std::fmt;
 use std::sync::Arc;
-use std::time::{Duration, Instant};
+use std::time::Duration;
 
 use itertools::Itertools;
 use quickwit_actors::{ActorExitStatus, Mailbox};
@@ -98,20 +98,18 @@ impl QueueCoordinator {
         message_type: MessageType,
         shard_max_age: Option<u32>,
         shard_max_count: Option<u32>,
+        shard_pruning_interval: Duration,
     ) -> Self {
         Self {
-            shared_state: QueueSharedState {
-                metastore: source_runtime.metastore,
-                source_id: source_runtime.pipeline_id.source_id.clone(),
-                index_uid: source_runtime.pipeline_id.index_uid.clone(),
-                reacquire_grace_period: Duration::from_secs(
-                    2 * source_runtime.indexing_setting.commit_timeout_secs as u64,
-                ),
-                last_initiated_pruning: Instant::now(),
-                max_age: shard_max_age,
-                max_count: shard_max_count,
-                pruning_interval: Duration::from_secs(60),
-            },
+            shared_state: QueueSharedState::new(
+                source_runtime.metastore,
+                source_runtime.pipeline_id.index_uid.clone(),
+                source_runtime.pipeline_id.source_id.clone(),
+                Duration::from_secs(2 * source_runtime.indexing_setting.commit_timeout_secs as u64),
+                shard_max_age,
+                shard_max_count,
+                shard_pruning_interval,
+            ),
             local_state: QueueLocalState::default(),
             pipeline_id: source_runtime.pipeline_id,
             source_type: source_runtime.source_config.source_type(),
@@ -143,8 +141,9 @@ impl QueueCoordinator {
             source_runtime,
             Arc::new(queue),
             message_type,
-            Some(config.deduplication_window_duration_sec),
+            Some(config.deduplication_window_duration_secs),
             Some(config.deduplication_window_max_messages),
+            Duration::from_secs(config.checkpoint_cleanup_interval_secs as u64),
         ))
     }
 

quickwit/quickwit-indexing/src/source/queue_sources/coordinator.rs

@@ -18,7 +18,8 @@
 // along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 use std::collections::BTreeMap;
-use std::time::{Duration, Instant};
+use std::sync::Arc;
+use std::time::Duration;
 
 use anyhow::{bail, Context};
 use quickwit_metastore::checkpoint::PartitionId;
@@ -34,35 +35,78 @@ use super::message::PreProcessedMessage;
 
 #[derive(Clone)]
 pub struct QueueSharedState {
-    pub metastore: MetastoreServiceClient,
+    metastore: MetastoreServiceClient,
     pub index_uid: IndexUid,
     pub source_id: String,
     /// Duration after which the processing of a shard is considered stale and
     /// should be reacquired
-    pub reacquire_grace_period: Duration,
-    pub max_age: Option<u32>,
-    pub max_count: Option<u32>,
-    pub last_initiated_pruning: Instant,
-    pub pruning_interval: Duration,
+    reacquire_grace_period: Duration,
+    _cleanup_handle: Arc<()>,
 }
 
 impl QueueSharedState {
-    async fn clean_partitions(&self) {
-        if self.max_count.is_none() && self.max_age.is_none() {
-            return;
+    /// Create a shared state service and runs a cleanup task that prunes shards
+    /// in the background
+    pub fn new(
+        metastore: MetastoreServiceClient,
+        index_uid: IndexUid,
+        source_id: String,
+        reacquire_grace_period: Duration,
+        max_age: Option<u32>,
+        max_count: Option<u32>,
+        pruning_interval: Duration,
+    ) -> Self {
+        let cleanup_handle = Arc::new(());
+        tokio::spawn(Self::run_cleanup_task(
+            metastore.clone(),
+            index_uid.clone(),
+            source_id.clone(),
+            max_age,
+            max_count,
+            pruning_interval,
+            cleanup_handle.clone(),
+        ));
+        Self {
+            metastore,
+            index_uid,
+            source_id,
+            reacquire_grace_period,
+            _cleanup_handle: cleanup_handle,
         }
-        let result = self
-            .metastore
-            .prune_shards(PruneShardsRequest {
-                index_uid: Some(self.index_uid.clone()),
-                source_id: self.source_id.clone(),
-                max_age: self.max_age,
-                max_count: self.max_count,
-            })
-            .await;
-        if let Err(err) = result {
-            error!(error = ?err, "failed to prune shards");
+    }
+
+    async fn run_cleanup_task(
+        metastore: MetastoreServiceClient,
+        index_uid: IndexUid,
+        source_id: String,
+        max_age: Option<u32>,
+        max_count: Option<u32>,
+        pruning_interval: Duration,
+        owner_handle: Arc<()>,
+    ) {
+        if max_count.is_none() && max_age.is_none() {
+            return;
         }
+        tokio::spawn(async move {
+            let mut interval = tokio::time::interval(pruning_interval);
+            loop {
+                interval.tick().await;
+                if Arc::strong_count(&owner_handle) == 1 {
+                    break;
+                }
+                let result: Result<_, _> = metastore
+                    .prune_shards(PruneShardsRequest {
+                        index_uid: Some(index_uid.clone()),
+                        source_id: source_id.clone(),
+                        max_age,
+                        max_count,
+                    })
+                    .await;
+                if let Err(err) = result {
+                    error!(error = ?err, "failed to prune shards");
+                }
+            }
+        });
     }
 
     /// Tries to acquire the ownership for the provided messages from the global
@@ -75,13 +119,6 @@ impl QueueSharedState {
         publish_token: &str,
         partitions: Vec<PartitionId>,
     ) -> anyhow::Result<Vec<(PartitionId, Position)>> {
-        if self.last_initiated_pruning.elapsed() > self.pruning_interval {
-            let self_cloned = self.clone();
-            tokio::spawn(async move {
-                self_cloned.clean_partitions().await;
-            });
-            self.last_initiated_pruning = Instant::now();
-        }
         let open_shard_subrequests = partitions
             .iter()
             .enumerate()
@@ -323,10 +360,7 @@ pub mod shared_state_for_tests {
             index_uid,
             source_id: "test-queue-src".to_string(),
             reacquire_grace_period: Duration::from_secs(10),
-            last_initiated_pruning: Instant::now(),
-            max_age: None,
-            max_count: None,
-            pruning_interval: Duration::from_secs(10),
+            _cleanup_handle: Arc::new(()),
         }
     }
 }
@@ -378,10 +412,7 @@ mod tests {
             index_uid,
             source_id: "test-sqs-source".to_string(),
             reacquire_grace_period: Duration::from_secs(10),
-            last_initiated_pruning: Instant::now(),
-            max_age: None,
-            max_count: None,
-            pruning_interval: Duration::from_secs(10),
+            _cleanup_handle: Arc::new(()),
         };
 
         let aquired = shared_state
@@ -411,10 +442,7 @@ mod tests {
             index_uid,
             source_id: "test-sqs-source".to_string(),
             reacquire_grace_period: Duration::from_secs(10),
-            last_initiated_pruning: Instant::now(),
-            max_age: None,
-            max_count: None,
-            pruning_interval: Duration::from_secs(10),
+            _cleanup_handle: Arc::new(()),
         };
 
         let acquired = shared_state
@@ -444,10 +472,7 @@ mod tests {
             index_uid,
             source_id: "test-sqs-source".to_string(),
             reacquire_grace_period: Duration::from_secs(10),
-            last_initiated_pruning: Instant::now(),
-            max_age: None,
-            max_count: None,
-            pruning_interval: Duration::from_secs(10),
+            _cleanup_handle: Arc::new(()),
         };
 
         let aquired = shared_state
@@ -481,10 +506,7 @@ mod tests {
             index_uid,
             source_id: "test-sqs-source".to_string(),
             reacquire_grace_period: Duration::from_secs(10),
-            last_initiated_pruning: Instant::now(),
-            max_age: None,
-            max_count: None,
-            pruning_interval: Duration::from_secs(10),
+            _cleanup_handle: Arc::new(()),
         };
 
         let checkpointed_msg = checkpoint_messages(&mut shared_state, "token1", source_messages)