feat(taskbroker): Retry activations that miss their heartbeat

Cargo.toml

@@ -50,7 +50,7 @@ sentry = { version = "0.41.0", default-features = false, features = [
     "tracing",
     "logs",
 ] }
-sentry_protos = "0.22.0"
+sentry_protos = "0.106.0"
 # TLS for sentry's transport. Not using sentry's `native-tls` feature, which pulls ureq + webpki-root-certs into Cargo.lock.
 # Keep this on the same reqwest version as the `sentry` crate. When upgrading sentry (newer releases use reqwest 0.13),
 # bump this too, otherwise Cargo builds two reqwest copies and sentry's copy loses this TLS feature.

migrations/postgres/0005_add_heartbeat.sql

@@ -0,0 +1,2 @@
+ALTER TABLE inflight_taskactivations ADD COLUMN heartbeat INTEGER;
+ALTER TABLE inflight_taskactivations ADD COLUMN heartbeat_deadline TIMESTAMPTZ;

migrations/sqlite/0007_add_heartbeat.sql

@@ -0,0 +1,2 @@
+ALTER TABLE inflight_taskactivations ADD COLUMN heartbeat INTEGER;
+ALTER TABLE inflight_taskactivations ADD COLUMN heartbeat_deadline INTEGER;

src/canary_tasks.rs

@@ -101,6 +101,8 @@ fn build_activations(config: &Config, topic: &str) -> Result<Vec<Activation>, Er
                 delay_until: None,
                 processing_deadline: None,
                 claim_expires_at: None,
+                heartbeat: None,
+                heartbeat_deadline: None,
                 on_attempts_exceeded: OnAttemptsExceeded::Discard,
                 at_most_once: false,
                 bucket: bucket_from_id(&id),

src/fetch/tests.rs

@@ -189,6 +189,10 @@ impl ActivationStore for MockStore {
         unimplemented!()
     }
 
+    async fn update_heartbeats(&self, _ids: &[String]) -> Result<Vec<String>, Error> {
+        unimplemented!()
+    }
+
     async fn delete_activation(&self, _id: &str) -> Result<(), Error> {
         unimplemented!()
     }

src/grpc/server.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 use std::time::Instant;
 
@@ -7,9 +7,10 @@ use chrono::Utc;
 use prost::Message;
 use sentry_protos::taskbroker::v1::consumer_service_server::ConsumerService;
 use sentry_protos::taskbroker::v1::{
-    FetchNextTask, GetTaskRequest, GetTaskResponse, SetBatchActivationStatusRequest,
-    SetBatchActivationStatusResponse, SetTaskStatusRequest, SetTaskStatusResponse, TaskActivation,
-    TaskActivationStatus,
+    BatchUpdateHeartbeatRequest, BatchUpdateHeartbeatResponse, FetchNextTask, GetTaskRequest,
+    GetTaskResponse, SetBatchActivationStatusRequest, SetBatchActivationStatusResponse,
+    SetTaskStatusRequest, SetTaskStatusResponse, TaskActivation, TaskActivationStatus,
+    UpdateHeartbeatRequest, UpdateHeartbeatResponse,
 };
 use tokio::sync::mpsc::Sender;
 use tonic::{Request, Response, Status};
@@ -366,6 +367,82 @@ impl ConsumerService for TaskbrokerServer {
             .record(start_time.elapsed());
         return Ok(Response::new(SetBatchActivationStatusResponse {}));
     }
+
+    #[instrument(skip_all)]
+    async fn update_heartbeat(
+        &self,
+        request: Request<UpdateHeartbeatRequest>,
+    ) -> Result<Response<UpdateHeartbeatResponse>, Status> {
+        let start_time = Instant::now();
+        let id = &request.get_ref().id;
+
+        let result = self.store.update_heartbeats(std::slice::from_ref(id)).await;
+        metrics::histogram!("grpc_server.update_heartbeat.duration").record(start_time.elapsed());
+
+        match result {
+            Ok(updated) if !updated.is_empty() => {
+                metrics::counter!("grpc_server.update_heartbeat", "result" => "ok").increment(1);
+                Ok(Response::new(UpdateHeartbeatResponse {}))
+            }
+
+            Ok(_) => {
+                metrics::counter!("grpc_server.update_heartbeat", "result" => "not_found")
+                    .increment(1);
+                Err(Status::not_found(
+                    "No processing activation with a heartbeat",
+                ))
+            }
+
+            Err(e) => {
+                metrics::counter!("grpc_server.update_heartbeat", "result" => "error").increment(1);
+                error!(?id, "Unable to update heartbeat of activation: {:?}", e);
+                Err(Status::internal(format!(
+                    "Unable to update heartbeat of {id:?}"
+                )))
+            }
+        }
+    }
+
+    #[instrument(skip_all)]
+    async fn batch_update_heartbeat(
+        &self,
+        request: Request<BatchUpdateHeartbeatRequest>,
+    ) -> Result<Response<BatchUpdateHeartbeatResponse>, Status> {
+        let start_time = Instant::now();
+        let ids = request.into_inner().ids;
+        metrics::histogram!("grpc_server.batch_update_heartbeat.ids").record(ids.len() as f64);
+
+        let result = self.store.update_heartbeats(&ids).await;
+        metrics::histogram!("grpc_server.batch_update_heartbeat.duration")
+            .record(start_time.elapsed());
+
+        match result {
+            Ok(updated) => {
+                let updated: HashSet<String> = updated.into_iter().collect();
+                let not_found_ids: Vec<String> =
+                    ids.into_iter().filter(|id| !updated.contains(id)).collect();
+
+                metrics::counter!("grpc_server.batch_update_heartbeat", "result" => "ok")
+                    .increment(updated.len() as u64);
+                metrics::counter!("grpc_server.batch_update_heartbeat", "result" => "not_found")
+                    .increment(not_found_ids.len() as u64);
+                Ok(Response::new(BatchUpdateHeartbeatResponse {
+                    not_found_ids,
+                }))
+            }
+
+            Err(e) => {
+                metrics::counter!("grpc_server.batch_update_heartbeat", "result" => "error")
+                    .increment(1);
+                error!(
+                    "Unable to update heartbeats of {} activations: {:?}",
+                    ids.len(),
+                    e
+                );
+                Err(Status::internal("Unable to update heartbeats"))
+            }
+        }
+    }
 }
 
 pub type StatusUpdate = (String, ActivationStatus);

