feat(taskbroker): Add Application to Worker Pool Mapping (#604)

* Add Application to Worker Mapping

* Add Counter for Missing Worker Mapping

* Add Missing Worker Mapping Metric to Push Channel Drain

* Add Application (Worker Pool) Tags to Metrics, Change `getsentry` to `sentry`

* Add Link Back to Comment on Signature Helper

* Fix Tests Hopefully

* Add Application Tag to Drain Loop Missing Mapping Metric

* Try Fixing Stuck Unit Test Due to Shared DLQ Topic

* Fix Lint

Author: george-sentry

src/config.rs

@@ -290,20 +290,14 @@ pub struct Config {
     /// Maximum time in milliseconds for a single push RPC to the worker service. This should be greater than the worker's internal timeout.
     pub push_timeout_ms: u64,
 
-    /// The worker service endpoint.
-    pub worker_endpoint: String,
-
     /// The hostname used to construct `callback_url` for task push requests.
     pub callback_addr: String,
 
     /// The port used to construct `callback_url` for task push requests.
     pub callback_port: u32,
 
-    /// Application filter for push mode. When set, only pending activations for this application are considered.
-    pub application: Option<String>,
-
-    /// List of namespaces for push mode. When set, application must also be set (store requirement).
-    pub namespaces: Option<Vec<String>>,
+    /// Maps every application to its worker endpoint, both represented as strings.
+    pub worker_map: BTreeMap<String, String>,
 }
 
 impl Default for Config {
@@ -385,11 +379,9 @@ impl Default for Config {
             push_queue_size: 1,
             push_queue_timeout_ms: 5000,
             push_timeout_ms: 30000,
-            worker_endpoint: "http://127.0.0.1:50052".into(),
             callback_addr: "0.0.0.0".into(),
             callback_port: 50051,
-            application: None,
-            namespaces: None,
+            worker_map: [("sentry".into(), "http://127.0.0.1:50052".into())].into(),
         }
     }
 }
@@ -527,6 +519,10 @@ mod tests {
         assert_eq!(config.max_pending_count, 2048);
         assert_eq!(config.max_processing_count, 2048);
         assert_eq!(config.vacuum_page_count, None);
+        assert_eq!(
+            config.worker_map.get("sentry").map(String::as_str),
+            Some("http://127.0.0.1:50052")
+        );
     }
 
     #[test]
@@ -554,6 +550,9 @@ mod tests {
                 max_processing_attempts: 5
                 vacuum_page_count: 1000
                 full_vacuum_on_start: true
+                worker_map:
+                    sentry: http://worker-sentry:50052
+                    launchpad: http://worker-launchpad:50053
             "#,
             )?;
             // Env vars always override config file
@@ -588,6 +587,16 @@ mod tests {
             assert_eq!(config.vacuum_page_count, Some(1000));
             assert_eq!(config.db_max_size, Some(3_000_000_000));
             assert!(config.full_vacuum_on_start);
+            assert_eq!(
+                config.worker_map,
+                BTreeMap::from([
+                    ("sentry".to_owned(), "http://worker-sentry:50052".to_owned(),),
+                    (
+                        "launchpad".to_owned(),
+                        "http://worker-launchpad:50053".to_owned(),
+                    ),
+                ])
+            );
 
             Ok(())
         });
@@ -632,6 +641,35 @@ mod tests {
                 config.default_metrics_tags,
                 BTreeMap::from([("key".to_owned(), "value".to_owned())])
             );
+            assert_eq!(
+                config.worker_map.get("sentry").map(String::as_str),
+                Some("http://127.0.0.1:50052"),
+                "partial env override must not drop worker_map defaults"
+            );
+
+            Ok(())
+        });
+    }
+
+    /// `worker_map` uses the same env map encoding as `default_metrics_tags` (brace `key=value` pairs).
+    #[test]
+    fn test_worker_map_from_env() {
+        Jail::expect_with(|jail| {
+            jail.set_env("TASKBROKER_LOG_FILTER", "error");
+            jail.set_env(
+                "TASKBROKER_WORKER_MAP",
+                "{sentry=http://127.0.0.1:60052,launchpad=http://127.0.0.1:60053}",
+            );
+
+            let args = Args { config: None };
+            let config = Config::from_args(&args).unwrap();
+            assert_eq!(
+                config.worker_map,
+                BTreeMap::from([
+                    ("sentry".to_owned(), "http://127.0.0.1:60052".to_owned(),),
+                    ("launchpad".to_owned(), "http://127.0.0.1:60053".to_owned(),),
+                ])
+            );
 
             Ok(())
         });
@@ -876,49 +914,4 @@ mod tests {
             Ok(())
         });
     }
