From 2c5bcbb0ac389ec53ddd0c36637ad8484f15767d Mon Sep 17 00:00:00 2001 From: Evan Hicks Date: Thu, 5 Nov 2026 10:42:01 -0500 Subject: [PATCH 1/7] feat(taskbroker): Retry activations that miss their heartbeat Activations with a heartbeat get a heartbeat_deadline when they start processing, and workers push it back with the UpdateHeartbeat RPC. Upkeep retries any activation past its heartbeat_deadline the same way it handles a missed processing deadline. Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 4 +- Cargo.toml | 2 +- migrations/postgres/0005_add_heartbeat.sql | 2 + migrations/sqlite/0007_add_heartbeat.sql | 2 + src/canary_tasks.rs | 2 + src/fetch/tests.rs | 4 + src/grpc/server.rs | 33 +++++- src/grpc/server_tests.rs | 40 ++++++- src/kafka/deserialize_activation.rs | 2 + src/kafka/deserialize_raw.rs | 3 + src/push/tests.rs | 4 + src/store/activation.rs | 8 ++ src/store/adapters/postgres.rs | 120 ++++++++++++++------- src/store/adapters/sqlite.rs | 116 +++++++++++++------- src/store/tests.rs | 88 +++++++++++++++ src/store/traits.rs | 4 + src/test_utils.rs | 9 ++ src/worker.rs | 1 + 18 files changed, 364 insertions(+), 80 deletions(-) create mode 100644 migrations/postgres/0005_add_heartbeat.sql create mode 100644 migrations/sqlite/0007_add_heartbeat.sql diff --git a/Cargo.lock b/Cargo.lock index 2a258356..a6947bc7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2996,9 +2996,9 @@ dependencies = [ [[package]] name = "sentry_protos" -version = "0.22.3" +version = "0.104.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6214388c55cb613be274cfdbb14cbb8ec47ac24ddbbf848ab08d4eefc793c00" +checksum = "568784007fd30c7309ac343976458bd76d037c70f3c4a4a37bf36591a0b7a306" dependencies = [ "prost", "prost-types", diff --git a/Cargo.toml b/Cargo.toml index 0470aef2..88097f79 100644 --- a/Cargo.toml +++ b/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.104.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. diff --git a/migrations/postgres/0005_add_heartbeat.sql b/migrations/postgres/0005_add_heartbeat.sql new file mode 100644 index 00000000..8a155a50 --- /dev/null +++ b/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; diff --git a/migrations/sqlite/0007_add_heartbeat.sql b/migrations/sqlite/0007_add_heartbeat.sql new file mode 100644 index 00000000..f6ad4c22 --- /dev/null +++ b/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; diff --git a/src/canary_tasks.rs b/src/canary_tasks.rs index bbe52df4..14fee5da 100644 --- a/src/canary_tasks.rs +++ b/src/canary_tasks.rs @@ -101,6 +101,8 @@ fn build_activations(config: &Config, topic: &str) -> Result, 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), diff --git a/src/fetch/tests.rs b/src/fetch/tests.rs index a079cf51..065270e3 100644 --- a/src/fetch/tests.rs +++ b/src/fetch/tests.rs @@ -189,6 +189,10 @@ impl ActivationStore for MockStore { unimplemented!() } + async fn update_heartbeat(&self, _id: &str) -> Result { + unimplemented!() + } + async fn delete_activation(&self, _id: &str) -> Result<(), Error> { unimplemented!() } diff --git a/src/grpc/server.rs b/src/grpc/server.rs index 1b39025f..16c76d12 100644 --- a/src/grpc/server.rs +++ b/src/grpc/server.rs @@ -9,7 +9,7 @@ use sentry_protos::taskbroker::v1::consumer_service_server::ConsumerService; use sentry_protos::taskbroker::v1::{ FetchNextTask, GetTaskRequest, GetTaskResponse, SetBatchActivationStatusRequest, SetBatchActivationStatusResponse, SetTaskStatusRequest, SetTaskStatusResponse, TaskActivation, - TaskActivationStatus, + TaskActivationStatus, UpdateHeartbeatRequest, UpdateHeartbeatResponse, }; use tokio::sync::mpsc::Sender; use tonic::{Request, Response, Status}; @@ -366,6 +366,37 @@ impl ConsumerService for TaskbrokerServer { .record(start_time.elapsed()); return Ok(Response::new(SetBatchActivationStatusResponse {})); } + + #[instrument(skip_all)] + async fn update_heartbeat( + &self, + request: Request, + ) -> Result, Status> { + let id = &request.get_ref().id; + + match self.store.update_heartbeat(id).await { + Ok(true) => { + metrics::counter!("grpc_server.update_heartbeat", "result" => "ok").increment(1); + Ok(Response::new(UpdateHeartbeatResponse {})) + } + + Ok(false) => { + 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:?}" + ))) + } + } + } } pub type StatusUpdate = (String, ActivationStatus); diff --git a/src/grpc/server_tests.rs b/src/grpc/server_tests.rs index 3a6bafda..c9813740 100644 --- a/src/grpc/server_tests.rs +++ b/src/grpc/server_tests.rs @@ -5,7 +5,7 @@ use rstest::rstest; use sentry_protos::taskbroker::v1::consumer_service_server::ConsumerService; use sentry_protos::taskbroker::v1::{ FetchNextTask, GetTaskRequest, SetBatchActivationStatusRequest, SetTaskStatusRequest, - TaskActivation, + TaskActivation, UpdateHeartbeatRequest, }; use tokio::sync::mpsc; use tonic::{Code, Request}; @@ -785,3 +785,41 @@ 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); +} diff --git a/src/kafka/deserialize_activation.rs b/src/kafka/deserialize_activation.rs index b373c31f..ff65b1c0 100644 --- a/src/kafka/deserialize_activation.rs +++ b/src/kafka/deserialize_activation.rs @@ -102,6 +102,8 @@ pub fn new( received_at: activation_time, processing_deadline: None, claim_expires_at: None, + heartbeat: activation.heartbeat.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, diff --git a/src/kafka/deserialize_raw.rs b/src/kafka/deserialize_raw.rs index 3808d0a8..728fbaab 100644 --- a/src/kafka/deserialize_raw.rs +++ b/src/kafka/deserialize_raw.rs @@ -154,6 +154,7 @@ pub fn new(config: RawConfig) -> impl Fn(&MessageBackend) -> Result impl Fn(&MessageBackend) -> Result Result { + Ok(false) + } + async fn delete_activation(&self, _id: &str) -> Result<()> { Ok(()) } diff --git a/src/store/activation.rs b/src/store/activation.rs index 2b27254c..c8a32c07 100644 --- a/src/store/activation.rs +++ b/src/store/activation.rs @@ -152,6 +152,14 @@ pub struct Activation { #[builder(default = None, setter(strip_option))] pub processing_deadline: Option>, + /// Seconds a worker has between heartbeats. None means the task does not heartbeat. + #[builder(default = None, setter(strip_option))] + pub heartbeat: Option, + + /// If a heartbeat task is still processing after this time, upkeep retries it. + #[builder(default = None, setter(strip_option))] + pub heartbeat_deadline: Option>, + /// 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>, diff --git a/src/store/adapters/postgres.rs b/src/store/adapters/postgres.rs index 6ad81ee8..75f2d5df 100644 --- a/src/store/adapters/postgres.rs +++ b/src/store/adapters/postgres.rs @@ -112,6 +112,8 @@ struct TableRow<'a> { pub processing_deadline_duration: i32, pub processing_deadline: Option>, pub claim_expires_at: Option>, + pub heartbeat: Option, + pub heartbeat_deadline: Option>, pub status: Cow<'a, str>, pub at_most_once: bool, pub application: Cow<'a, str>, @@ -137,6 +139,8 @@ impl<'a> From<&'a Activation> for TableRow<'a> { processing_deadline_duration: value.processing_deadline_duration, processing_deadline: value.processing_deadline, claim_expires_at: value.claim_expires_at, + heartbeat: value.heartbeat, + heartbeat_deadline: value.heartbeat_deadline, status: Cow::Owned(value.status.to_string()), at_most_once: value.at_most_once, application: Cow::Borrowed(&value.application), @@ -167,6 +171,8 @@ impl From> for Activation { delay_until: value.delay_until, processing_deadline: value.processing_deadline, claim_expires_at: value.claim_expires_at, + heartbeat: value.heartbeat, + heartbeat_deadline: value.heartbeat_deadline, at_most_once: value.at_most_once, application: value.application.into_owned(), namespace: value.namespace.into_owned(), @@ -196,6 +202,8 @@ impl FromRow<'_, PgRow> for TableRow<'static> { processing_deadline_duration: row.try_get("processing_deadline_duration")?, processing_deadline: row.try_get("processing_deadline")?, claim_expires_at: row.try_get("claim_expires_at")?, + heartbeat: row.try_get("heartbeat")?, + heartbeat_deadline: row.try_get("heartbeat_deadline")?, status: Cow::Owned(row.try_get::("status")?), at_most_once: row.try_get("at_most_once")?, application: Cow::Owned(row.try_get::("application")?), @@ -404,6 +412,8 @@ impl ActivationStore for PostgresStore { processing_deadline_duration, processing_deadline, claim_expires_at, + heartbeat, + heartbeat_deadline, status, at_most_once, application, @@ -468,6 +478,8 @@ impl ActivationStore for PostgresStore { processing_deadline_duration, processing_deadline, claim_expires_at, + heartbeat, + heartbeat_deadline, status, at_most_once, application, @@ -505,6 +517,8 @@ impl ActivationStore for PostgresStore { } else { b.push("null"); } + b.push_bind(row.heartbeat); + b.push_bind(row.heartbeat_deadline); b.push_bind(row.status); b.push_bind(row.at_most_once); b.push_bind(row.application); @@ -588,6 +602,7 @@ impl ActivationStore for PostgresStore { query_builder.push(format!( "UPDATE inflight_taskactivations SET processing_deadline = now() + (processing_deadline_duration * interval '1 second') + (interval '{grace_period} seconds'), + heartbeat_deadline = now() + (heartbeat * interval '1 second'), claim_expires_at = NULL, status = " )); @@ -641,6 +656,7 @@ impl ActivationStore for PostgresStore { UPDATE inflight_taskactivations SET status = $1, processing_deadline = now() + (processing_deadline_duration * interval '1 second') + (interval '{grace_period} seconds'), + heartbeat_deadline = now() + (heartbeat * interval '1 second'), claim_expires_at = NULL FROM selected WHERE inflight_taskactivations.id = selected.id", @@ -695,6 +711,7 @@ impl ActivationStore for PostgresStore { UPDATE inflight_taskactivations SET status = $1, processing_deadline = now() + (processing_deadline_duration * interval '1 second') + (interval '{grace_period} seconds'), + heartbeat_deadline = now() + (heartbeat * interval '1 second'), claim_expires_at = NULL FROM selected WHERE inflight_taskactivations.id = selected.id", @@ -1056,6 +1073,25 @@ impl ActivationStore for PostgresStore { .await } + #[instrument(skip_all)] + #[framed] + async fn update_heartbeat(&self, id: &str) -> Result { + retry_query(&self.config.retry, "update_heartbeat", || async { + let mut conn = self.acquire_write_conn_metric("update_heartbeat").await?; + let result = sqlx::query( + "UPDATE inflight_taskactivations + SET heartbeat_deadline = now() + (heartbeat * interval '1 second') + WHERE id = $1 AND status = $2 AND heartbeat IS NOT NULL", + ) + .bind(id) + .bind(ActivationStatus::Processing.to_string()) + .execute(&mut *conn) + .await?; + Ok(result.rows_affected() > 0) + }) + .await + } + #[instrument(skip_all)] #[framed] async fn delete_activation(&self, id: &str) -> Result<(), Error> { @@ -1104,6 +1140,8 @@ impl ActivationStore for PostgresStore { processing_deadline_duration, processing_deadline, claim_expires_at, + heartbeat, + heartbeat_deadline, status, at_most_once, application, @@ -1175,52 +1213,58 @@ impl ActivationStore for PostgresStore { async fn handle_processing_deadline(&self) -> Result { let now = Utc::now(); let mut atomic = self.write_pool.begin().await?; + let mut modified_rows = 0; + + for (deadline, reason) in [ + ("processing_deadline", "processing_deadline"), + ("heartbeat_deadline", "heartbeat"), + ] { + // At-most-once tasks that fail their deadlines go directly to failure + // there are no retries, as the worker will reject the task due to at_most_once keys. + let mut query_builder = QueryBuilder::new( + "UPDATE inflight_taskactivations + SET processing_deadline = null, heartbeat_deadline = null, status = ", + ); + query_builder.push_bind(ActivationStatus::Failure.to_string()); + query_builder.push(format!(" WHERE {deadline} < ")); + query_builder.push_bind(now); + query_builder.push(" AND at_most_once = TRUE AND status = "); + query_builder.push_bind(ActivationStatus::Processing.to_string()); - // At-most-once tasks that fail their processing deadlines go directly to failure - // there are no retries, as the worker will reject the task due to at_most_once keys. - let mut query_builder = QueryBuilder::new( - "UPDATE inflight_taskactivations - SET processing_deadline = null, status = ", - ); - query_builder.push_bind(ActivationStatus::Failure.to_string()); - query_builder.push(" WHERE processing_deadline < "); - query_builder.push_bind(now); - query_builder.push(" AND at_most_once = TRUE AND status = "); - query_builder.push_bind(ActivationStatus::Processing.to_string()); - - self.add_partition_condition(&mut query_builder, false); + self.add_partition_condition(&mut query_builder, false); - let most_once_result = query_builder.build().execute(&mut *atomic).await; + let mut reset = 0; + if let Ok(query_res) = query_builder.build().execute(&mut *atomic).await { + reset = query_res.rows_affected(); + } - let mut processing_deadline_modified_rows = 0; - if let Ok(query_res) = most_once_result { - processing_deadline_modified_rows = query_res.rows_affected(); - } + // Update regular tasks. + // Increment processing_attempts by 1 and reset deadlines to null. + let mut query_builder = QueryBuilder::new( + "UPDATE inflight_taskactivations + SET processing_deadline = null, heartbeat_deadline = null, status = ", + ); + query_builder.push_bind(ActivationStatus::Pending.to_string()); + query_builder.push(", processing_attempts = processing_attempts + 1"); + query_builder.push(format!(" WHERE {deadline} < ")); + query_builder.push_bind(now); + query_builder.push(" AND status = "); + query_builder.push_bind(ActivationStatus::Processing.to_string()); + self.add_partition_condition(&mut query_builder, false); - // Update regular tasks. - // Increment processing_attempts by 1 and reset processing_deadline to null. - let mut query_builder = QueryBuilder::new( - "UPDATE inflight_taskactivations - SET processing_deadline = null, status = ", - ); - query_builder.push_bind(ActivationStatus::Pending.to_string()); - query_builder.push(", processing_attempts = processing_attempts + 1"); - query_builder.push(" WHERE processing_deadline < "); - query_builder.push_bind(now); - query_builder.push(" AND status = "); - query_builder.push_bind(ActivationStatus::Processing.to_string()); - self.add_partition_condition(&mut query_builder, false); + let Ok(query_res) = query_builder.build().execute(&mut *atomic).await else { + return Err(anyhow!("Could not update tasks past {deadline}")); + }; + reset += query_res.rows_affected(); - let result = query_builder.build().execute(&mut *atomic).await; + metrics::counter!("upkeep.handle_processing_deadline.reset", "reason" => reason) + .increment(reset); + modified_rows += reset; + } atomic.commit().await?; - if let Ok(query_res) = result { - processing_deadline_modified_rows += query_res.rows_affected(); - return Ok(processing_deadline_modified_rows); - } - - Err(anyhow!("Could not update tasks past processing_deadline")) + Ok(modified_rows) } /// Update tasks that have exceeded their max processing attempts. diff --git a/src/store/adapters/sqlite.rs b/src/store/adapters/sqlite.rs index 9dd19656..f0f197b8 100644 --- a/src/store/adapters/sqlite.rs +++ b/src/store/adapters/sqlite.rs @@ -54,6 +54,8 @@ pub struct TableRow<'a> { pub processing_deadline_duration: i32, pub processing_deadline: Option>, pub claim_expires_at: Option>, + pub heartbeat: Option, + pub heartbeat_deadline: Option>, pub status: Cow<'a, str>, pub at_most_once: bool, pub application: Cow<'a, str>, @@ -78,6 +80,8 @@ impl<'a> From<&'a Activation> for TableRow<'a> { processing_deadline_duration: value.processing_deadline_duration, processing_deadline: value.processing_deadline, claim_expires_at: value.claim_expires_at, + heartbeat: value.heartbeat, + heartbeat_deadline: value.heartbeat_deadline, status: Cow::Owned(value.status.to_string()), at_most_once: value.at_most_once, application: Cow::Borrowed(&value.application), @@ -112,6 +116,8 @@ impl From> for Activation { delay_until: value.delay_until, processing_deadline: value.processing_deadline, claim_expires_at: value.claim_expires_at, + heartbeat: value.heartbeat, + heartbeat_deadline: value.heartbeat_deadline, at_most_once: value.at_most_once, application: value.application.into_owned(), namespace: value.namespace.into_owned(), @@ -140,6 +146,8 @@ impl FromRow<'_, SqliteRow> for TableRow<'static> { processing_deadline_duration: row.try_get("processing_deadline_duration")?, processing_deadline: row.try_get("processing_deadline")?, claim_expires_at: row.try_get("claim_expires_at")?, + heartbeat: row.try_get("heartbeat")?, + heartbeat_deadline: row.try_get("heartbeat_deadline")?, status: Cow::Owned(row.try_get::("status")?), at_most_once: row.try_get("at_most_once")?, application: Cow::Owned(row.try_get::("application")?), @@ -444,6 +452,8 @@ impl ActivationStore for SqliteStore { processing_deadline_duration, processing_deadline, claim_expires_at, + heartbeat, + heartbeat_deadline, status, at_most_once, application, @@ -504,6 +514,8 @@ impl ActivationStore for SqliteStore { processing_deadline_duration, processing_deadline, claim_expires_at, + heartbeat, + heartbeat_deadline, status, at_most_once, application, @@ -543,6 +555,9 @@ impl ActivationStore for SqliteStore { b.push("null"); } + b.push_bind(row.heartbeat); + b.push_bind(row.heartbeat_deadline.map(|t| t.timestamp())); + b.push_bind(row.status); b.push_bind(row.at_most_once); b.push_bind(row.application); @@ -595,7 +610,7 @@ impl ActivationStore for SqliteStore { if mark_processing { query_builder.push(format!( - "processing_deadline = unixepoch('now', '+' || (processing_deadline_duration + {grace_period}) || ' seconds'), claim_expires_at = NULL, status = " + "processing_deadline = unixepoch('now', '+' || (processing_deadline_duration + {grace_period}) || ' seconds'), heartbeat_deadline = unixepoch('now', '+' || heartbeat || ' seconds'), claim_expires_at = NULL, status = " )); query_builder.push_bind(ActivationStatus::Processing); @@ -660,7 +675,7 @@ impl ActivationStore for SqliteStore { let result = sqlx::query(&format!( "UPDATE inflight_taskactivations SET status = $1, - processing_deadline = unixepoch('now', '+' || (processing_deadline_duration + {grace_period}) || ' seconds'), + processing_deadline = unixepoch('now', '+' || (processing_deadline_duration + {grace_period}) || ' seconds'), heartbeat_deadline = unixepoch('now', '+' || heartbeat || ' seconds'), claim_expires_at = NULL WHERE id = $2 AND status = $3", )) @@ -699,7 +714,7 @@ impl ActivationStore for SqliteStore { let mut query_builder = QueryBuilder::new("UPDATE inflight_taskactivations SET status = "); query_builder.push_bind(ActivationStatus::Processing); query_builder.push(format!( - ", processing_deadline = unixepoch('now', '+' || (processing_deadline_duration + {grace_period}) || ' seconds'), claim_expires_at = NULL WHERE status = ", + ", processing_deadline = unixepoch('now', '+' || (processing_deadline_duration + {grace_period}) || ' seconds'), heartbeat_deadline = unixepoch('now', '+' || heartbeat || ' seconds'), claim_expires_at = NULL WHERE status = ", )); query_builder.push_bind(ActivationStatus::Claimed); query_builder.push(" AND id IN ("); @@ -920,6 +935,21 @@ impl ActivationStore for SqliteStore { Ok(()) } + #[instrument(skip_all)] + async fn update_heartbeat(&self, id: &str) -> Result { + let mut conn = self.acquire_write_conn_metric("update_heartbeat").await?; + let result = sqlx::query( + "UPDATE inflight_taskactivations + SET heartbeat_deadline = unixepoch('now', '+' || heartbeat || ' seconds') + WHERE id = $1 AND status = $2 AND heartbeat IS NOT NULL", + ) + .bind(id) + .bind(ActivationStatus::Processing) + .execute(&mut *conn) + .await?; + Ok(result.rows_affected() > 0) + } + #[instrument(skip_all)] async fn delete_activation(&self, id: &str) -> Result<(), Error> { let mut conn = self.acquire_write_conn_metric("delete_activation").await?; @@ -967,6 +997,8 @@ impl ActivationStore for SqliteStore { processing_deadline_duration, processing_deadline, claim_expires_at, + heartbeat, + heartbeat_deadline, status, at_most_once, application, @@ -1026,46 +1058,56 @@ impl ActivationStore for SqliteStore { async fn handle_processing_deadline(&self) -> Result { let now = Utc::now(); let mut atomic = self.write_pool.begin().await?; + let mut modified_rows = 0; + + for (deadline, reason) in [ + ("processing_deadline", "processing_deadline"), + ("heartbeat_deadline", "heartbeat"), + ] { + // Idempotent tasks that fail their deadlines go directly to failure + // there are no retries, as the worker will reject the task due to idempotency keys. + let most_once_result = sqlx::query(&format!( + "UPDATE inflight_taskactivations + SET processing_deadline = null, heartbeat_deadline = null, status = $1 + WHERE {deadline} < $2 AND at_most_once = TRUE AND status = $3", + )) + .bind(ActivationStatus::Failure) + .bind(now.timestamp()) + .bind(ActivationStatus::Processing) + .execute(&mut *atomic) + .await; - // Idempotent tasks that fail their processing deadlines go directly to failure - // there are no retries, as the worker will reject the task due to idempotency keys. - let most_once_result = sqlx::query( - "UPDATE inflight_taskactivations - SET processing_deadline = null, status = $1 - WHERE processing_deadline < $2 AND at_most_once = TRUE AND status = $3", - ) - .bind(ActivationStatus::Failure) - .bind(now.timestamp()) - .bind(ActivationStatus::Processing) - .execute(&mut *atomic) - .await; - - let mut processing_deadline_modified_rows = 0; - if let Ok(query_res) = most_once_result { - processing_deadline_modified_rows = query_res.rows_affected(); - } + let mut reset = 0; + if let Ok(query_res) = most_once_result { + reset = query_res.rows_affected(); + } - // Update non-idempotent tasks. - // Increment processing_attempts by 1 and reset processing_deadline to null. - let result = sqlx::query( - "UPDATE inflight_taskactivations - SET processing_deadline = null, status = $1, processing_attempts = processing_attempts + 1 - WHERE processing_deadline < $2 AND status = $3", - ) - .bind(ActivationStatus::Pending) - .bind(now.timestamp()) - .bind(ActivationStatus::Processing) - .execute(&mut *atomic) - .await; + // Update non-idempotent tasks. + // Increment processing_attempts by 1 and reset deadlines to null. + let result = sqlx::query(&format!( + "UPDATE inflight_taskactivations + SET processing_deadline = null, heartbeat_deadline = null, status = $1, processing_attempts = processing_attempts + 1 + WHERE {deadline} < $2 AND status = $3", + )) + .bind(ActivationStatus::Pending) + .bind(now.timestamp()) + .bind(ActivationStatus::Processing) + .execute(&mut *atomic) + .await; - atomic.commit().await?; + let Ok(query_res) = result else { + return Err(anyhow!("Could not update tasks past {deadline}")); + }; + reset += query_res.rows_affected(); - if let Ok(query_res) = result { - processing_deadline_modified_rows += query_res.rows_affected(); - return Ok(processing_deadline_modified_rows); + metrics::counter!("upkeep.handle_processing_deadline.reset", "reason" => reason) + .increment(reset); + modified_rows += reset; } - Err(anyhow!("Could not update tasks past processing_deadline")) + atomic.commit().await?; + + Ok(modified_rows) } /// Update tasks that have exceeded their max processing attempts. diff --git a/src/store/tests.rs b/src/store/tests.rs index 8d2e1598..4044dd8a 100644 --- a/src/store/tests.rs +++ b/src/store/tests.rs @@ -1330,6 +1330,94 @@ async fn test_handle_processing_deadline(#[case] adapter: &str) { store.remove_db().await.unwrap(); } +#[tokio::test] +#[rstest] +#[case::sqlite("sqlite")] +#[case::postgres("postgres")] +async fn test_handle_processing_deadline_heartbeat(#[case] adapter: &str) { + let store = create_test_store(adapter).await; + let past = Utc.with_ymd_and_hms(2024, 11, 14, 21, 22, 23).unwrap(); + let future = Utc::now() + Duration::from_secs(600); + + let mut batch = make_activations(2); + for activation in batch.iter_mut() { + activation.status = ActivationStatus::Processing; + activation.processing_deadline = Some(future); + activation.heartbeat = Some(10); + } + batch[0].heartbeat_deadline = Some(past); + batch[1].heartbeat_deadline = Some(future); + assert!(store.store(&batch).await.is_ok()); + + assert_eq!(store.handle_processing_deadline().await.unwrap(), 1); + + let task = store.get_by_id(&batch[0].id).await.unwrap().unwrap(); + assert_eq!(task.status, ActivationStatus::Pending); + assert_eq!(task.processing_attempts, 1); + assert_eq!(task.heartbeat_deadline, None); + assert_eq!(task.processing_deadline, None); + + let task = store.get_by_id(&batch[1].id).await.unwrap().unwrap(); + assert_eq!(task.status, ActivationStatus::Processing); + store.remove_db().await.unwrap(); +} + +#[tokio::test] +#[rstest] +#[case::sqlite("sqlite")] +#[case::postgres("postgres")] +async fn test_heartbeat_deadline_set_and_updated(#[case] adapter: &str) { + let store = create_test_store(adapter).await; + + let mut batch = make_activations(3); + batch[0].heartbeat = Some(10); + batch[1].heartbeat = Some(10); + + // Pull mode sets the heartbeat deadline when claiming + assert!(store.store(&batch[0..1]).await.is_ok()); + assert!( + store + .claim_activation_for_pull(None, None) + .await + .unwrap() + .is_some() + ); + + // Push mode sets it when marking processing + assert!(store.store(&batch[1..]).await.is_ok()); + assert_eq!( + store + .claim_activations_for_push(None, None) + .await + .unwrap() + .len(), + 2 + ); + store + .mark_activation_processing(&batch[1].id) + .await + .unwrap(); + store + .mark_processing_batch(&[batch[2].id.clone()]) + .await + .unwrap(); + + let now = Utc::now(); + for activation in &batch[0..2] { + let task = store.get_by_id(&activation.id).await.unwrap().unwrap(); + assert_eq!(task.status, ActivationStatus::Processing); + assert!(task.heartbeat_deadline.unwrap() > now); + assert!(store.update_heartbeat(&activation.id).await.unwrap()); + } + + let task = store.get_by_id(&batch[2].id).await.unwrap().unwrap(); + assert_eq!(task.status, ActivationStatus::Processing); + assert_eq!(task.heartbeat_deadline, None); + assert!(!store.update_heartbeat(&batch[2].id).await.unwrap()); + assert!(!store.update_heartbeat("missing").await.unwrap()); + store.remove_db().await.unwrap(); +} + #[tokio::test] #[rstest] #[case::sqlite("sqlite")] diff --git a/src/store/traits.rs b/src/store/traits.rs index f70bb47c..34955b24 100644 --- a/src/store/traits.rs +++ b/src/store/traits.rs @@ -209,6 +209,10 @@ pub trait ActivationStore: Send + Sync { deadline: Option>, ) -> Result<(), Error>; + /// Push back the heartbeat deadline of a processing activation. Returns false + /// if no processing activation with a heartbeat has this id. + async fn update_heartbeat(&self, id: &str) -> Result; + /// Delete an activation by id async fn delete_activation(&self, id: &str) -> Result<(), Error>; diff --git a/src/test_utils.rs b/src/test_utils.rs index 7fc6bcc5..f24b5b2d 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -40,6 +40,7 @@ pub struct TaskActivationBuilder { pub processing_deadline_duration: Option, pub expires: Option, pub delay: Option, + pub heartbeat: Option, } impl TaskActivationBuilder { @@ -56,6 +57,7 @@ impl TaskActivationBuilder { processing_deadline_duration: None, expires: None, delay: None, + heartbeat: None, } } @@ -114,6 +116,11 @@ impl TaskActivationBuilder { self } + pub fn heartbeat(mut self, heartbeat: u64) -> Self { + self.heartbeat = Some(heartbeat); + self + } + #[allow(deprecated)] pub fn build(self) -> v1::TaskActivation { v1::TaskActivation { @@ -132,6 +139,7 @@ impl TaskActivationBuilder { retry_state: self.retry_state, expires: self.expires, delay: self.delay, + heartbeat: self.heartbeat, } } } @@ -191,6 +199,7 @@ impl ActivationBuilder { // Set 'expiration' and 'delay' fields manually after activation has been build activation.expires = expires; activation.delay = delay; + activation.heartbeat = self.heartbeat.flatten().map(|h| h as u64); self.activation = Some(activation.encode_to_vec()); diff --git a/src/worker.rs b/src/worker.rs index 3495d23d..cf4ae15c 100644 --- a/src/worker.rs +++ b/src/worker.rs @@ -86,6 +86,7 @@ impl WorkerClient for Worker { let request = PushTaskRequest { task: Some(task), + #[allow(deprecated)] callback_url: "".into(), // Workers don't use this anymore, so use an empty string until field is removed from schema }; From e25bcae94b3d3dfc9b893ce6b26921550cfa4744 Mon Sep 17 00:00:00 2001 From: Evan Hicks Date: Thu, 5 Nov 2026 16:02:39 -0500 Subject: [PATCH 2/7] feat(taskbroker): Add BatchUpdateHeartbeat handler Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 3 +- Cargo.toml | 2 +- src/fetch/tests.rs | 2 +- src/grpc/server.rs | 55 +++++++++++++++++++++++++++++----- src/grpc/server_tests.rs | 18 +++++++++-- src/push/tests.rs | 4 +-- src/store/adapters/postgres.rs | 18 +++++------ src/store/adapters/sqlite.rs | 27 ++++++++++------- src/store/tests.rs | 8 +++-- src/store/traits.rs | 6 ++-- 10 files changed, 103 insertions(+), 40 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a6947bc7..531b0850 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2997,8 +2997,7 @@ dependencies = [ [[package]] name = "sentry_protos" version = "0.104.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "568784007fd30c7309ac343976458bd76d037c70f3c4a4a37bf36591a0b7a306" +source = "git+https://github.com/getsentry/sentry-protos?rev=6a1fdb404953404abfce6a4d0a0d900a44c8a253#6a1fdb404953404abfce6a4d0a0d900a44c8a253" dependencies = [ "prost", "prost-types", diff --git a/Cargo.toml b/Cargo.toml index 88097f79..49a01d80 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -50,7 +50,7 @@ sentry = { version = "0.41.0", default-features = false, features = [ "tracing", "logs", ] } -sentry_protos = "0.104.0" +sentry_protos = { git = "https://github.com/getsentry/sentry-protos", rev = "6a1fdb404953404abfce6a4d0a0d900a44c8a253" } # 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. diff --git a/src/fetch/tests.rs b/src/fetch/tests.rs index 065270e3..b41fdd8c 100644 --- a/src/fetch/tests.rs +++ b/src/fetch/tests.rs @@ -189,7 +189,7 @@ impl ActivationStore for MockStore { unimplemented!() } - async fn update_heartbeat(&self, _id: &str) -> Result { + async fn update_heartbeats(&self, _ids: &[String]) -> Result, Error> { unimplemented!() } diff --git a/src/grpc/server.rs b/src/grpc/server.rs index 16c76d12..6cc4581c 100644 --- a/src/grpc/server.rs +++ b/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, UpdateHeartbeatRequest, UpdateHeartbeatResponse, + BatchUpdateHeartbeatRequest, BatchUpdateHeartbeatResponse, FetchNextTask, GetTaskRequest, + GetTaskResponse, SetBatchActivationStatusRequest, SetBatchActivationStatusResponse, + SetTaskStatusRequest, SetTaskStatusResponse, TaskActivation, TaskActivationStatus, + UpdateHeartbeatRequest, UpdateHeartbeatResponse, }; use tokio::sync::mpsc::Sender; use tonic::{Request, Response, Status}; @@ -374,13 +375,13 @@ impl ConsumerService for TaskbrokerServer { ) -> Result, Status> { let id = &request.get_ref().id; - match self.store.update_heartbeat(id).await { - Ok(true) => { + match self.store.update_heartbeats(std::slice::from_ref(id)).await { + Ok(updated) if !updated.is_empty() => { metrics::counter!("grpc_server.update_heartbeat", "result" => "ok").increment(1); Ok(Response::new(UpdateHeartbeatResponse {})) } - Ok(false) => { + Ok(_) => { metrics::counter!("grpc_server.update_heartbeat", "result" => "not_found") .increment(1); Err(Status::not_found( @@ -397,6 +398,46 @@ impl ConsumerService for TaskbrokerServer { } } } + + #[instrument(skip_all)] + async fn batch_update_heartbeat( + &self, + request: Request, + ) -> Result, Status> { + let ids: Vec = request + .into_inner() + .updates + .into_iter() + .map(|u| u.id) + .collect(); + + match self.store.update_heartbeats(&ids).await { + Ok(updated) => { + let updated: HashSet = updated.into_iter().collect(); + let not_found_ids: Vec = + 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); diff --git a/src/grpc/server_tests.rs b/src/grpc/server_tests.rs index c9813740..7e48943d 100644 --- a/src/grpc/server_tests.rs +++ b/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, UpdateHeartbeatRequest, + BatchUpdateHeartbeatRequest, FetchNextTask, GetTaskRequest, SetBatchActivationStatusRequest, + SetTaskStatusRequest, TaskActivation, UpdateHeartbeatRequest, }; use tokio::sync::mpsc; use tonic::{Code, Request}; @@ -822,4 +822,18 @@ async fn test_update_heartbeat(#[case] adapter: &str) { .await .unwrap_err(); assert_eq!(e.code(), Code::NotFound); + + let response = service + .batch_update_heartbeat(Request::new(BatchUpdateHeartbeatRequest { + updates: vec![ + request(&activations[0].id).into_inner(), + request("missing").into_inner(), + ], + })) + .await + .unwrap(); + assert_eq!( + response.get_ref().not_found_ids, + vec!["missing".to_string()] + ); } diff --git a/src/push/tests.rs b/src/push/tests.rs index 35c06dda..5e9e6c57 100644 --- a/src/push/tests.rs +++ b/src/push/tests.rs @@ -122,8 +122,8 @@ impl ActivationStore for MockStore { Ok(()) } - async fn update_heartbeat(&self, _id: &str) -> Result { - Ok(false) + async fn update_heartbeats(&self, _ids: &[String]) -> Result> { + Ok(vec![]) } async fn delete_activation(&self, _id: &str) -> Result<()> { diff --git a/src/store/adapters/postgres.rs b/src/store/adapters/postgres.rs index 75f2d5df..63695c53 100644 --- a/src/store/adapters/postgres.rs +++ b/src/store/adapters/postgres.rs @@ -1075,19 +1075,19 @@ impl ActivationStore for PostgresStore { #[instrument(skip_all)] #[framed] - async fn update_heartbeat(&self, id: &str) -> Result { - retry_query(&self.config.retry, "update_heartbeat", || async { - let mut conn = self.acquire_write_conn_metric("update_heartbeat").await?; - let result = sqlx::query( + async fn update_heartbeats(&self, ids: &[String]) -> Result, Error> { + retry_query(&self.config.retry, "update_heartbeats", || async { + let mut conn = self.acquire_write_conn_metric("update_heartbeats").await?; + Ok(sqlx::query_scalar( "UPDATE inflight_taskactivations SET heartbeat_deadline = now() + (heartbeat * interval '1 second') - WHERE id = $1 AND status = $2 AND heartbeat IS NOT NULL", + WHERE id = ANY($1) AND status = $2 AND heartbeat IS NOT NULL + RETURNING id", ) - .bind(id) + .bind(ids) .bind(ActivationStatus::Processing.to_string()) - .execute(&mut *conn) - .await?; - Ok(result.rows_affected() > 0) + .fetch_all(&mut *conn) + .await?) }) .await } diff --git a/src/store/adapters/sqlite.rs b/src/store/adapters/sqlite.rs index f0f197b8..9eec0672 100644 --- a/src/store/adapters/sqlite.rs +++ b/src/store/adapters/sqlite.rs @@ -936,18 +936,25 @@ impl ActivationStore for SqliteStore { } #[instrument(skip_all)] - async fn update_heartbeat(&self, id: &str) -> Result { - let mut conn = self.acquire_write_conn_metric("update_heartbeat").await?; - let result = sqlx::query( + async fn update_heartbeats(&self, ids: &[String]) -> Result, Error> { + let mut conn = self.acquire_write_conn_metric("update_heartbeats").await?; + let mut query_builder = QueryBuilder::new( "UPDATE inflight_taskactivations SET heartbeat_deadline = unixepoch('now', '+' || heartbeat || ' seconds') - WHERE id = $1 AND status = $2 AND heartbeat IS NOT NULL", - ) - .bind(id) - .bind(ActivationStatus::Processing) - .execute(&mut *conn) - .await?; - Ok(result.rows_affected() > 0) + WHERE heartbeat IS NOT NULL AND status = ", + ); + query_builder.push_bind(ActivationStatus::Processing); + query_builder.push(" AND id IN ("); + let mut separated = query_builder.separated(", "); + for id in ids { + separated.push_bind(id); + } + separated.push_unseparated(") RETURNING id"); + + Ok(query_builder + .build_query_scalar() + .fetch_all(&mut *conn) + .await?) } #[instrument(skip_all)] diff --git a/src/store/tests.rs b/src/store/tests.rs index 4044dd8a..e662d4d6 100644 --- a/src/store/tests.rs +++ b/src/store/tests.rs @@ -1407,14 +1407,16 @@ async fn test_heartbeat_deadline_set_and_updated(#[case] adapter: &str) { let task = store.get_by_id(&activation.id).await.unwrap().unwrap(); assert_eq!(task.status, ActivationStatus::Processing); assert!(task.heartbeat_deadline.unwrap() > now); - assert!(store.update_heartbeat(&activation.id).await.unwrap()); } let task = store.get_by_id(&batch[2].id).await.unwrap().unwrap(); assert_eq!(task.status, ActivationStatus::Processing); assert_eq!(task.heartbeat_deadline, None); - assert!(!store.update_heartbeat(&batch[2].id).await.unwrap()); - assert!(!store.update_heartbeat("missing").await.unwrap()); + + let ids = [&batch[0].id, &batch[1].id, &batch[2].id, "missing"].map(String::from); + let mut updated = store.update_heartbeats(&ids).await.unwrap(); + updated.sort(); + assert_eq!(updated, ids[0..2]); store.remove_db().await.unwrap(); } diff --git a/src/store/traits.rs b/src/store/traits.rs index 34955b24..f33d2d45 100644 --- a/src/store/traits.rs +++ b/src/store/traits.rs @@ -209,9 +209,9 @@ pub trait ActivationStore: Send + Sync { deadline: Option>, ) -> Result<(), Error>; - /// Push back the heartbeat deadline of a processing activation. Returns false - /// if no processing activation with a heartbeat has this id. - async fn update_heartbeat(&self, id: &str) -> Result; + /// Push back the heartbeat deadline of processing activations. Returns the ids + /// that were updated; ids that are not processing or have no heartbeat are left out. + async fn update_heartbeats(&self, ids: &[String]) -> Result, Error>; /// Delete an activation by id async fn delete_activation(&self, id: &str) -> Result<(), Error>; From 2a60241f38621e0d9185cc79f7352107e5317608 Mon Sep 17 00:00:00 2001 From: Evan Hicks Date: Thu, 5 Nov 2026 16:09:13 -0500 Subject: [PATCH 3/7] ref(taskbroker): Read ids from BatchUpdateHeartbeat request Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 2 +- Cargo.toml | 2 +- src/grpc/server.rs | 7 +------ src/grpc/server_tests.rs | 5 +---- 4 files changed, 4 insertions(+), 12 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 531b0850..94c575d1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2997,7 +2997,7 @@ dependencies = [ [[package]] name = "sentry_protos" version = "0.104.0" -source = "git+https://github.com/getsentry/sentry-protos?rev=6a1fdb404953404abfce6a4d0a0d900a44c8a253#6a1fdb404953404abfce6a4d0a0d900a44c8a253" +source = "git+https://github.com/getsentry/sentry-protos?rev=0383e4fc75e5eeccd5803de9cca509653681de5c#0383e4fc75e5eeccd5803de9cca509653681de5c" dependencies = [ "prost", "prost-types", diff --git a/Cargo.toml b/Cargo.toml index 49a01d80..c20ee357 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -50,7 +50,7 @@ sentry = { version = "0.41.0", default-features = false, features = [ "tracing", "logs", ] } -sentry_protos = { git = "https://github.com/getsentry/sentry-protos", rev = "6a1fdb404953404abfce6a4d0a0d900a44c8a253" } +sentry_protos = { git = "https://github.com/getsentry/sentry-protos", rev = "0383e4fc75e5eeccd5803de9cca509653681de5c" } # 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. diff --git a/src/grpc/server.rs b/src/grpc/server.rs index 6cc4581c..73343f30 100644 --- a/src/grpc/server.rs +++ b/src/grpc/server.rs @@ -404,12 +404,7 @@ impl ConsumerService for TaskbrokerServer { &self, request: Request, ) -> Result, Status> { - let ids: Vec = request - .into_inner() - .updates - .into_iter() - .map(|u| u.id) - .collect(); + let ids = request.into_inner().ids; match self.store.update_heartbeats(&ids).await { Ok(updated) => { diff --git a/src/grpc/server_tests.rs b/src/grpc/server_tests.rs index 7e48943d..998d1f76 100644 --- a/src/grpc/server_tests.rs +++ b/src/grpc/server_tests.rs @@ -825,10 +825,7 @@ async fn test_update_heartbeat(#[case] adapter: &str) { let response = service .batch_update_heartbeat(Request::new(BatchUpdateHeartbeatRequest { - updates: vec![ - request(&activations[0].id).into_inner(), - request("missing").into_inner(), - ], + ids: vec![activations[0].id.clone(), "missing".to_string()], })) .await .unwrap(); From 24459dca76cb48241b849d4a286044c4f3a78696 Mon Sep 17 00:00:00 2001 From: Evan Hicks Date: Fri, 6 Nov 2026 14:38:20 -0500 Subject: [PATCH 4/7] build(taskbroker): Bump sentry_protos to 0.106.0 Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 5 +++-- Cargo.toml | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 94c575d1..0c12a1fa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2996,8 +2996,9 @@ dependencies = [ [[package]] name = "sentry_protos" -version = "0.104.0" -source = "git+https://github.com/getsentry/sentry-protos?rev=0383e4fc75e5eeccd5803de9cca509653681de5c#0383e4fc75e5eeccd5803de9cca509653681de5c" +version = "0.106.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88f21dd8883662197f94aeb1fbff2f88c69829ac83605061f715071313dc936f" dependencies = [ "prost", "prost-types", diff --git a/Cargo.toml b/Cargo.toml index c20ee357..fafb4226 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -50,7 +50,7 @@ sentry = { version = "0.41.0", default-features = false, features = [ "tracing", "logs", ] } -sentry_protos = { git = "https://github.com/getsentry/sentry-protos", rev = "0383e4fc75e5eeccd5803de9cca509653681de5c" } +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. From f58d33ad0cdf1551712d6311f462c0773d262aa5 Mon Sep 17 00:00:00 2001 From: Evan Hicks Date: Mon, 9 Nov 2026 16:02:12 -0500 Subject: [PATCH 5/7] fix(taskbroker): Ignore zero heartbeat and test deadline bump Co-Authored-By: Claude Opus 5.5 --- src/kafka/deserialize_activation.rs | 6 +++++- src/store/tests.rs | 22 ++++++++++++++++++---- 2 files changed, 23 insertions(+), 5 deletions(-) diff --git a/src/kafka/deserialize_activation.rs b/src/kafka/deserialize_activation.rs index ff65b1c0..0a0ca3d3 100644 --- a/src/kafka/deserialize_activation.rs +++ b/src/kafka/deserialize_activation.rs @@ -102,7 +102,11 @@ pub fn new( received_at: activation_time, processing_deadline: None, claim_expires_at: None, - heartbeat: activation.heartbeat.and_then(|h| i32::try_from(h).ok()), + // 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, diff --git a/src/store/tests.rs b/src/store/tests.rs index e662d4d6..10e16dfc 100644 --- a/src/store/tests.rs +++ b/src/store/tests.rs @@ -1369,9 +1369,12 @@ async fn test_handle_processing_deadline_heartbeat(#[case] adapter: &str) { async fn test_heartbeat_deadline_set_and_updated(#[case] adapter: &str) { let store = create_test_store(adapter).await; - let mut batch = make_activations(3); + let mut batch = make_activations(4); batch[0].heartbeat = Some(10); batch[1].heartbeat = Some(10); + batch[3].status = ActivationStatus::Processing; + batch[3].heartbeat = Some(10); + batch[3].heartbeat_deadline = Some(Utc::now() - Duration::from_secs(60)); // Pull mode sets the heartbeat deadline when claiming assert!(store.store(&batch[0..1]).await.is_ok()); @@ -1384,7 +1387,7 @@ async fn test_heartbeat_deadline_set_and_updated(#[case] adapter: &str) { ); // Push mode sets it when marking processing - assert!(store.store(&batch[1..]).await.is_ok()); + assert!(store.store(&batch[1..3]).await.is_ok()); assert_eq!( store .claim_activations_for_push(None, None) @@ -1413,10 +1416,21 @@ async fn test_heartbeat_deadline_set_and_updated(#[case] adapter: &str) { assert_eq!(task.status, ActivationStatus::Processing); assert_eq!(task.heartbeat_deadline, None); - let ids = [&batch[0].id, &batch[1].id, &batch[2].id, "missing"].map(String::from); + assert!(store.store(&batch[3..]).await.is_ok()); + let ids = [ + &batch[0].id, + &batch[1].id, + &batch[3].id, + &batch[2].id, + "missing", + ] + .map(String::from); let mut updated = store.update_heartbeats(&ids).await.unwrap(); updated.sort(); - assert_eq!(updated, ids[0..2]); + assert_eq!(updated, ids[0..3]); + + let task = store.get_by_id(&batch[3].id).await.unwrap().unwrap(); + assert!(task.heartbeat_deadline.unwrap() > Utc::now()); store.remove_db().await.unwrap(); } From 5aaa1b085f98841c9b98f9b1561836971eab8006 Mon Sep 17 00:00:00 2001 From: Evan Hicks Date: Mon, 9 Nov 2026 16:35:13 -0500 Subject: [PATCH 6/7] ref(taskbroker): Check both deadlines in one upkeep query and add heartbeat metrics Co-Authored-By: Claude Opus 5.5 --- src/grpc/server.rs | 14 +++++- src/store/adapters/postgres.rs | 81 ++++++++++++++++------------------ src/store/adapters/sqlite.rs | 75 ++++++++++++++----------------- 3 files changed, 81 insertions(+), 89 deletions(-) diff --git a/src/grpc/server.rs b/src/grpc/server.rs index 73343f30..ff6e339a 100644 --- a/src/grpc/server.rs +++ b/src/grpc/server.rs @@ -373,9 +373,13 @@ impl ConsumerService for TaskbrokerServer { &self, request: Request, ) -> Result, Status> { + let start_time = Instant::now(); let id = &request.get_ref().id; - match self.store.update_heartbeats(std::slice::from_ref(id)).await { + 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 {})) @@ -404,9 +408,15 @@ impl ConsumerService for TaskbrokerServer { &self, request: Request, ) -> Result, 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 self.store.update_heartbeats(&ids).await { + match result { Ok(updated) => { let updated: HashSet = updated.into_iter().collect(); let not_found_ids: Vec = diff --git a/src/store/adapters/postgres.rs b/src/store/adapters/postgres.rs index 63695c53..8e2ddf9b 100644 --- a/src/store/adapters/postgres.rs +++ b/src/store/adapters/postgres.rs @@ -1213,54 +1213,47 @@ impl ActivationStore for PostgresStore { async fn handle_processing_deadline(&self) -> Result { let now = Utc::now(); let mut atomic = self.write_pool.begin().await?; - let mut modified_rows = 0; - - for (deadline, reason) in [ - ("processing_deadline", "processing_deadline"), - ("heartbeat_deadline", "heartbeat"), - ] { - // At-most-once tasks that fail their deadlines go directly to failure - // there are no retries, as the worker will reject the task due to at_most_once keys. - let mut query_builder = QueryBuilder::new( - "UPDATE inflight_taskactivations - SET processing_deadline = null, heartbeat_deadline = null, status = ", - ); - query_builder.push_bind(ActivationStatus::Failure.to_string()); - query_builder.push(format!(" WHERE {deadline} < ")); - query_builder.push_bind(now); - query_builder.push(" AND at_most_once = TRUE AND status = "); - query_builder.push_bind(ActivationStatus::Processing.to_string()); - - self.add_partition_condition(&mut query_builder, false); + // At-most-once tasks that miss a deadline go directly to failure + // there are no retries, as the worker will reject the task due to at_most_once keys. + let mut query_builder = QueryBuilder::new( + "UPDATE inflight_taskactivations + SET processing_deadline = null, heartbeat_deadline = null, status = ", + ); + query_builder.push_bind(ActivationStatus::Failure.to_string()); + query_builder.push(" WHERE (processing_deadline < "); + query_builder.push_bind(now); + query_builder.push(" OR heartbeat_deadline < "); + query_builder.push_bind(now); + query_builder.push(") AND at_most_once = TRUE AND status = "); + query_builder.push_bind(ActivationStatus::Processing.to_string()); - let mut reset = 0; - if let Ok(query_res) = query_builder.build().execute(&mut *atomic).await { - reset = query_res.rows_affected(); - } + self.add_partition_condition(&mut query_builder, false); - // Update regular tasks. - // Increment processing_attempts by 1 and reset deadlines to null. - let mut query_builder = QueryBuilder::new( - "UPDATE inflight_taskactivations - SET processing_deadline = null, heartbeat_deadline = null, status = ", - ); - query_builder.push_bind(ActivationStatus::Pending.to_string()); - query_builder.push(", processing_attempts = processing_attempts + 1"); - query_builder.push(format!(" WHERE {deadline} < ")); - query_builder.push_bind(now); - query_builder.push(" AND status = "); - query_builder.push_bind(ActivationStatus::Processing.to_string()); - self.add_partition_condition(&mut query_builder, false); + let mut modified_rows = 0; + if let Ok(query_res) = query_builder.build().execute(&mut *atomic).await { + modified_rows = query_res.rows_affected(); + } - let Ok(query_res) = query_builder.build().execute(&mut *atomic).await else { - return Err(anyhow!("Could not update tasks past {deadline}")); - }; - reset += query_res.rows_affected(); + // Update regular tasks. + // Increment processing_attempts by 1 and reset deadlines to null. + let mut query_builder = QueryBuilder::new( + "UPDATE inflight_taskactivations + SET processing_deadline = null, heartbeat_deadline = null, status = ", + ); + query_builder.push_bind(ActivationStatus::Pending.to_string()); + query_builder.push(", processing_attempts = processing_attempts + 1"); + query_builder.push(" WHERE (processing_deadline < "); + query_builder.push_bind(now); + query_builder.push(" OR heartbeat_deadline < "); + query_builder.push_bind(now); + query_builder.push(") AND status = "); + query_builder.push_bind(ActivationStatus::Processing.to_string()); + self.add_partition_condition(&mut query_builder, false); - metrics::counter!("upkeep.handle_processing_deadline.reset", "reason" => reason) - .increment(reset); - modified_rows += reset; - } + let Ok(query_res) = query_builder.build().execute(&mut *atomic).await else { + return Err(anyhow!("Could not update tasks past their deadlines")); + }; + modified_rows += query_res.rows_affected(); atomic.commit().await?; diff --git a/src/store/adapters/sqlite.rs b/src/store/adapters/sqlite.rs index 9eec0672..90e10777 100644 --- a/src/store/adapters/sqlite.rs +++ b/src/store/adapters/sqlite.rs @@ -1065,52 +1065,41 @@ impl ActivationStore for SqliteStore { async fn handle_processing_deadline(&self) -> Result { let now = Utc::now(); let mut atomic = self.write_pool.begin().await?; - let mut modified_rows = 0; - - for (deadline, reason) in [ - ("processing_deadline", "processing_deadline"), - ("heartbeat_deadline", "heartbeat"), - ] { - // Idempotent tasks that fail their deadlines go directly to failure - // there are no retries, as the worker will reject the task due to idempotency keys. - let most_once_result = sqlx::query(&format!( - "UPDATE inflight_taskactivations - SET processing_deadline = null, heartbeat_deadline = null, status = $1 - WHERE {deadline} < $2 AND at_most_once = TRUE AND status = $3", - )) - .bind(ActivationStatus::Failure) - .bind(now.timestamp()) - .bind(ActivationStatus::Processing) - .execute(&mut *atomic) - .await; - - let mut reset = 0; - if let Ok(query_res) = most_once_result { - reset = query_res.rows_affected(); - } + // Idempotent tasks that miss a deadline go directly to failure + // there are no retries, as the worker will reject the task due to idempotency keys. + let most_once_result = sqlx::query( + "UPDATE inflight_taskactivations + SET processing_deadline = null, heartbeat_deadline = null, status = $1 + WHERE (processing_deadline < $2 OR heartbeat_deadline < $2) AND at_most_once = TRUE AND status = $3", + ) + .bind(ActivationStatus::Failure) + .bind(now.timestamp()) + .bind(ActivationStatus::Processing) + .execute(&mut *atomic) + .await; - // Update non-idempotent tasks. - // Increment processing_attempts by 1 and reset deadlines to null. - let result = sqlx::query(&format!( - "UPDATE inflight_taskactivations - SET processing_deadline = null, heartbeat_deadline = null, status = $1, processing_attempts = processing_attempts + 1 - WHERE {deadline} < $2 AND status = $3", - )) - .bind(ActivationStatus::Pending) - .bind(now.timestamp()) - .bind(ActivationStatus::Processing) - .execute(&mut *atomic) - .await; + let mut modified_rows = 0; + if let Ok(query_res) = most_once_result { + modified_rows = query_res.rows_affected(); + } - let Ok(query_res) = result else { - return Err(anyhow!("Could not update tasks past {deadline}")); - }; - reset += query_res.rows_affected(); + // Update non-idempotent tasks. + // Increment processing_attempts by 1 and reset deadlines to null. + let result = sqlx::query( + "UPDATE inflight_taskactivations + SET processing_deadline = null, heartbeat_deadline = null, status = $1, processing_attempts = processing_attempts + 1 + WHERE (processing_deadline < $2 OR heartbeat_deadline < $2) AND status = $3", + ) + .bind(ActivationStatus::Pending) + .bind(now.timestamp()) + .bind(ActivationStatus::Processing) + .execute(&mut *atomic) + .await; - metrics::counter!("upkeep.handle_processing_deadline.reset", "reason" => reason) - .increment(reset); - modified_rows += reset; - } + let Ok(query_res) = result else { + return Err(anyhow!("Could not update tasks past their deadlines")); + }; + modified_rows += query_res.rows_affected(); atomic.commit().await?; From 8f5b09d55362ff904d30213f8bc5ca81c3af000f Mon Sep 17 00:00:00 2001 From: Evan Hicks Date: Mon, 9 Nov 2026 16:51:13 -0500 Subject: [PATCH 7/7] ref(taskbroker): Restore per-reason upkeep queries and skip rows without heartbeat Co-Authored-By: Claude Opus 5.5 --- src/store/adapters/postgres.rs | 87 +++++++++++++++++++--------------- src/store/adapters/sqlite.rs | 79 +++++++++++++++++------------- 2 files changed, 97 insertions(+), 69 deletions(-) diff --git a/src/store/adapters/postgres.rs b/src/store/adapters/postgres.rs index 8e2ddf9b..4320ece3 100644 --- a/src/store/adapters/postgres.rs +++ b/src/store/adapters/postgres.rs @@ -1213,47 +1213,60 @@ impl ActivationStore for PostgresStore { async fn handle_processing_deadline(&self) -> Result { let now = Utc::now(); let mut atomic = self.write_pool.begin().await?; - // At-most-once tasks that miss a deadline go directly to failure - // there are no retries, as the worker will reject the task due to at_most_once keys. - let mut query_builder = QueryBuilder::new( - "UPDATE inflight_taskactivations - SET processing_deadline = null, heartbeat_deadline = null, status = ", - ); - query_builder.push_bind(ActivationStatus::Failure.to_string()); - query_builder.push(" WHERE (processing_deadline < "); - query_builder.push_bind(now); - query_builder.push(" OR heartbeat_deadline < "); - query_builder.push_bind(now); - query_builder.push(") AND at_most_once = TRUE AND status = "); - query_builder.push_bind(ActivationStatus::Processing.to_string()); + let mut modified_rows = 0; - self.add_partition_condition(&mut query_builder, false); + for (deadline, filter, reason) in [ + ("processing_deadline", "", "processing_deadline"), + ( + "heartbeat_deadline", + " AND heartbeat IS NOT NULL", + "heartbeat", + ), + ] { + // At-most-once tasks that fail their deadlines go directly to failure + // there are no retries, as the worker will reject the task due to at_most_once keys. + let mut query_builder = QueryBuilder::new( + "UPDATE inflight_taskactivations + SET processing_deadline = null, heartbeat_deadline = null, status = ", + ); + query_builder.push_bind(ActivationStatus::Failure.to_string()); + query_builder.push(format!(" WHERE {deadline} < ")); + query_builder.push_bind(now); + query_builder.push(filter); + query_builder.push(" AND at_most_once = TRUE AND status = "); + query_builder.push_bind(ActivationStatus::Processing.to_string()); - let mut modified_rows = 0; - if let Ok(query_res) = query_builder.build().execute(&mut *atomic).await { - modified_rows = query_res.rows_affected(); - } + self.add_partition_condition(&mut query_builder, false); - // Update regular tasks. - // Increment processing_attempts by 1 and reset deadlines to null. - let mut query_builder = QueryBuilder::new( - "UPDATE inflight_taskactivations - SET processing_deadline = null, heartbeat_deadline = null, status = ", - ); - query_builder.push_bind(ActivationStatus::Pending.to_string()); - query_builder.push(", processing_attempts = processing_attempts + 1"); - query_builder.push(" WHERE (processing_deadline < "); - query_builder.push_bind(now); - query_builder.push(" OR heartbeat_deadline < "); - query_builder.push_bind(now); - query_builder.push(") AND status = "); - query_builder.push_bind(ActivationStatus::Processing.to_string()); - self.add_partition_condition(&mut query_builder, false); + let mut reset = 0; + if let Ok(query_res) = query_builder.build().execute(&mut *atomic).await { + reset = query_res.rows_affected(); + } - let Ok(query_res) = query_builder.build().execute(&mut *atomic).await else { - return Err(anyhow!("Could not update tasks past their deadlines")); - }; - modified_rows += query_res.rows_affected(); + // Update regular tasks. + // Increment processing_attempts by 1 and reset deadlines to null. + let mut query_builder = QueryBuilder::new( + "UPDATE inflight_taskactivations + SET processing_deadline = null, heartbeat_deadline = null, status = ", + ); + query_builder.push_bind(ActivationStatus::Pending.to_string()); + query_builder.push(", processing_attempts = processing_attempts + 1"); + query_builder.push(format!(" WHERE {deadline} < ")); + query_builder.push_bind(now); + query_builder.push(filter); + query_builder.push(" AND status = "); + query_builder.push_bind(ActivationStatus::Processing.to_string()); + self.add_partition_condition(&mut query_builder, false); + + let Ok(query_res) = query_builder.build().execute(&mut *atomic).await else { + return Err(anyhow!("Could not update tasks past {deadline}")); + }; + reset += query_res.rows_affected(); + + metrics::counter!("upkeep.handle_processing_deadline.reset", "reason" => reason) + .increment(reset); + modified_rows += reset; + } atomic.commit().await?; diff --git a/src/store/adapters/sqlite.rs b/src/store/adapters/sqlite.rs index 90e10777..2d4f097e 100644 --- a/src/store/adapters/sqlite.rs +++ b/src/store/adapters/sqlite.rs @@ -1065,41 +1065,56 @@ impl ActivationStore for SqliteStore { async fn handle_processing_deadline(&self) -> Result { let now = Utc::now(); let mut atomic = self.write_pool.begin().await?; - // Idempotent tasks that miss a deadline go directly to failure - // there are no retries, as the worker will reject the task due to idempotency keys. - let most_once_result = sqlx::query( - "UPDATE inflight_taskactivations - SET processing_deadline = null, heartbeat_deadline = null, status = $1 - WHERE (processing_deadline < $2 OR heartbeat_deadline < $2) AND at_most_once = TRUE AND status = $3", - ) - .bind(ActivationStatus::Failure) - .bind(now.timestamp()) - .bind(ActivationStatus::Processing) - .execute(&mut *atomic) - .await; - let mut modified_rows = 0; - if let Ok(query_res) = most_once_result { - modified_rows = query_res.rows_affected(); - } - // Update non-idempotent tasks. - // Increment processing_attempts by 1 and reset deadlines to null. - let result = sqlx::query( - "UPDATE inflight_taskactivations - SET processing_deadline = null, heartbeat_deadline = null, status = $1, processing_attempts = processing_attempts + 1 - WHERE (processing_deadline < $2 OR heartbeat_deadline < $2) AND status = $3", - ) - .bind(ActivationStatus::Pending) - .bind(now.timestamp()) - .bind(ActivationStatus::Processing) - .execute(&mut *atomic) - .await; + for (deadline, filter, reason) in [ + ("processing_deadline", "", "processing_deadline"), + ( + "heartbeat_deadline", + " AND heartbeat IS NOT NULL", + "heartbeat", + ), + ] { + // Idempotent tasks that fail their deadlines go directly to failure + // there are no retries, as the worker will reject the task due to idempotency keys. + let most_once_result = sqlx::query(&format!( + "UPDATE inflight_taskactivations + SET processing_deadline = null, heartbeat_deadline = null, status = $1 + WHERE {deadline} < $2{filter} AND at_most_once = TRUE AND status = $3", + )) + .bind(ActivationStatus::Failure) + .bind(now.timestamp()) + .bind(ActivationStatus::Processing) + .execute(&mut *atomic) + .await; - let Ok(query_res) = result else { - return Err(anyhow!("Could not update tasks past their deadlines")); - }; - modified_rows += query_res.rows_affected(); + let mut reset = 0; + if let Ok(query_res) = most_once_result { + reset = query_res.rows_affected(); + } + + // Update non-idempotent tasks. + // Increment processing_attempts by 1 and reset deadlines to null. + let result = sqlx::query(&format!( + "UPDATE inflight_taskactivations + SET processing_deadline = null, heartbeat_deadline = null, status = $1, processing_attempts = processing_attempts + 1 + WHERE {deadline} < $2{filter} AND status = $3", + )) + .bind(ActivationStatus::Pending) + .bind(now.timestamp()) + .bind(ActivationStatus::Processing) + .execute(&mut *atomic) + .await; + + let Ok(query_res) = result else { + return Err(anyhow!("Could not update tasks past {deadline}")); + }; + reset += query_res.rows_affected(); + + metrics::counter!("upkeep.handle_processing_deadline.reset", "reason" => reason) + .increment(reset); + modified_rows += reset; + } atomic.commit().await?;