src/grpc/server_tests.rs

@@ -4,8 +4,8 @@ use prost::Message;
 use rstest::rstest;
 use sentry_protos::taskbroker::v1::consumer_service_server::ConsumerService;
 use sentry_protos::taskbroker::v1::{
-    FetchNextTask, GetTaskRequest, SetBatchActivationStatusRequest, SetTaskStatusRequest,
-    TaskActivation,
+    BatchUpdateHeartbeatRequest, FetchNextTask, GetTaskRequest, SetBatchActivationStatusRequest,
+    SetTaskStatusRequest, TaskActivation, UpdateHeartbeatRequest,
 };
 use tokio::sync::mpsc;
 use tonic::{Code, Request};
@@ -785,3 +785,52 @@ async fn test_set_batch_activation_status_mixed(#[case] adapter: &str) {
     let row2 = store.get_by_id("id_2").await.unwrap().expect("row exists");
     assert_eq!(row2.status, ActivationStatus::Retry);
 }
+
+#[tokio::test]
+#[rstest]
+#[case::sqlite("sqlite")]
+#[case::postgres("postgres")]
+async fn test_update_heartbeat(#[case] adapter: &str) {
+    let store = create_test_store(adapter).await;
+    let mut activations = make_activations(1);
+    activations[0].status = ActivationStatus::Processing;
+    activations[0].heartbeat = Some(10);
+    store.store(&activations).await.unwrap();
+
+    let service = TaskbrokerServer {
+        store,
+        config: create_config(),
+        update_tx: None,
+    };
+
+    let request = |id: &str| {
+        Request::new(UpdateHeartbeatRequest {
+            id: id.to_string(),
+            heartbeat_interval: 10,
+        })
+    };
+
+    assert!(
+        service
+            .update_heartbeat(request(&activations[0].id))
+            .await
+            .is_ok()
+    );
+
+    let e = service
+        .update_heartbeat(request("missing"))
+        .await
+        .unwrap_err();
+    assert_eq!(e.code(), Code::NotFound);
+
+    let response = service
+        .batch_update_heartbeat(Request::new(BatchUpdateHeartbeatRequest {
+            ids: vec![activations[0].id.clone(), "missing".to_string()],
+        }))
+        .await
+        .unwrap();
+    assert_eq!(
+        response.get_ref().not_found_ids,
+        vec!["missing".to_string()]
+    );
+}

src/kafka/deserialize_activation.rs

@@ -102,6 +102,12 @@ pub fn new(
             received_at: activation_time,
             processing_deadline: None,
             claim_expires_at: None,
+            // A zero heartbeat would make the task miss its deadline as soon as it starts.
+            heartbeat: activation
+                .heartbeat
+                .filter(|&h| h > 0)
+                .and_then(|h| i32::try_from(h).ok()),
+            heartbeat_deadline: None,
             processing_deadline_duration: activation.processing_deadline_duration as i32,
             processing_attempts: 0,
             expires_at,

src/kafka/deserialize_raw.rs

@@ -154,6 +154,7 @@ pub fn new(config: RawConfig) -> impl Fn(&MessageBackend) -> Result<Activation,
             processing_deadline_duration: config.processing_deadline_duration.into(),
             expires: None,
             delay: None,
+            heartbeat: None,
         };
 
         let activation_bytes = activation.encode_to_vec();
@@ -188,6 +189,8 @@ pub fn new(config: RawConfig) -> impl Fn(&MessageBackend) -> Result<Activation,
             received_at: received_at_time,
             processing_deadline: None,
             claim_expires_at: None,
+            heartbeat: None,
+            heartbeat_deadline: None,
             processing_deadline_duration: config.processing_deadline_duration.into(),
             processing_attempts: 0,
             expires_at: None,

src/push/tests.rs

@@ -122,6 +122,10 @@ impl ActivationStore for MockStore {
         Ok(())
     }
 
+    async fn update_heartbeats(&self, _ids: &[String]) -> Result<Vec<String>> {
+        Ok(vec![])
+    }
+
     async fn delete_activation(&self, _id: &str) -> Result<()> {
         Ok(())
     }

src/store/activation.rs

@@ -152,6 +152,14 @@ pub struct Activation {
     #[builder(default = None, setter(strip_option))]
     pub processing_deadline: Option<DateTime<Utc>>,
 
+    /// Seconds a worker has between heartbeats. None means the task does not heartbeat.
+    #[builder(default = None, setter(strip_option))]
+    pub heartbeat: Option<i32>,
+
+    /// If a heartbeat task is still processing after this time, upkeep retries it.
+    #[builder(default = None, setter(strip_option))]
+    pub heartbeat_deadline: Option<DateTime<Utc>>,
+
     /// If a task is still claimed after this time, upkeep may release the claim.
     #[builder(default = None, setter(strip_option))]
     pub claim_expires_at: Option<DateTime<Utc>>,