-
-    #[test]
-    fn test_default_application_and_namespaces() {
-        let config = Config::default();
-        assert_eq!(config.application, None);
-        assert_eq!(config.namespaces, None);
-    }
-
-    #[test]
-    fn test_from_args_application_from_env() {
-        Jail::expect_with(|jail| {
-            jail.set_env("TASKBROKER_APPLICATION", "getsentry");
-
-            let args = Args { config: None };
-            let config = Config::from_args(&args).unwrap();
-            assert_eq!(config.application.as_deref(), Some("getsentry"));
-            assert_eq!(config.namespaces, None);
-
-            Ok(())
-        });
-    }
-
-    #[test]
-    fn test_from_args_application_and_namespaces_from_config_file() {
-        Jail::expect_with(|jail| {
-            jail.create_file(
-                "config.yaml",
-                r#"
-                application: getsentry
-                namespaces:
-                  - ns1
-                  - ns2
-            "#,
-            )?;
-
-            let args = Args {
-                config: Some("config.yaml".to_owned()),
-            };
-            let config = Config::from_args(&args).unwrap();
-            assert_eq!(config.application.as_deref(), Some("getsentry"));
-            assert_eq!(config.namespaces, Some(vec!["ns1".into(), "ns2".into()]));
-
-            Ok(())
-        });
-    }
 }

src/fetch/mod.rs

@@ -115,13 +115,7 @@ impl<T: TaskPusher + Send + Sync + 'static> FetchPool<T> {
                             let start = Instant::now();
                             let mut backoff = false;
 
