feat(taskbroker): Batch Status Updates and Delete Completed Tasks Immediately

src/config.rs

@@ -293,6 +293,12 @@ pub struct Config {
     /// Maximum time in milliseconds for a single push RPC to the worker service. This should be greater than the worker's internal timeout.
     pub push_timeout_ms: u64,
 
+    /// The size of a batch of status updates. Only active in push mode.
+    pub status_flush_batch_size: usize,
+
+    /// Maximum milliseconds to wait before flushing a batch of status updates.
+    pub status_flush_interval_ms: u64,
+
     /// The hostname used to construct `callback_url` for task push requests.
     pub callback_addr: String,
 
@@ -383,6 +389,8 @@ impl Default for Config {
             push_queue_size: 1,
             push_queue_timeout_ms: 5000,
             push_timeout_ms: 30000,
+            status_flush_batch_size: 1,
+            status_flush_interval_ms: 100,
             callback_addr: "0.0.0.0".into(),
             callback_port: 50051,
             worker_map: [("sentry".into(), "http://127.0.0.1:50052".into())].into(),

src/fetch/mod.rs

@@ -109,13 +109,19 @@ impl<T: TaskPusher + Send + Sync + 'static> FetchPool<T> {
                         }
 
                         _ = async {
+                            let start = Instant::now();
+
                             debug!("Fetching next batch of pending activations...");
                             metrics::counter!("fetch.loop.count").increment(1);
 
-                            let start = Instant::now();
                             let mut backoff = false;
 
-                            match store.claim_activations_for_push(limit, bucket).await {
+                            let start_claim = Instant::now();
+                            let result = store.claim_activations_for_push(limit, bucket).await;
+                            metrics::histogram!("fetch.claim_activations_for_push.duration")
+                                .record(start_claim.elapsed());
+
+                            match result {
                                 Ok(activations) if activations.is_empty() => {
                                     metrics::counter!("fetch.empty").increment(1);
                                     debug!("No pending activations");

src/fetch/tests.rs

@@ -122,6 +122,18 @@ impl InflightActivationStore for MockStore {
         unimplemented!()
     }
 
+    async fn set_status_batch(
+        &self,
+        _ids: &[String],
+        _status: InflightActivationStatus,
+    ) -> Result<(), Error> {
+        Ok(())
+    }
+
+    async fn delete_activation_batch(&self, _ids: &[String]) -> Result<u64, Error> {
+        Ok(0)
+    }
+
     async fn set_processing_deadline(
         &self,
         _id: &str,

src/flusher.rs

@@ -0,0 +1,62 @@
+use std::future::Future;
+use std::pin::Pin;
+use std::time::Duration;
+
+use anyhow::Result;
+use tokio::sync::mpsc::Receiver;
+
+/// Run flusher that receives values of type T from a channel and flushes
+/// them using the provided async `flush` function either when the batch is
+/// full or when the max flush interval has elapsed.
+pub async fn run_flusher<T, F>(
+    mut rx: Receiver<T>,
+    batch_size: usize,
+    interval_ms: u64,
+    mut flush: F,
+) -> Result<()>
+where
+    F: for<'a> FnMut(&'a mut Vec<T>) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>>,
+{
+    let batch_size = batch_size.max(1);
+    let interval_ms = interval_ms.max(1);
+
+    let period = Duration::from_millis(interval_ms);
+    let mut interval = tokio::time::interval(period);
+    interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
+
+    let mut buffer: Vec<T> = Vec::with_capacity(batch_size);
+
+    loop {
+        tokio::select! {
+            msg = rx.recv() => {
+                match msg {
+                    Some(v) => {
+                        buffer.push(v);
+
+                        while let Ok(update) = rx.try_recv() {
+                            buffer.push(update);
+                        }
+
+                        if buffer.len() >= batch_size {
+                            flush(&mut buffer).await;
+                        }
+                    }
+
+                    None => {
+                        // Channel closed (shutdown), flush remaining and exit
+                        flush(&mut buffer).await;
+                        break;
+                    }
+                }
+            }
+
+            _ = interval.tick() => {
+                if !buffer.is_empty() {
+                    flush(&mut buffer).await;
+                }
+            }
+        }
+    }
+
+    Ok(())
+}

src/grpc/mod.rs

@@ -1,5 +1,7 @@
 pub mod auth_middleware;
 pub mod metrics_middleware;
 pub mod server;
+pub mod status_flusher;
+
 #[cfg(test)]
 mod server_tests;

src/grpc/server.rs

@@ -1,15 +1,18 @@
+use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Instant;
 
+use anyhow::Result;
 use chrono::Utc;
 use prost::Message;
 use sentry_protos::taskbroker::v1::consumer_service_server::ConsumerService;
 use sentry_protos::taskbroker::v1::{
     FetchNextTask, GetTaskRequest, GetTaskResponse, SetTaskStatusRequest, SetTaskStatusResponse,
     TaskActivation, TaskActivationStatus,
 };
+use tokio::sync::mpsc;
 use tonic::{Request, Response, Status};
-use tracing::{error, instrument, warn};
+use tracing::{debug, error, instrument, warn};
 
 use crate::config::{Config, DeliveryMode};
 use crate::store::activation::InflightActivationStatus;
@@ -18,6 +21,7 @@ use crate::store::traits::InflightActivationStore;
 pub struct TaskbrokerServer {
     pub store: Arc<dyn InflightActivationStore>,
     pub config: Arc<Config>,
+    pub status_tx: Option<mpsc::Sender<StatusUpdate>>,
 }
 
 #[tonic::async_trait]
@@ -97,10 +101,20 @@ impl ConsumerService for TaskbrokerServer {
                 "Invalid status, expects 3 (Failure), 4 (Retry), or 5 (Complete), but got: {status:?}"
             )));
         }
+
         if status == InflightActivationStatus::Failure {
             metrics::counter!("grpc_server.set_status.failure").increment(1);
         }
 
+        if let Some(ref tx) = self.status_tx {
+            tx.send((id, status))
+                .await
+                .map_err(|_| Status::internal("Status update channel closed"))?;
+
+            metrics::histogram!("grpc_server.set_status.duration").record(start_time.elapsed());
+            return Ok(Response::new(SetTaskStatusResponse { task: None }));
+        }
+
         match self.store.set_status(&id, status).await {
             Ok(Some(_)) => metrics::counter!(
                 "grpc_server.set_status",
@@ -194,3 +208,58 @@ impl ConsumerService for TaskbrokerServer {
         res
     }
 }
+
+pub type StatusUpdate = (String, InflightActivationStatus);
+
+pub async fn flush_status_updates(
+    store: Arc<dyn InflightActivationStore>,
+    buffer: &mut Vec<StatusUpdate>,
+) {
+    if buffer.is_empty() {
+        return;
+    }
+
+    let updates = std::mem::take(buffer);
+    let mut by_status: HashMap<InflightActivationStatus, Vec<String>> = HashMap::new();
+
+    for (id, status) in updates {
+        by_status.entry(status).or_default().push(id);
+    }
+
+    let mut success = 0;
+    let mut fail = 0;
+
+    for (status, ids) in by_status {
+        let count = ids.len() as u64;
+
+        match store
+            .set_status_batch(&ids, status)
+            .await
+            .map(|()| ids.len() as u64)
+        {
+            Ok(count) => {
+                success += count;
+                debug!(?status, ?count, "Flushed status batch");
+            }
+
+            Err(e) => {
+                fail += count;
+
+                error!(
+                    ?status,
+                    ?count,
+                    error = ?e,
+                    "Failed to flush status batch"
+                );
+
+                // Push failed updates back into the buffer so they can be retried on next flush
+                for id in ids {
+                    buffer.push((id, status));
+                }
+            }
+        }
+    }
+
+    metrics::gauge!(format!("grpc_server.flush_status_updates.success")).set(success as f64);
+    metrics::gauge!(format!("grpc_server.flush_status_updates.fail")).set(fail as f64);
+}