feat(s3-push): delete(ack) event messages after processed

src/execution/live_updater.rs

@@ -31,6 +31,39 @@ struct StatsReportState {
 const MIN_REPORT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(5);
 const REPORT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(10);
 
+struct SharedAckFn {
+    count: usize,
+    ack_fn: Option<Box<dyn FnOnce() -> BoxFuture<'static, Result<()>> + Send + Sync>>,
+}
+
+impl SharedAckFn {
+    fn new(
+        count: usize,
+        ack_fn: Box<dyn FnOnce() -> BoxFuture<'static, Result<()>> + Send + Sync>,
+    ) -> Self {
+        Self {
+            count,
+            ack_fn: Some(ack_fn),
+        }
+    }
+
+    async fn ack(v: &Mutex<Self>) -> Result<()> {
+        let ack_fn = {
+            let mut v = v.lock().unwrap();
+            v.count -= 1;
+            if v.count > 0 {
+                None
+            } else {
+                v.ack_fn.take()
+            }
+        };
+        if let Some(ack_fn) = ack_fn {
+            ack_fn().await?;
+        }
+        Ok(())
+    }
+}
+
 async fn update_source(
     flow_ctx: Arc<FlowContext>,
     plan: Arc<plan::ExecutionPlan>,
@@ -92,13 +125,50 @@ async fn update_source(
             futs.push(
                 async move {
                     let mut change_stream = change_stream;
-                    while let Some(change) = change_stream.next().await {
-                        tokio::spawn(source_context.clone().process_source_key(
-                            change.key,
-                            change.data,
-                            source_update_stats.clone(),
-                            pool.clone(),
-                        ));
+                    let retry_options = retriable::RetryOptions {
+                        max_retries: None,
+                        initial_backoff: std::time::Duration::from_secs(5),
+                        max_backoff: std::time::Duration::from_secs(60),
+                    };
+                    loop {
+                        // Workaround as AsyncFnMut isn't mature yet.
+                        // Should be changed to use AsyncFnMut once it is.
+                        let change_stream = tokio::sync::Mutex::new(&mut change_stream);
+                        let change_msg = retriable::run(
+                            || async {
+                                let mut change_stream = change_stream.lock().await;
+                                change_stream
+                                    .next()
+                                    .await
+                                    .transpose()
+                                    .map_err(retriable::Error::always_retryable)
+                            },
+                            &retry_options,
+                        )
+                        .await?;
+                        let change_msg = if let Some(change_msg) = change_msg {
+                            change_msg
+                        } else {
+                            break;
+                        };
+                        let ack_fn = change_msg.ack_fn.map(|ack_fn| {
+                            Arc::new(Mutex::new(SharedAckFn::new(
+                                change_msg.changes.iter().len(),
+                                ack_fn,
+                            )))
+                        });
+                        for change in change_msg.changes {
+                            let ack_fn = ack_fn.clone();
+                            tokio::spawn(source_context.clone().process_source_key(
+                                change.key,
+                                change.data,
+                                source_update_stats.clone(),
+                                ack_fn.map(|ack_fn| {
+                                    move || async move { SharedAckFn::ack(&ack_fn).await }
+                                }),
+                                pool.clone(),
+                            ));
+                        }
                     }
                     Ok(())
                 }

src/execution/source_indexer.rs

@@ -1,5 +1,6 @@
 use crate::prelude::*;
 
+use futures::future::Ready;
 use sqlx::PgPool;
 use std::collections::{hash_map, HashMap};
 use tokio::{sync::Semaphore, task::JoinSet};
@@ -36,6 +37,8 @@ pub struct SourceIndexingContext {
     state: Mutex<SourceIndexingState>,
 }
 
+pub const NO_ACK: Option<fn() -> Ready<Result<()>>> = None;
+
 impl SourceIndexingContext {
     pub async fn load(
         flow: Arc<builder::AnalyzedFlow>,
@@ -79,11 +82,15 @@ impl SourceIndexingContext {
         })
     }
 
-    pub async fn process_source_key(
+    pub async fn process_source_key<
+        AckFut: Future<Output = Result<()>> + Send + 'static,
+        AckFn: FnOnce() -> AckFut,
+    >(
         self: Arc<Self>,
         key: value::KeyValue,
         source_data: Option<interface::SourceData>,
         update_stats: Arc<stats::UpdateStats>,
+        ack_fn: Option<AckFn>,
         pool: PgPool,
     ) {
         let process = async {
@@ -173,11 +180,20 @@ impl SourceIndexingContext {
                 }
             }
             drop(permit);
+            if let Some(ack_fn) = ack_fn {
+                ack_fn().await?;
+            }
             anyhow::Ok(())
         };
         if let Err(e) = process.await {
             update_stats.num_errors.inc(1);
-            error!("{:?}", e.context("Error in processing a source row"));
+            error!(
+                "{:?}",
+                e.context(format!(
+                    "Error in processing row from source `{source}` with key: {key}",
+                    source = self.flow.flow_instance.import_ops[self.source_idx].name
+                ))
+            );
         }
     }
 
@@ -203,7 +219,7 @@ impl SourceIndexingContext {
         }
         Some(
             self.clone()
-                .process_source_key(key, None, update_stats.clone(), pool.clone()),
+                .process_source_key(key, None, update_stats.clone(), NO_ACK, pool.clone()),
         )
     }
 
@@ -269,6 +285,7 @@ impl SourceIndexingContext {
                 key,
                 source_data,
                 update_stats.clone(),
+                NO_ACK,
                 pool.clone(),
             ));
         }

src/ops/interface.rs

@@ -95,6 +95,11 @@ pub struct SourceChange {
     pub data: Option<SourceData>,
 }
 
+pub struct SourceChangeMessage {
+    pub changes: Vec<SourceChange>,
+    pub ack_fn: Option<Box<dyn FnOnce() -> BoxFuture<'static, Result<()>> + Send + Sync>>,
+}
+
 #[derive(Debug, Default)]
 pub struct SourceExecutorListOptions {
     pub include_ordinal: bool,
@@ -141,7 +146,9 @@ pub trait SourceExecutor: Send + Sync {
         options: &SourceExecutorGetOptions,
     ) -> Result<PartialSourceRowData>;
 
-    async fn change_stream(&self) -> Result<Option<BoxStream<'async_trait, SourceChange>>> {
+    async fn change_stream(
+        &self,
+    ) -> Result<Option<BoxStream<'async_trait, Result<SourceChangeMessage>>>> {
         Ok(None)
     }
 }

src/ops/sources/amazon_s3.rs

@@ -3,7 +3,6 @@ use async_stream::try_stream;
 use aws_config::BehaviorVersion;
 use aws_sdk_s3::Client;
 use globset::{Glob, GlobSet, GlobSetBuilder};
-use log::warn;
 use std::sync::Arc;
 
 use crate::base::field_attrs;
@@ -23,6 +22,20 @@ struct SqsContext {
     client: aws_sdk_sqs::Client,
     queue_url: String,
 }
+
+impl SqsContext {
+    async fn delete_message(&self, receipt_handle: String) -> Result<()> {
+        error!("Deleting message: {}", receipt_handle);
+        self.client
+            .delete_message()
+            .queue_url(&self.queue_url)
+            .receipt_handle(receipt_handle)
+            .send()
+            .await?;
+        Ok(())
+    }
+}
+
 struct Executor {
     client: Client,
     bucket_name: String,
@@ -152,24 +165,26 @@ impl SourceExecutor for Executor {
         Ok(PartialSourceRowData { value, ordinal })
     }
 
-    async fn change_stream(&self) -> Result<Option<BoxStream<'async_trait, SourceChange>>> {
+    async fn change_stream(
+        &self,
+    ) -> Result<Option<BoxStream<'async_trait, Result<SourceChangeMessage>>>> {
         let sqs_context = if let Some(sqs_context) = &self.sqs_context {
             sqs_context
         } else {
             return Ok(None);
         };
         let stream = stream! {
             loop {
-                let changes = match self.poll_sqs(&sqs_context).await {
-                    Ok(changes) => changes,
+                 match self.poll_sqs(&sqs_context).await {
+                    Ok(messages) => {
+                        for message in messages {
+                            yield Ok(message);
+                        }
+                    }
                     Err(e) => {
-                        warn!("Failed to poll SQS: {}", e);
-                        continue;
+                        yield Err(e);
                     }
                 };
-                for change in changes {
-                    yield change;
-                }
             }
         };
         Ok(Some(stream.boxed()))
@@ -206,7 +221,7 @@ pub struct S3Object {
 }
 
 impl Executor {
-    async fn poll_sqs(&self, sqs_context: &Arc<SqsContext>) -> Result<Vec<SourceChange>> {
+    async fn poll_sqs(&self, sqs_context: &Arc<SqsContext>) -> Result<Vec<SourceChangeMessage>> {
         let resp = sqs_context
             .client
             .receive_message()
@@ -220,36 +235,53 @@ impl Executor {
         } else {
             return Ok(Vec::new());
         };
-        let mut changes = vec![];
-        for message in messages.into_iter().filter_map(|m| m.body) {
-            let notification: S3EventNotification = serde_json::from_str(&message)?;
-            for record in notification.records {
-                let s3 = if let Some(s3) = record.s3 {
-                    s3
-                } else {
-                    continue;
-                };
-                if s3.bucket.name != self.bucket_name {
-                    continue;
-                }
-                if !self
-                    .prefix
-                    .as_ref()
-                    .map_or(true, |prefix| s3.object.key.starts_with(prefix))
-                {
-                    continue;
+        let mut change_messages = vec![];
+        for message in messages.into_iter() {
+            if let Some(body) = message.body {
+                let notification: S3EventNotification = serde_json::from_str(&body)?;
+                let mut changes = vec![];
+                for record in notification.records {
+                    let s3 = if let Some(s3) = record.s3 {
+                        s3
+                    } else {
+                        continue;
+                    };
+                    if s3.bucket.name != self.bucket_name {
+                        continue;
+                    }
+                    if !self
+                        .prefix
+                        .as_ref()
+                        .map_or(true, |prefix| s3.object.key.starts_with(prefix))
+                    {
+                        continue;
+                    }
+                    if record.event_name.starts_with("ObjectCreated:")
+                        || record.event_name.starts_with("ObjectDeleted:")
+                    {
+                        changes.push(SourceChange {
+                            key: KeyValue::Str(s3.object.key.into()),
+                            data: None,
+                        });
+                    }
                 }
-                if record.event_name.starts_with("ObjectCreated:")
-                    || record.event_name.starts_with("ObjectDeleted:")
-                {
-                    changes.push(SourceChange {
-                        key: KeyValue::Str(s3.object.key.into()),
-                        data: None,
-                    });
+                if let Some(receipt_handle) = message.receipt_handle {
+                    if !changes.is_empty() {
+                        let sqs_context = sqs_context.clone();
+                        change_messages.push(SourceChangeMessage {
+                            changes,
+                            ack_fn: Some(Box::new(move || {
+                                async move { sqs_context.delete_message(receipt_handle).await }
+                                    .boxed()
+                            })),
+                        });
+                    } else {
+                        sqs_context.delete_message(receipt_handle).await?;
+                    }
                 }
             }
         }
-        Ok(changes)
+        Ok(change_messages)
     }
 }
 

src/ops/sources/google_drive.rs

@@ -190,7 +190,7 @@ impl Executor {
     async fn get_recent_updates(
         &self,
         cutoff_time: &mut DateTime<Utc>,
-    ) -> Result<Vec<SourceChange>> {
+    ) -> Result<SourceChangeMessage> {
         let mut page_size: i32 = 10;
         let mut next_page_token: Option<String> = None;
         let mut changes = Vec::new();
@@ -234,7 +234,10 @@ impl Executor {
             page_size = 100;
         }
         *cutoff_time = Self::make_cutoff_time(most_recent_modified_time, start_time);
-        Ok(changes)
+        Ok(SourceChangeMessage {
+            changes,
+            ack_fn: None,
+        })
     }
 
     async fn is_file_covered(&self, file_id: &str) -> Result<bool> {
@@ -416,7 +419,9 @@ impl SourceExecutor for Executor {
         Ok(PartialSourceRowData { value, ordinal })
     }
 
-    async fn change_stream(&self) -> Result<Option<BoxStream<'async_trait, SourceChange>>> {
+    async fn change_stream(
+        &self,
+    ) -> Result<Option<BoxStream<'async_trait, Result<SourceChangeMessage>>>> {
         let poll_interval = if let Some(poll_interval) = self.recent_updates_poll_interval {
             poll_interval
         } else {
@@ -428,17 +433,7 @@ impl SourceExecutor for Executor {
         let stream = stream! {
             loop {
                 interval.tick().await;
-                let changes = self.get_recent_updates(&mut cutoff_time).await;
-                match changes {
-                    Ok(changes) => {
-                        for change in changes {
-                            yield change;
-                        }
-                    }
-                    Err(e) => {
-                        error!("Error getting recent updates: {e}");
-                    }
-                }
+                yield self.get_recent_updates(&mut cutoff_time).await;
             }
         };
         Ok(Some(stream.boxed()))