feat(taskbroker): Add Canary Task Sending to Taskbroker

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!")

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:

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:

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<String, ()>,
+}
+
+/// 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<Vec<Activation>, 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<String, ()>,
+    }
+
+    #[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::<CanaryParameters>(&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());
+    }
+}

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,
@@ -268,6 +272,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,
@@ -957,6 +962,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!(config.worker_map.is_empty());
     }
 
@@ -1066,6 +1072,7 @@ mod tests {
                 max_processing_attempts: 5
                 vacuum_page_count: 1000
                 full_vacuum_on_start: true
+                canary_tasks: 3
                 worker_map:
                     sentry: http://sentry:50052
                     launchpad: http://launchpad:50052
@@ -1104,6 +1111,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([
@@ -1149,6 +1157,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,
@@ -1158,6 +1167,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(())
         });

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;

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;
@@ -108,7 +109,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