-                            let application = config.application.as_deref();
-                            let namespaces = config.namespaces.as_deref();
-
-                            match store
-                                .claim_activations_for_push(application, namespaces, limit, bucket)
-                                .await
-                            {
+                            match store.claim_activations_for_push(limit, bucket).await {
                                 Ok(activations) if activations.is_empty() => {
                                     metrics::counter!("fetch.empty").increment(1);
                                     debug!("No pending activations");

src/push/mod.rs

@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 
@@ -113,7 +114,7 @@ impl PushPool {
         let mut push_pool: JoinSet<Result<()>> = crate::tokio::spawn_pool(
             self.config.push_threads,
             |_| {
-                let endpoint = self.config.worker_endpoint.clone();
+                let worker_map = self.config.worker_map.clone();
                 let receiver = self.receiver.clone();
                 let store = store.clone();
 
@@ -130,21 +131,28 @@ impl PushPool {
                 async move {
                     metrics::counter!("push.worker.connect.attempt").increment(1);
 
-                    let mut worker = match WorkerServiceClient::connect(endpoint).await {
-                        Ok(w) => {
-                            metrics::counter!("push.worker.connect", "result" => "ok").increment(1);
-                            w
-                        }
+                    let mut workers = HashMap::new();
 
-                        Err(e) => {
-                            metrics::counter!("push.worker.connect", "result" => "error")
-                                .increment(1);
-                            error!("Failed to connect to worker - {:?}", e);
+                    for (application, endpoint) in worker_map.clone() {
+                        let worker = match WorkerServiceClient::connect(endpoint).await {
+                            Ok(w) => {
+                                metrics::counter!("push.worker.connect", "result" => "ok", "application" => application.clone())
+                                    .increment(1);
+                                w
+                            }
 
-                            // When we fail to connect, the taskbroker will shut down, but this may change in the future
-                            return Err(e.into());
-                        }
-                    };
+                            Err(e) => {
+                                metrics::counter!("push.worker.connect", "result" => "error", "application" => application.clone())
+                                    .increment(1);
+                                error!(error = ?e, "Failed to connect to worker");
+
+                                // When we fail to connect, the taskbroker will shut down, but this may change in the future
+                                return Err(e.into());
+                            }
+                        };
+
+                        workers.insert(application, worker);
+                    }
 
                     loop {
                         tokio::select! {
@@ -165,8 +173,20 @@ impl PushPool {
                                 let id = activation.id.clone();
                                 let callback_url = callback_url.clone();
 
+                                let Some(worker) = workers.get_mut(&activation.application) else {
+                                    metrics::counter!("push.missing_worker_mapping", "application" => activation.application.clone()).increment(1);
+
+                                    error!(
+                                        task_id = %id,
+                                        application = activation.application,
+                                        "Task application has no worker pool mapping"
+                                    );
+
+                                    continue
+                                };
+
                                 match push_task(
-                                    &mut worker,
+                                    worker,
                                     activation,
                                     callback_url,
                                     timeout,
@@ -209,8 +229,20 @@ impl PushPool {
                         let id = activation.id.clone();
                         let callback_url = callback_url.clone();
 
+                        let Some(worker) = workers.get_mut(&activation.application) else {
+                            metrics::counter!("push.missing_worker_mapping", "application" => activation.application.clone()).increment(1);
+
+                            error!(
+                                task_id = %id,
+                                application = activation.application,
+                                "Task application has no worker pool mapping"
+                            );
+
+                            continue;
+                        };
+
                         match push_task(
-                            &mut worker,
+                            worker,
                             activation,
                             callback_url,
                             timeout,

src/store/tests.rs

@@ -265,7 +265,7 @@ async fn test_get_pending_activation_bucket_filter(#[case] adapter: &str) {
     assert!(store.store(batch).await.is_ok());
 
     let mut first = store
-        .claim_activations_for_push(None, None, Some(1), Some((15, 25)))
+        .claim_activations_for_push(Some(1), Some((15, 25)))
         .await
         .unwrap();
     assert_eq!(first.len(), 1);
@@ -274,7 +274,7 @@ async fn test_get_pending_activation_bucket_filter(#[case] adapter: &str) {
     assert_eq!(first.bucket, 20);
 
     let mut second = store
-        .claim_activations_for_push(None, None, Some(1), Some((0, 15)))
+        .claim_activations_for_push(Some(1), Some((0, 15)))
         .await
         .unwrap();
     assert_eq!(second.len(), 1);
@@ -284,7 +284,7 @@ async fn test_get_pending_activation_bucket_filter(#[case] adapter: &str) {
 
     assert!(
         store
-            .claim_activations_for_push(None, None, Some(1), Some((15, 25)))
+            .claim_activations_for_push(Some(1), Some((15, 25)))
             .await
             .unwrap()
             .is_empty()
@@ -606,10 +606,7 @@ async fn test_get_pending_activations_no_limit(#[case] adapter: &str) {
     let batch = make_activations(N as u32);
     assert!(store.store(batch).await.is_ok());
 
-    let got = store
-        .claim_activations_for_push(None, None, None, None)
-        .await
-        .unwrap();
+    let got = store.claim_activations_for_push(None, None).await.unwrap();
     assert_eq!(got.len(), N);
     assert!(
         got.iter()
@@ -641,7 +638,7 @@ async fn test_get_pending_activations_limit_below_pending(#[case] adapter: &str)
     assert!(store.store(batch).await.is_ok());
 
     let got = store
-        .claim_activations_for_push(None, None, Some(X), None)
+        .claim_activations_for_push(Some(X), None)
         .await
         .unwrap();
     assert_eq!(got.len(), X as usize);
@@ -678,7 +675,7 @@ async fn test_get_pending_activations_limit_above_pending(#[case] adapter: &str)
     assert!(store.store(batch).await.is_ok());
 
     let got = store
-        .claim_activations_for_push(None, None, Some(X), None)
+        .claim_activations_for_push(Some(X), None)
         .await
         .unwrap();
     assert_eq!(got.len(), Y);

src/store/traits.rs

@@ -30,22 +30,10 @@ pub trait InflightActivationStore: Send + Sync {
     /// Claims `limit` activations within the `bucket` range. Push mode uses status `Claimed` until `mark_activation_processing` moves to `Processing`.
     async fn claim_activations_for_push(
         &self,
-        application: Option<&str>,
-        namespaces: Option<&[String]>,
         limit: Option<i32>,
         bucket: Option<BucketRange>,
     ) -> Result<Vec<InflightActivation>, Error> {
-        // If a namespace filter is used, an application must also be used
-        if namespaces.is_some() && application.is_none() {
-            warn!(
-                ?namespaces,
-                "Received request for namespaced task without application"
-            );
-
-            return Ok(vec![]);
-        }
-
-        self.claim_activations(application, namespaces, limit, bucket, false)
+        self.claim_activations(None, None, limit, bucket, false)
             .await
     }
 

src/test_utils.rs

@@ -336,8 +336,8 @@ pub fn create_integration_config_with_ssl() -> Arc<Config> {
     Arc::new(config)
 }
 
-pub fn create_integration_config_with_topic(topic: String) -> Arc<Config> {
-    let config = Config {
+pub fn create_integration_config_with_topic(topic: String) -> Config {
+    Config {
         pg_host: get_pg_host(),
         pg_port: get_pg_port(),
         pg_username: get_pg_username(),
@@ -347,9 +347,7 @@ pub fn create_integration_config_with_topic(topic: String) -> Arc<Config> {
         kafka_topic: topic,
         kafka_auto_offset_reset: "earliest".into(),
         ..Config::default()
-    };
-
-    Arc::new(config)
+    }
 }
 
 /// Create a kafka producer for a given config