diff --git a/Cargo.lock b/Cargo.lock index 2a258356..0c12a1fa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2996,9 +2996,9 @@ dependencies = [ [[package]] name = "sentry_protos" -version = "0.22.3" +version = "0.106.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6214388c55cb613be274cfdbb14cbb8ec47ac24ddbbf848ab08d4eefc793c00" +checksum = "88f21dd8883662197f94aeb1fbff2f88c69829ac83605061f715071313dc936f" dependencies = [ "prost", "prost-types", diff --git a/Cargo.toml b/Cargo.toml index 0470aef2..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 = "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. 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..b41fdd8c 100644 --- a/src/fetch/tests.rs +++ b/src/fetch/tests.rs @@ -189,6 +189,10 @@ impl ActivationStore for MockStore { unimplemented!() } + async fn update_heartbeats(&self, _ids: &[String]) -> Result, Error> { + 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..ff6e339a 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, + 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, + ) -> Result, 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, + ) -> 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 result { + 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 3a6bafda..998d1f76 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, + 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()] + ); +} diff --git a/src/kafka/deserialize_activation.rs b/src/kafka/deserialize_activation.rs index b373c31f..0a0ca3d3 100644 --- a/src/kafka/deserialize_activation.rs +++ b/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, 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(vec![]) + } + 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..4320ece3 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_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 = ANY($1) AND status = $2 AND heartbeat IS NOT NULL + RETURNING id", + ) + .bind(ids) + .bind(ActivationStatus::Processing.to_string()) + .fetch_all(&mut *conn) + .await?) + }) + .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,64 @@ 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, 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()); - // 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(filter); + 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..2d4f097e 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,28 @@ impl ActivationStore for SqliteStore { Ok(()) } + #[instrument(skip_all)] + 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 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)] async fn delete_activation(&self, id: &str) -> Result<(), Error> { let mut conn = self.acquire_write_conn_metric("delete_activation").await?; @@ -967,6 +1004,8 @@ impl ActivationStore for SqliteStore { processing_deadline_duration, processing_deadline, claim_expires_at, + heartbeat, + heartbeat_deadline, status, at_most_once, application, @@ -1026,46 +1065,60 @@ 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, 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; - // 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{filter} 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..10e16dfc 100644 --- a/src/store/tests.rs +++ b/src/store/tests.rs @@ -1330,6 +1330,110 @@ 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(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()); + assert!( + store + .claim_activation_for_pull(None, None) + .await + .unwrap() + .is_some() + ); + + // Push mode sets it when marking processing + assert!(store.store(&batch[1..3]).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); + } + + 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.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..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(); +} + #[tokio::test] #[rstest] #[case::sqlite("sqlite")] diff --git a/src/store/traits.rs b/src/store/traits.rs index f70bb47c..f33d2d45 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 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>; 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 };