From 8092b31dae1fd3a8f23dfb7a6ddb8b640630c1b3 Mon Sep 17 00:00:00 2001 From: james-mcnulty Date: Mon, 13 Jul 2026 16:45:46 -0400 Subject: [PATCH 1/3] Add Canary Task Sending to Taskbroker --- .../python/src/taskbroker_client/canary.py | 2 +- src/canary_tasks.rs | 179 ++++++++++++++++++ src/config/mod.rs | 10 + src/lib.rs | 1 + src/main.rs | 6 +- 5 files changed, 196 insertions(+), 2 deletions(-) create mode 100644 src/canary_tasks.rs diff --git a/clients/python/src/taskbroker_client/canary.py b/clients/python/src/taskbroker_client/canary.py index c2d1c536..e9a2661f 100644 --- a/clients/python/src/taskbroker_client/canary.py +++ b/clients/python/src/taskbroker_client/canary.py @@ -9,6 +9,6 @@ def canary_task() -> None: - logger.info("Running canary task...") + print("Running canary task...") sleep(0.1) print("Done running canary task!") diff --git a/src/canary_tasks.rs b/src/canary_tasks.rs new file mode 100644 index 00000000..bbe52df4 --- /dev/null +++ b/src/canary_tasks.rs @@ -0,0 +1,179 @@ +use std::collections::BTreeMap; + +use anyhow::{Context, Error, anyhow}; +use chrono::Utc; +use prost::Message as _; +use sentry_protos::taskbroker::v1::{OnAttemptsExceeded, TaskActivation}; +use serde::Serialize; +use tracing::info; +use uuid::Uuid; + +use crate::config::Config; +use crate::kafka::deserialize_activation::bucket_from_id; +use crate::store::activation::{Activation, ActivationStatus}; +use crate::store::traits::ActivationStore; + +const CANARY_NAMESPACE: &str = "internal"; +const CANARY_TASKNAME: &str = "canary_task"; +const CANARY_PROCESSING_DEADLINE_SECONDS: i32 = 10; + +#[derive(Serialize)] +struct CanaryParameters { + args: Vec<()>, + kwargs: BTreeMap, +} + +/// Add the configured number of canary tasks to the activation store for every +/// application in `worker_map`. +pub async fn enqueue(config: &Config, store: &dyn ActivationStore) -> Result<(), Error> { + if config.canary_tasks == 0 || config.worker_map.is_empty() { + return Ok(()); + } + + let topic = config + .consumable_topics() + .map_err(|error| anyhow!(error.to_string()))? + .first() + .map(|(name, _)| *name) + .ok_or_else(|| anyhow!("No consumable topic configured"))?; + + let activations = build_activations(config, topic)?; + let batch_size = config.store.insert_batch_max_length.max(1); + let mut stored = 0; + + for batch in activations.chunks(batch_size) { + stored += store.store(batch).await?; + } + + if stored > 0 { + info!( + count = stored, + worker_pools = config.worker_map.len(), + "Stored startup canary tasks" + ); + } + + Ok(()) +} + +fn build_activations(config: &Config, topic: &str) -> Result, Error> { + let parameters_bytes = rmp_serde::to_vec_named(&CanaryParameters { + args: Vec::new(), + kwargs: BTreeMap::new(), + }) + .context("Failed to serialize canary task parameters")?; + + let mut activations = Vec::new(); + + for application in config.worker_map.keys() { + for offset in 0..config.canary_tasks { + let now = Utc::now(); + let id = Uuid::new_v4().to_string(); + let task_activation = TaskActivation { + id: id.clone(), + application: Some(application.clone()), + namespace: CANARY_NAMESPACE.to_owned(), + taskname: CANARY_TASKNAME.to_owned(), + parameters_bytes: parameters_bytes.clone(), + processing_deadline_duration: CANARY_PROCESSING_DEADLINE_SECONDS as u64, + received_at: Some(prost_types::Timestamp { + seconds: now.timestamp(), + nanos: now.timestamp_subsec_nanos() as i32, + }), + ..Default::default() + }; + + activations.push(Activation { + id: id.clone(), + application: application.clone(), + namespace: CANARY_NAMESPACE.to_owned(), + taskname: CANARY_TASKNAME.to_owned(), + activation: task_activation.encode_to_vec(), + status: ActivationStatus::Pending, + topic: topic.to_owned(), + partition: 0, + offset: offset.into(), + added_at: now, + received_at: now, + processing_attempts: 0, + processing_deadline_duration: CANARY_PROCESSING_DEADLINE_SECONDS, + expires_at: None, + delay_until: None, + processing_deadline: None, + claim_expires_at: None, + on_attempts_exceeded: OnAttemptsExceeded::Discard, + at_most_once: false, + bucket: bucket_from_id(&id), + }); + } + } + + Ok(activations) +} + +#[cfg(test)] +mod tests { + use std::collections::{BTreeMap, HashMap}; + + use prost::Message as _; + use sentry_protos::taskbroker::v1::TaskActivation; + use serde::Deserialize; + + use super::{CANARY_NAMESPACE, CANARY_TASKNAME, build_activations}; + use crate::config::Config; + + #[derive(Debug, Deserialize, PartialEq)] + struct CanaryParameters { + args: Vec<()>, + kwargs: BTreeMap, + } + + #[test] + fn builds_configured_number_of_canaries_per_worker_pool() { + let config = Config { + canary_tasks: 3, + worker_map: BTreeMap::from([ + ("launchpad".to_owned(), "http://launchpad".to_owned()), + ("sentry".to_owned(), "http://sentry".to_owned()), + ]), + ..Default::default() + }; + + let activations = build_activations(&config, "taskworker").unwrap(); + assert_eq!(activations.len(), 6); + + let mut counts = HashMap::new(); + for activation in activations { + *counts.entry(activation.application.clone()).or_insert(0) += 1; + assert_eq!(activation.namespace, CANARY_NAMESPACE); + assert_eq!(activation.taskname, CANARY_TASKNAME); + assert_eq!(activation.topic, "taskworker"); + assert_eq!(activation.processing_deadline_duration, 10); + + let task = TaskActivation::decode(activation.activation.as_slice()).unwrap(); + assert_eq!(task.id, activation.id); + assert_eq!( + task.application.as_deref(), + Some(activation.application.as_str()) + ); + assert_eq!(task.namespace, CANARY_NAMESPACE); + assert_eq!(task.taskname, CANARY_TASKNAME); + assert_eq!( + rmp_serde::from_slice::(&task.parameters_bytes).unwrap(), + CanaryParameters { + args: Vec::new(), + kwargs: BTreeMap::new(), + } + ); + } + + assert_eq!(counts.get("launchpad"), Some(&3)); + assert_eq!(counts.get("sentry"), Some(&3)); + } + + #[test] + fn builds_no_canaries_when_disabled() { + let config = Config::default(); + assert!(build_activations(&config, "taskworker").unwrap().is_empty()); + } +} diff --git a/src/config/mod.rs b/src/config/mod.rs index 92731003..36eeb23b 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -166,6 +166,10 @@ pub struct Config { /// when it starts up, but before the GRPC server, consumer and upkeep begin. pub full_vacuum_on_start: bool, + /// Number of internal canary tasks to add to the store for each configured + /// worker pool when the broker starts. Set to zero to disable. + pub canary_tasks: u32, + /// Enable the upkeep thread to perforam a full `VACUUM` on the database /// periodically. pub full_vacuum_on_upkeep: bool, @@ -267,6 +271,7 @@ impl Default for Config { max_message_size: 5000000, grpc_max_message_size: 52 * 1024 * 1024, // 52MB full_vacuum_on_start: true, + canary_tasks: 0, full_vacuum_on_upkeep: true, vacuum_interval_ms: 30000, log_async_backtrace: false, @@ -915,6 +920,7 @@ mod tests { assert_eq!(config.store.max_pending_count, 2048); assert_eq!(config.store.max_processing_count, 2048); assert_eq!(config.store.sqlite.vacuum_page_count, None); + assert_eq!(config.canary_tasks, 0); assert_eq!( config.worker_map.get("sentry").map(String::as_str), Some("http://127.0.0.1:50052") @@ -1026,6 +1032,7 @@ mod tests { max_processing_attempts: 5 vacuum_page_count: 1000 full_vacuum_on_start: true + canary_tasks: 3 worker_map: sentry: http://worker-sentry:50052 launchpad: http://worker-launchpad:50053 @@ -1064,6 +1071,7 @@ mod tests { assert_eq!(config.store.sqlite.vacuum_page_count, Some(1000)); assert_eq!(config.store.max_size, Some(3_000_000_000)); assert!(config.full_vacuum_on_start); + assert_eq!(config.canary_tasks, 3); assert_eq!( config.worker_map, BTreeMap::from([ @@ -1085,6 +1093,7 @@ mod tests { jail.set_env("TASKBROKER_LOG_FILTER", "error"); jail.set_env("TASKBROKER_DATABASE_ADAPTER", "postgres"); jail.set_env("TASKBROKER_MAX_PROCESSING_ATTEMPTS", "5"); + jail.set_env("TASKBROKER_CANARY_TASKS", "2"); let args = Args { run: Run::Broker, @@ -1094,6 +1103,7 @@ mod tests { assert_eq!(config.log_filter, "error"); assert_eq!(config.store.adapter, DatabaseAdapter::Postgres); assert_eq!(config.store.max_processing_attempts, 5); + assert_eq!(config.canary_tasks, 2); Ok(()) }); diff --git a/src/lib.rs b/src/lib.rs index c8af6983..76b221e2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,7 @@ use clap::{Parser, ValueEnum}; use std::fs; +pub mod canary_tasks; pub mod config; pub mod fetch; pub mod flusher; diff --git a/src/main.rs b/src/main.rs index 4acc4491..4c127b40 100644 --- a/src/main.rs +++ b/src/main.rs @@ -15,6 +15,7 @@ use tonic::transport::Server; use tonic_health::ServingStatus; use tracing::{debug, error, info, warn}; +use taskbroker::canary_tasks; use taskbroker::config::store::DatabaseAdapter; use taskbroker::config::{Config, DeliveryMode}; use taskbroker::fetch::FetchPool; @@ -106,7 +107,10 @@ async fn main() -> Result<(), Error> { Err(err) => error!("Failed to run full vacuum on startup: {:?}", err), } } - // Get startup time after migrations and vacuum + + canary_tasks::enqueue(&config, store.as_ref()).await?; + + // Get startup time after migrations, vacuum, and startup canaries. let startup_time = Utc::now(); // Taskbroker exposes a grpc.v1.health endpoint. We use upkeep to track the health From b8e4036d0d0dd75222af1e7587fef0ae04a3553e Mon Sep 17 00:00:00 2001 From: james-mcnulty Date: Mon, 13 Jul 2026 16:51:58 -0400 Subject: [PATCH 2/3] Minor Comment Tweak --- src/main.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main.rs b/src/main.rs index 4c127b40..23f19524 100644 --- a/src/main.rs +++ b/src/main.rs @@ -110,7 +110,7 @@ async fn main() -> Result<(), Error> { canary_tasks::enqueue(&config, store.as_ref()).await?; - // Get startup time after migrations, vacuum, and startup canaries. + // Get startup time after migrations, vacuum, and startup canaries let startup_time = Utc::now(); // Taskbroker exposes a grpc.v1.health endpoint. We use upkeep to track the health From b95c1a064b4e6bb4982649a2476cb48a1489842b Mon Sep 17 00:00:00 2001 From: james-mcnulty Date: Wed, 15 Jul 2026 16:16:18 -0400 Subject: [PATCH 3/3] Fix Canary Task Tests --- clients/python/tests/test_app.py | 8 ++----- clients/python/tests/worker/test_worker.py | 26 ++++++++++------------ 2 files changed, 14 insertions(+), 20 deletions(-) diff --git a/clients/python/tests/test_app.py b/clients/python/tests/test_app.py index cc14d264..e4686681 100644 --- a/clients/python/tests/test_app.py +++ b/clients/python/tests/test_app.py @@ -1,5 +1,3 @@ -from unittest.mock import patch - import msgpack import pytest from sentry_protos.taskbroker.v1.taskbroker_pb2 import TaskActivation @@ -51,11 +49,9 @@ def test_registers_internal_canary_task(capsys: pytest.CaptureFixture[str]) -> N assert namespace.name == INTERNAL_NAMESPACE assert isinstance(task, Task) assert task.name == CANARY_TASK_NAME - with patch("taskbroker_client.canary.logger") as mock_logger: - task() + task() - mock_logger.info.assert_called_once_with("Running canary task...") - assert capsys.readouterr().out == "Done running canary task!\n" + assert capsys.readouterr().out == "Running canary task...\nDone running canary task!\n" def test_should_attempt_at_most_once() -> None: diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 7253a74e..8e849002 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -1053,24 +1053,22 @@ def test_child_process_canary_task(capsys: pytest.CaptureFixture[str]) -> None: shutdown = Event() todo.put(CANARY_TASK) - with mock.patch("taskbroker_client.canary.logger") as mock_logger: - child_process( - "examples.app:app", - todo, - processed, - shutdown, - max_task_count=1, - processing_pool_name="test", - process_type="fork", - skip_awaiting_futures=False, - future_checking_frequency=0.1, - ) + child_process( + "examples.app:app", + todo, + processed, + shutdown, + max_task_count=1, + processing_pool_name="test", + process_type="fork", + skip_awaiting_futures=False, + future_checking_frequency=0.1, + ) result = processed.get() assert result.task_id == CANARY_TASK.activation.id assert result.status == TASK_ACTIVATION_STATUS_COMPLETE - mock_logger.info.assert_called_once_with("Running canary task...") - assert capsys.readouterr().out == "Done running canary task!\n" + assert capsys.readouterr().out == "Running canary task...\nDone running canary task!\n" def test_child_process_emits_running_message() -> None: