Shash sriv/stream 1807 discard tasks from incorrect applications

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<String>);
+
+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<String> for Applications {
+    fn from_iter<T: IntoIterator<Item = String>>(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<String, String>,
 
+    /// 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<String>,
 
@@ -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,42 @@ impl Config {
             }
         }
 
+        self.normalize_applications();
+        self.validate_applications()?;
+
+        Ok(())
+    }
+
+    /// 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;
+        }
+
+        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 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() || self.delivery_mode != DeliveryMode::Push {
+            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 +1317,112 @@ 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(())
+        });
+    }
+
+    /// `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:
+                    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 {

src/grpc/server.rs

@@ -43,6 +43,25 @@ impl ConsumerService for TaskbrokerServer {
         let application = &request.get_ref().application;
         let namespace = &request.get_ref().namespace;
 
+        // Worker pointed at the wrong pool fails
+        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}'"
+            )));
+        }
+
         let inflight = self
             .store
             .claim_activation_for_pull(application.as_deref(), namespace.as_deref())
@@ -177,6 +196,24 @@ 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)
+        {
+            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 }));
+        }
+
         let start_time = Instant::now();
 
         let res = match self

src/grpc/server_tests.rs

@@ -202,6 +202,73 @@ 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);
+}
+
+/// 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")]