From 34ac651b9a3b91f95e79923ba63ef8e2691a8b60 Mon Sep 17 00:00:00 2001 From: ShashSriv Date: Thu, 1 Oct 2026 15:53:18 -0700 Subject: [PATCH 1/5] feat(consumer): Discard activations from applications a pool cannot run A pool can only execute tasks for the applications its workers serve. Nothing checked that, so a misrouted activation was immortal: in pull mode no worker ever claimed it and it sat Pending until the lag alert paged someone, and in push mode it was claimed, found no worker, had its claim expire, and was fetched again forever. Add an `applications` allowlist to broker config and discard anything outside it in ActivationBatcher::reduce, alongside the existing killswitch and expiry filters, publishing the payload to the pool's deadletter topic so a misconfiguration is recoverable rather than silent data loss. Push pools derive the set from `worker_map`, which already lists every application they can deliver to. Pull pools are never given a `worker_map`, so an unset `applications` admits everything and they keep running until ops opts them in; a pool whose `worker_map` names an application outside the allowlist now fails to start, since its tasks would be discarded before reaching that worker. Two supporting changes for activations that are already stored, which the consumer filter cannot reach: the push thread now sets Failure instead of returning silently, handing the task to upkeep's discard/deadletter path, and get_task returns failed_precondition for an unserved application so a worker pointed at the wrong pool fails instead of polling an empty queue. The producer swap and batch send in flush are extracted into helpers and shared with the forwarding path rather than copied. Co-Authored-By: Claude Opus 5 --- src/config/mod.rs | 152 +++++++++++++++++++++- src/grpc/server.rs | 10 ++ src/grpc/server_tests.rs | 26 ++++ src/kafka/activation_batcher.rs | 216 ++++++++++++++++++++++++++------ src/push/tests.rs | 63 +++++++++- src/push/thread.rs | 14 +++ 6 files changed, 442 insertions(+), 39 deletions(-) diff --git a/src/config/mod.rs b/src/config/mod.rs index 11df7567..60b7dff5 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -1,6 +1,6 @@ #![allow(clippy::result_large_err)] use std::borrow::Cow; -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use std::time::Duration; use anyhow::{Result, anyhow}; @@ -49,6 +49,27 @@ pub enum DeliveryMode { Push, } +/// The applications a pool admits. An empty set admits every application. +#[derive(Clone, Debug, Default, PartialEq, Deserialize, Serialize)] +#[serde(transparent)] +pub struct Applications(BTreeSet); + +impl Applications { + pub fn admits(&self, application: &str) -> bool { + self.0.is_empty() || self.0.contains(application) + } + + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } +} + +impl FromIterator for Applications { + fn from_iter>(iter: T) -> Self { + Self(iter.into_iter().collect()) + } +} + #[derive(Clone, PartialEq, Debug, Deserialize, Serialize, Validate)] pub struct Config { /// Deprecated configuration options. Not meant to be used. @@ -216,6 +237,12 @@ pub struct Config { #[validate(length(min = 1))] pub worker_map: BTreeMap, + /// The applications this pool can execute. Activations for any other + /// application are discarded by the consumer. Normalization derives this + /// from `worker_map` in push mode; left empty, the check is disabled. + #[serde(default)] + pub applications: Applications, + /// The namespace to assign to raw mode activations. pub raw_namespace: Option, @@ -291,6 +318,7 @@ impl Default for Config { status_update_batch_size: 1, status_update_interval_ms: 100, worker_map: [].into(), + applications: Applications::default(), raw_namespace: None, raw_application: None, raw_taskname: None, @@ -810,6 +838,43 @@ impl Config { } } + self.normalize_applications(); + self.validate_applications()?; + + Ok(()) + } + + /// Push pools already declare every application they can deliver to in + /// `worker_map`, so they never have to repeat it. Pull pools have no such + /// list, and an empty `applications` admits everything, so a pool that + /// hasn't opted in keeps running rather than discarding all of its work. + fn normalize_applications(&mut self) { + if !self.applications.is_empty() { + return; + } + + if self.delivery_mode == DeliveryMode::Push { + self.applications = self.worker_map.keys().cloned().collect(); + } else { + warn!("applications is unset, activations will not be checked against it"); + } + } + + /// A `worker_map` entry outside `applications` is unreachable: its activations + /// are discarded before they are ever pushed to that worker. + fn validate_applications(&self) -> Result<()> { + if self.applications.is_empty() { + return Ok(()); + } + + for application in self.worker_map.keys() { + if !self.applications.admits(application) { + return Err(anyhow!( + "worker_map application '{application}' is not in applications" + )); + } + } + Ok(()) } @@ -1253,6 +1318,91 @@ mod tests { }); } + #[test] + fn test_applications_derived_from_worker_map_in_push_mode() { + Jail::expect_with(|jail| { + jail.set_env("TASKBROKER_DELIVERY_MODE", "push"); + jail.set_env( + "TASKBROKER_WORKER_MAP", + "{launchpad=http://127.0.0.1:50052}", + ); + + let args = Args { + run: Run::Broker, + config: None, + }; + let config = Config::from_args(&args).unwrap(); + + assert!(config.applications.admits("launchpad")); + assert!(!config.applications.admits("sentry")); + + Ok(()) + }); + } + + /// Pull pools are not given a `worker_map`, so there is nothing to derive from + /// and everything stays admitted until the pool opts in. + #[test] + fn test_applications_unset_in_pull_mode_admits_everything() { + Jail::expect_with(|_jail| { + let args = Args { + run: Run::Broker, + config: None, + }; + let config = Config::from_args(&args).unwrap(); + + assert_eq!(config.delivery_mode, DeliveryMode::Pull); + assert!(config.applications.is_empty()); + assert!(config.applications.admits("launchpad")); + + Ok(()) + }); + } + + /// `applications` uses the same env list encoding as other sequences. + #[test] + fn test_applications_from_env() { + Jail::expect_with(|jail| { + jail.set_env("TASKBROKER_APPLICATIONS", "[sentry,launchpad]"); + + let args = Args { + run: Run::Broker, + config: None, + }; + let config = Config::from_args(&args).unwrap(); + + assert!(config.applications.admits("sentry")); + assert!(config.applications.admits("launchpad")); + assert!(!config.applications.admits("seer")); + + Ok(()) + }); + } + + #[test] + fn test_applications_rejects_unreachable_worker_map_entry() { + Jail::expect_with(|jail| { + jail.create_file( + "config.yaml", + r#" + applications: + - launchpad + worker_map: + sentry: http://sentry:50052 + "#, + )?; + + let args = Args { + run: Run::Broker, + config: Some("config.yaml".to_owned()), + }; + + assert!(Config::from_args(&args).is_err()); + + Ok(()) + }); + } + #[test] fn test_kafka_consumer_config() { let args = Args { diff --git a/src/grpc/server.rs b/src/grpc/server.rs index ef39c507..8ae404ac 100644 --- a/src/grpc/server.rs +++ b/src/grpc/server.rs @@ -43,6 +43,16 @@ impl ConsumerService for TaskbrokerServer { let application = &request.get_ref().application; let namespace = &request.get_ref().namespace; + // Without this a worker pointed at the wrong pool polls an empty queue + // forever instead of failing. + if let Some(application) = application + && !self.config.applications.admits(application) + { + return Err(Status::failed_precondition(format!( + "Broker does not serve application '{application}'" + ))); + } + let inflight = self .store .claim_activation_for_pull(application.as_deref(), namespace.as_deref()) diff --git a/src/grpc/server_tests.rs b/src/grpc/server_tests.rs index 720c6a7a..833cb0b5 100644 --- a/src/grpc/server_tests.rs +++ b/src/grpc/server_tests.rs @@ -202,6 +202,32 @@ async fn test_get_task_with_application_success(#[case] adapter: &str) { assert_eq!(task.application, Some("hammers".into())); } +#[tokio::test] +async fn test_get_task_unserved_application_returns_failed_precondition() { + let store = create_test_store("sqlite").await; + let mut config = Config { + applications: ["sentry".to_owned()].into_iter().collect(), + ..Config::default() + }; + config.normalize_and_validate().unwrap(); + + let service = TaskbrokerServer { + store, + config: Arc::new(config), + update_tx: None, + }; + + let request = GetTaskRequest { + namespace: None, + application: Some("hammers".into()), + }; + + let response = service.get_task(Request::new(request)).await; + + assert!(response.is_err()); + assert_eq!(response.unwrap_err().code(), Code::FailedPrecondition); +} + #[tokio::test] #[rstest] #[case::sqlite("sqlite")] diff --git a/src/kafka/activation_batcher.rs b/src/kafka/activation_batcher.rs index aa205125..0214de79 100644 --- a/src/kafka/activation_batcher.rs +++ b/src/kafka/activation_batcher.rs @@ -1,3 +1,4 @@ +use std::collections::BTreeMap; use std::mem::replace; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -7,8 +8,9 @@ use futures::future::join_all; use rdkafka::config::ClientConfig; use rdkafka::producer::{FutureProducer, FutureRecord}; use rdkafka::util::Timeout; +use tracing::error; -use crate::config::Config; +use crate::config::{Applications, Config}; use crate::runtime_config::RuntimeConfigManager; use crate::store::activation::Activation; @@ -21,6 +23,8 @@ pub struct ActivationBatcherConfig { pub producer_config: ClientConfig, pub kafka_topic: String, pub kafka_long_topic: String, + pub kafka_deadletter_topic: String, + pub applications: Applications, pub send_timeout_ms: u64, pub max_batch_time_ms: u64, pub max_batch_len: usize, @@ -36,6 +40,8 @@ impl ActivationBatcherConfig { producer_config: config.kafka_producer_config(), kafka_topic: topic_name.to_owned(), kafka_long_topic: config.kafka_long_topic.clone(), + kafka_deadletter_topic: config.kafka_deadletter_topic.clone(), + applications: config.applications.clone(), send_timeout_ms: config.kafka_send_timeout_ms, max_batch_time_ms: config.store.insert_batch_max_time_ms, max_batch_len: config.store.insert_batch_max_length, @@ -47,7 +53,8 @@ impl ActivationBatcherConfig { pub struct ActivationBatcher { batch: Vec, batch_size: usize, - forward_batch: Vec>, // payload + forward_batch: Vec>, // payload + deadletter_batch: Vec<(String, Vec)>, // application, payload config: ActivationBatcherConfig, runtime_config_manager: Arc, producer: Arc, @@ -74,12 +81,59 @@ impl ActivationBatcher { batch: Vec::with_capacity(config.max_batch_len), batch_size: 0, forward_batch: Vec::with_capacity(config.max_batch_len), + deadletter_batch: Vec::new(), config, runtime_config_manager, producer, producer_cluster, } } + + /// The cluster the deadletter topic lives on. The forwarding producer + /// authenticates against it, so it is also the forwarding default. + fn deadletter_cluster(&self) -> String { + self.config + .producer_config + .get("bootstrap.servers") + .expect("producer config always sets bootstrap.servers") + .to_string() + } + + /// The producer, repointed when the target cluster changed. + fn producer_for(&mut self, cluster: String) -> Arc { + if self.producer_cluster != cluster { + let mut config = self.config.producer_config.clone(); + config.set("bootstrap.servers", &cluster); + self.producer = Arc::new( + config + .create() + .expect("Could not create kafka producer in activation batcher"), + ); + self.producer_cluster = cluster; + } + + self.producer.clone() + } + + /// Produce every payload to `topic`, returning (attempts, successes). + async fn send_all<'a>( + &self, + producer: &FutureProducer, + topic: &str, + payloads: impl Iterator>, + ) -> (usize, usize) { + let sends = payloads.map(|payload| { + producer.send( + FutureRecord::<(), Vec>::to(topic).payload(payload), + Timeout::After(Duration::from_millis(self.config.send_timeout_ms)), + ) + }); + + let results = join_all(sends).await; + let successes = results.iter().filter(|r| r.is_ok()).count(); + + (results.len(), successes) + } } impl Reducer for ActivationBatcher { @@ -137,6 +191,18 @@ impl Reducer for ActivationBatcher { return Ok(()); } + // A pool can only run tasks for the applications its workers serve. Anything + // else would sit in the store until it paged someone, so discard it here. + if !self.config.applications.admits(&t.application) { + metrics::counter!( + "filter.unknown_application", + "topic" => self.config.kafka_topic.clone(), + ) + .increment(1); + self.deadletter_batch.push((t.application, t.activation)); + return Ok(()); + } + if runtime_config.demoted_namespaces.contains(namespace) { if forward_topic == self.config.kafka_topic { metrics::counter!( @@ -166,7 +232,10 @@ impl Reducer for ActivationBatcher { } async fn flush(&mut self) -> Result, anyhow::Error> { - if self.batch.is_empty() && self.forward_batch.is_empty() { + if self.batch.is_empty() + && self.forward_batch.is_empty() + && self.deadletter_batch.is_empty() + { return Ok(None); } @@ -181,52 +250,64 @@ impl Reducer for ActivationBatcher { // The forwarding producer authenticates against the deadletter // cluster, so default demoted forwarding there too (and consistently // with upkeep) when no demoted_topic_cluster is configured. - let forward_cluster = - runtime_config - .demoted_topic_cluster - .clone() - .unwrap_or_else(|| { - self.config - .producer_config - .get("bootstrap.servers") - .expect("producer config always sets bootstrap.servers") - .to_string() - }); - if self.producer_cluster != forward_cluster { - let mut new_config = self.config.producer_config.clone(); - new_config.set("bootstrap.servers", &forward_cluster); - self.producer = Arc::new( - new_config - .create() - .expect("Could not create kafka producer in activation batcher"), - ); - self.producer_cluster = forward_cluster; - } + let forward_cluster = runtime_config + .demoted_topic_cluster + .clone() + .unwrap_or_else(|| self.deadletter_cluster()); + let producer = self.producer_for(forward_cluster); let forward_topic = runtime_config .demoted_topic .clone() .unwrap_or(self.config.kafka_long_topic.clone()); - let sends = self.forward_batch.iter().map(|payload| { - self.producer.send( - FutureRecord::<(), Vec>::to(&forward_topic).payload(payload), - Timeout::After(Duration::from_millis(self.config.send_timeout_ms)), - ) - }); - let results = join_all(sends).await; - let success_count = results.iter().filter(|r| r.is_ok()).count(); + let (attempts, successes) = self + .send_all(&producer, &forward_topic, self.forward_batch.iter()) + .await; let topic = self.config.kafka_topic.clone(); metrics::histogram!("consumer.forward_attempts", "topic" => topic.clone()) - .record(results.len() as f64); + .record(attempts as f64); metrics::histogram!("consumer.forward_successes", "topic" => topic.clone()) - .record(success_count as f64); + .record(successes as f64); metrics::histogram!("consumer.forward_failures", "topic" => topic) - .record((results.len() - success_count) as f64); + .record((attempts - successes) as f64); self.forward_batch.clear(); } + if !self.deadletter_batch.is_empty() { + let mut counts: BTreeMap<&str, usize> = BTreeMap::new(); + for (application, _) in &self.deadletter_batch { + *counts.entry(application.as_str()).or_default() += 1; + } + error!( + topic = self.config.kafka_topic.as_str(), + ?counts, + "Discarding activations for applications this pool does not serve" + ); + + let deadletter_cluster = self.deadletter_cluster(); + let producer = self.producer_for(deadletter_cluster); + let deadletter_topic = self.config.kafka_deadletter_topic.clone(); + let (attempts, successes) = self + .send_all( + &producer, + &deadletter_topic, + self.deadletter_batch.iter().map(|(_, payload)| payload), + ) + .await; + + let topic = self.config.kafka_topic.clone(); + metrics::histogram!("consumer.unknown_application_attempts", "topic" => topic.clone()) + .record(attempts as f64); + metrics::histogram!("consumer.unknown_application_successes", "topic" => topic.clone()) + .record(successes as f64); + metrics::histogram!("consumer.unknown_application_failures", "topic" => topic) + .record((attempts - successes) as f64); + + self.deadletter_batch.clear(); + } + self.batch_size = 0; Ok(Some(replace( @@ -238,6 +319,7 @@ impl Reducer for ActivationBatcher { fn reset(&mut self) { self.batch_size = 0; self.forward_batch.clear(); + self.deadletter_batch.clear(); self.batch.clear(); } @@ -387,6 +469,70 @@ demoted_namespaces: []"#; assert_eq!(batcher.batch.len(), 0); } + #[tokio::test] + async fn test_discard_task_from_unserved_application() { + let runtime_config = Arc::new(RuntimeConfigManager::new(None).await); + let mut config = Config { + applications: ["sentry".to_owned()].into_iter().collect(), + ..Default::default() + }; + config.normalize_and_validate().unwrap(); + let config = Arc::new(config); + let mut batcher = ActivationBatcher::new( + ActivationBatcherConfig::from_topic(&config, config.consumable_topics().unwrap()[0].0), + runtime_config, + ); + + let namespace = generate_unique_namespace(); + + let unserved = ActivationBuilder::new() + .id("0") + .taskname("taskname") + .namespace(&namespace) + .application("launchpad") + .build(TaskActivationBuilder::new()); + + batcher.reduce(unserved).await.unwrap(); + assert_eq!(batcher.batch.len(), 0); + assert_eq!(batcher.deadletter_batch.len(), 1); + assert_eq!(batcher.deadletter_batch[0].0, "launchpad"); + + let served = ActivationBuilder::new() + .id("1") + .taskname("taskname") + .namespace(&namespace) + .application("sentry") + .build(TaskActivationBuilder::new()); + + batcher.reduce(served).await.unwrap(); + assert_eq!(batcher.batch.len(), 1); + assert_eq!(batcher.deadletter_batch.len(), 1); + } + + /// A pool that has not opted in keeps running every application it is sent. + #[tokio::test] + async fn test_unset_applications_keeps_every_task() { + let runtime_config = Arc::new(RuntimeConfigManager::new(None).await); + let mut config = Config::default(); + config.normalize_and_validate().unwrap(); + let config = Arc::new(config); + let mut batcher = ActivationBatcher::new( + ActivationBatcherConfig::from_topic(&config, config.consumable_topics().unwrap()[0].0), + runtime_config, + ); + + let activation = ActivationBuilder::new() + .id("0") + .taskname("taskname") + .namespace(generate_unique_namespace()) + .application("launchpad") + .build(TaskActivationBuilder::new()); + + batcher.reduce(activation).await.unwrap(); + assert_eq!(batcher.batch.len(), 1); + assert_eq!(batcher.deadletter_batch.len(), 0); + } + #[tokio::test] async fn test_close_by_bytes_limit() { let runtime_config = Arc::new(RuntimeConfigManager::new(None).await); diff --git a/src/push/tests.rs b/src/push/tests.rs index 93e40d48..1ba624e6 100644 --- a/src/push/tests.rs +++ b/src/push/tests.rs @@ -14,7 +14,7 @@ use crate::store::activation::{Activation, ActivationStatus}; use crate::store::traits::ActivationStore; use crate::store::types::{FailedTasksForwarder, TopicPartition}; use crate::test_utils::make_activations; -use crate::worker::test_worker_map; +use crate::worker::{WorkerMap, test_worker_map}; use super::*; @@ -22,12 +22,17 @@ use super::*; #[derive(Default, Clone)] struct MockStore { marked_processing: Arc>>, + status_updates: Arc>>, } impl MockStore { fn marked_ids(&self) -> Vec { self.marked_processing.lock().unwrap().clone() } + + fn recorded_statuses(&self) -> Vec<(String, ActivationStatus)> { + self.status_updates.lock().unwrap().clone() + } } #[async_trait] @@ -72,11 +77,16 @@ impl ActivationStore for MockStore { async fn set_status( &self, - _id: &str, - _status: ActivationStatus, + id: &str, + status: ActivationStatus, _max_attempts: Option, _delay_on_retry: Option, ) -> Result> { + self.status_updates + .lock() + .unwrap() + .push((id.to_string(), status)); + Ok(None) } @@ -335,3 +345,50 @@ async fn push_pool_start_does_not_mark_activation_processing_on_push_failure() { "mark_activation_processing should not be called when push fails" ); } + +/// An activation the pool has no worker for must reach a terminal status, otherwise +/// the claim expires and it is fetched again forever. +#[tokio::test] +async fn push_pool_start_fails_activation_without_worker_mapping() { + let config = Arc::new(Config { + worker_map: [("sentry".into(), "unused".into())].into(), + push: PushConfig { + threads: 1, + queue: PushQueueConfig { + size: 10, + ..Default::default() + }, + ..Default::default() + }, + ..Default::default() + }); + + let store = Arc::new(MockStore::default()); + let (sender, receiver) = flume::bounded(config.push.queue.size); + let pool = Arc::new(PushPool::new(receiver, config)); + + let updater = test_eager_updater(store.clone()); + + tokio::spawn({ + let store = store.clone(); + async move { + pool.start(vec![WorkerMap::new()], updater, store.clone()) + .await + } + }); + + let activation = make_activations(1).remove(0); + let id = activation.id.clone(); + + sender + .send_async((activation, Instant::now())) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + + assert_eq!( + store.recorded_statuses(), + vec![(id, ActivationStatus::Failure)] + ); + assert!(store.marked_ids().is_empty()); +} diff --git a/src/push/thread.rs b/src/push/thread.rs index e75b8318..acc31886 100644 --- a/src/push/thread.rs +++ b/src/push/thread.rs @@ -68,6 +68,20 @@ impl PushThread { "Application has no worker mapping" ); + // Leaving the claim in place would revert to pending and be fetched + // again forever, so fail the task and let upkeep discard it. + if let Err(e) = self + .store + .set_status(id, ActivationStatus::Failure, None, None) + .await + { + error!( + task_id = %id, + error = ?e, + "Failed to fail activation without a worker mapping" + ); + } + return; }; From 275fb4b47db9e659c37af85905294182a24f4a9f Mon Sep 17 00:00:00 2001 From: ShashSriv Date: Thu, 1 Oct 2026 16:00:04 -0700 Subject: [PATCH 2/5] test(consumer): Cover the deadletter publish for discarded activations The existing test stops at `reduce`, so the publish in `flush` was never exercised. Produce a discarded activation through a real broker and consume it back off the pool's deadletter topic, following the pattern upkeep uses for its own deadletter test. Co-Authored-By: Claude Opus 5 --- src/kafka/activation_batcher.rs | 54 ++++++++++++++++++++++++++++++++- 1 file changed, 53 insertions(+), 1 deletion(-) diff --git a/src/kafka/activation_batcher.rs b/src/kafka/activation_batcher.rs index 0214de79..0271db5f 100644 --- a/src/kafka/activation_batcher.rs +++ b/src/kafka/activation_batcher.rs @@ -344,11 +344,17 @@ mod tests { use std::sync::Arc; use chrono::Utc; + use prost::Message; + use sentry_protos::taskbroker::v1::TaskActivation; use tempfile::NamedTempFile; + use crate::config::deprecated::DeprecatedConfig; use crate::config::store::StoreConfig; use crate::store::activation::ActivationBuilder; - use crate::test_utils::{TaskActivationBuilder, generate_unique_namespace}; + use crate::test_utils::{ + TaskActivationBuilder, consume_topic, create_integration_config_from_base, + generate_unique_namespace, reset_topic, + }; use super::{ ActivationBatcher, ActivationBatcherConfig, Config, Reducer, RuntimeConfigManager, @@ -509,6 +515,52 @@ demoted_namespaces: []"#; assert_eq!(batcher.deadletter_batch.len(), 1); } + #[tokio::test] + async fn test_discarded_task_is_published_to_deadletter_topic() { + let config = Arc::new(create_integration_config_from_base(Config { + deprecated: DeprecatedConfig { + kafka_topic: Some("taskbroker-test-unknown-application".into()), + ..DeprecatedConfig::default() + }, + kafka_deadletter_topic: "taskbroker-test-unknown-application-dlq".into(), + applications: ["sentry".to_owned()].into_iter().collect(), + ..Default::default() + })); + reset_topic(config.clone()).await; + + let runtime_config = Arc::new(RuntimeConfigManager::new(None).await); + let mut batcher = ActivationBatcher::new( + ActivationBatcherConfig::from_topic(&config, config.consumable_topics().unwrap()[0].0), + runtime_config, + ); + + let unserved = ActivationBuilder::new() + .id("0") + .taskname("taskname") + .namespace(generate_unique_namespace()) + .application("launchpad") + .build(TaskActivationBuilder::new()); + let expected = TaskActivation::decode(&unserved.activation as &[u8]).unwrap(); + + batcher.reduce(unserved).await.unwrap(); + assert_eq!(batcher.deadletter_batch.len(), 1); + + let flushed = batcher + .flush() + .await + .unwrap() + .expect("flush produced a batch"); + assert!(flushed.is_empty()); + assert_eq!(batcher.deadletter_batch.len(), 0); + + let messages = + consume_topic(config.clone(), config.kafka_deadletter_topic.as_ref(), 1).await; + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].id, expected.id); + assert_eq!(messages[0].application, Some("launchpad".to_owned())); + assert_eq!(messages[0].parameters_bytes, expected.parameters_bytes); + } + /// A pool that has not opted in keeps running every application it is sent. #[tokio::test] async fn test_unset_applications_keeps_every_task() { From 137904b0399933bda0c7507484e9a0ea07c4320e Mon Sep 17 00:00:00 2001 From: ShashSriv Date: Fri, 2 Oct 2026 11:02:52 -0700 Subject: [PATCH 3/5] ref(consumer): Trim comments on the application check Co-Authored-By: Claude Opus 5 --- src/config/mod.rs | 6 ++---- src/grpc/server.rs | 3 +-- src/kafka/activation_batcher.rs | 3 +-- src/push/thread.rs | 3 +-- 4 files changed, 5 insertions(+), 10 deletions(-) diff --git a/src/config/mod.rs b/src/config/mod.rs index 60b7dff5..1a884fe3 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -844,10 +844,8 @@ impl Config { Ok(()) } - /// Push pools already declare every application they can deliver to in - /// `worker_map`, so they never have to repeat it. Pull pools have no such - /// list, and an empty `applications` admits everything, so a pool that - /// hasn't opted in keeps running rather than discarding all of its work. + /// Push pools derive the set from `worker_map`. Pull pools have no such list and + /// are left empty, which admits everything. fn normalize_applications(&mut self) { if !self.applications.is_empty() { return; diff --git a/src/grpc/server.rs b/src/grpc/server.rs index 8ae404ac..f07f4305 100644 --- a/src/grpc/server.rs +++ b/src/grpc/server.rs @@ -43,8 +43,7 @@ impl ConsumerService for TaskbrokerServer { let application = &request.get_ref().application; let namespace = &request.get_ref().namespace; - // Without this a worker pointed at the wrong pool polls an empty queue - // forever instead of failing. + // Worker pointed at the wrong pool fails if let Some(application) = application && !self.config.applications.admits(application) { diff --git a/src/kafka/activation_batcher.rs b/src/kafka/activation_batcher.rs index 0271db5f..72636097 100644 --- a/src/kafka/activation_batcher.rs +++ b/src/kafka/activation_batcher.rs @@ -191,8 +191,7 @@ impl Reducer for ActivationBatcher { return Ok(()); } - // A pool can only run tasks for the applications its workers serve. Anything - // else would sit in the store until it paged someone, so discard it here. + // Discard tasks for applications this pool's workers do not serve. if !self.config.applications.admits(&t.application) { metrics::counter!( "filter.unknown_application", diff --git a/src/push/thread.rs b/src/push/thread.rs index acc31886..0f46bac7 100644 --- a/src/push/thread.rs +++ b/src/push/thread.rs @@ -68,8 +68,7 @@ impl PushThread { "Application has no worker mapping" ); - // Leaving the claim in place would revert to pending and be fetched - // again forever, so fail the task and let upkeep discard it. + // Fail the task so upkeep discards it instead of the claim expiring. if let Err(e) = self .store .set_status(id, ActivationStatus::Failure, None, None) From 11f866e2b1abcba553f5d283c768d07f8913b60d Mon Sep 17 00:00:00 2001 From: ShashSriv Date: Fri, 2 Oct 2026 14:31:08 -0700 Subject: [PATCH 4/5] fix(consumer): Address review findings on the application check - `worker_map` is only delivered through in push mode, and `from_args` gives pull pools a default entry they never read, so the subset check rejected a pull pool opting in to any other application. Scope the check to push. - `set_task_status` claims through the same path as `get_task`, so apply the check there too. It returns no task rather than an error, because an error makes the worker re-send the status update. - Discards and forwards bypass `batch`, so `is_full` never tripped on them and a misrouted topic only flushed on the timer. Co-Authored-By: Claude Opus 5 --- src/config/mod.rs | 28 +++++++++++++++++++--- src/grpc/server.rs | 8 +++++++ src/grpc/server_tests.rs | 41 +++++++++++++++++++++++++++++++++ src/kafka/activation_batcher.rs | 37 +++++++++++++++++++++++++++++ 4 files changed, 111 insertions(+), 3 deletions(-) diff --git a/src/config/mod.rs b/src/config/mod.rs index 1a884fe3..3c5289f3 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -858,10 +858,11 @@ impl Config { } } - /// A `worker_map` entry outside `applications` is unreachable: its activations - /// are discarded before they are ever pushed to that worker. + /// A `worker_map` entry outside `applications` is unreachable in push mode. Pull + /// pools never deliver through `worker_map` and are given a default entry from + /// `from_args` that they never read, so the check does not apply to them. fn validate_applications(&self) -> Result<()> { - if self.applications.is_empty() { + if self.applications.is_empty() || self.delivery_mode != DeliveryMode::Push { return Ok(()); } @@ -1377,12 +1378,33 @@ mod tests { }); } + /// `from_args` gives pull pools a default `worker_map` they never deliver through, + /// which must not block them from opting in to an unrelated application. + #[test] + fn test_pull_applications_ignore_default_worker_map() { + Jail::expect_with(|jail| { + jail.create_file("config.yaml", "applications:\n - launchpad\n")?; + + let args = Args { + run: Run::Broker, + config: Some("config.yaml".to_owned()), + }; + let config = Config::from_args(&args).unwrap(); + + assert!(config.applications.admits("launchpad")); + assert!(!config.applications.admits("sentry")); + + Ok(()) + }); + } + #[test] fn test_applications_rejects_unreachable_worker_map_entry() { Jail::expect_with(|jail| { jail.create_file( "config.yaml", r#" + delivery_mode: push applications: - launchpad worker_map: diff --git a/src/grpc/server.rs b/src/grpc/server.rs index f07f4305..331a31ad 100644 --- a/src/grpc/server.rs +++ b/src/grpc/server.rs @@ -186,6 +186,14 @@ impl ConsumerService for TaskbrokerServer { return Ok(Response::new(SetTaskStatusResponse { task: None })); }; + // Same check as get_task. Returning an error here would make the worker + // re-send the status update, so an unserved application just gets no task. + if let Some(application) = application + && !self.config.applications.admits(application) + { + return Ok(Response::new(SetTaskStatusResponse { task: None })); + } + let start_time = Instant::now(); let res = match self diff --git a/src/grpc/server_tests.rs b/src/grpc/server_tests.rs index 833cb0b5..3a6bafda 100644 --- a/src/grpc/server_tests.rs +++ b/src/grpc/server_tests.rs @@ -228,6 +228,47 @@ async fn test_get_task_unserved_application_returns_failed_precondition() { assert_eq!(response.unwrap_err().code(), Code::FailedPrecondition); } +/// A leftover activation for an unserved application must not be handed out on the +/// fetch-next path either. +#[tokio::test] +#[allow(deprecated)] +async fn test_fetch_next_task_skips_unserved_application() { + let store = create_test_store("sqlite").await; + let mut config = Config { + applications: ["sentry".to_owned()].into_iter().collect(), + ..Config::default() + }; + config.normalize_and_validate().unwrap(); + + let mut activations = make_activations(2); + let mut payload = TaskActivation::decode(&activations[1].activation as &[u8]).unwrap(); + payload.application = Some("launchpad".into()); + activations[1].activation = payload.encode_to_vec(); + activations[1].application = "launchpad".into(); + store.store(&activations).await.unwrap(); + + let service = TaskbrokerServer { + store, + config: Arc::new(config), + update_tx: None, + }; + + let request = SetTaskStatusRequest { + id: "id_0".to_string(), + status: 5, // Complete + fetch_next_task: Some(FetchNextTask { + namespace: None, + application: Some("launchpad".into()), + }), + max_attempts: None, + delay_on_retry: None, + }; + + let response = service.set_task_status(Request::new(request)).await; + assert!(response.is_ok()); + assert!(response.unwrap().get_ref().task.is_none()); +} + #[tokio::test] #[rstest] #[case::sqlite("sqlite")] diff --git a/src/kafka/activation_batcher.rs b/src/kafka/activation_batcher.rs index 7e382916..d45c4766 100644 --- a/src/kafka/activation_batcher.rs +++ b/src/kafka/activation_batcher.rs @@ -344,6 +344,8 @@ impl Reducer for ActivationBatcher { async fn is_full(&self) -> bool { self.batch.len() >= self.config.max_batch_len + || self.forward_batch.len() >= self.config.max_batch_len + || self.deadletter_batch.len() >= self.config.max_batch_len || self.batch_size >= self.config.max_batch_size } @@ -600,6 +602,41 @@ use_arroyo_producer: {use_arroyo_producer}"# assert_eq!(messages[0].parameters_bytes, expected.parameters_bytes); } + /// Discards bypass `batch`, so without them `is_full` never trips and a misrouted + /// topic only flushes on the timer. + #[tokio::test] + async fn test_deadletter_batch_counts_towards_full() { + let runtime_config = Arc::new(RuntimeConfigManager::new(None).await); + let mut config = Config { + applications: ["sentry".to_owned()].into_iter().collect(), + store: StoreConfig { + insert_batch_max_length: 2, + ..StoreConfig::default() + }, + ..Default::default() + }; + config.normalize_and_validate().unwrap(); + let config = Arc::new(config); + let mut batcher = ActivationBatcher::new( + ActivationBatcherConfig::from_topic(&config, config.consumable_topics().unwrap()[0].0), + runtime_config, + ); + + let namespace = generate_unique_namespace(); + for id in 0..2 { + let unserved = ActivationBuilder::new() + .id(id.to_string()) + .taskname("taskname") + .namespace(&namespace) + .application("launchpad") + .build(TaskActivationBuilder::new()); + batcher.reduce(unserved).await.unwrap(); + } + + assert!(batcher.batch.is_empty()); + assert!(batcher.is_full().await); + } + /// A pool that has not opted in keeps running every application it is sent. #[tokio::test] async fn test_unset_applications_keeps_every_task() { From 8ca96db4962c268dffa6f8c0d37f6105335ea55c Mon Sep 17 00:00:00 2001 From: ShashSriv Date: Mon, 5 Oct 2026 09:58:47 -0700 Subject: [PATCH 5/5] fix(grpc): Count and log requests for unserved applications The fetch-next path returned no task without any signal, so a worker pointed at the wrong pool looked identical to an idle queue. Both claim paths now emit `grpc_server.unserved_application` (tagged by method) and log the application. Co-Authored-By: Claude Opus 5 --- src/grpc/server.rs | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/src/grpc/server.rs b/src/grpc/server.rs index 331a31ad..1b39025f 100644 --- a/src/grpc/server.rs +++ b/src/grpc/server.rs @@ -47,6 +47,16 @@ impl ConsumerService for TaskbrokerServer { if let Some(application) = application && !self.config.applications.admits(application) { + metrics::counter!( + "grpc_server.unserved_application", + "method" => "get_task", + ) + .increment(1); + error!( + application, + "Requested application is not served by this pool" + ); + return Err(Status::failed_precondition(format!( "Broker does not serve application '{application}'" ))); @@ -191,6 +201,16 @@ impl ConsumerService for TaskbrokerServer { if let Some(application) = application && !self.config.applications.admits(application) { + metrics::counter!( + "grpc_server.unserved_application", + "method" => "fetch_next_task", + ) + .increment(1); + error!( + application, + "Requested application is not served by this pool" + ); + return Ok(Response::new(SetTaskStatusResponse { task: None })); }