diff --git a/lib/runtime/src/discovery/kube.rs b/lib/runtime/src/discovery/kube.rs index 89c1e03a838b..40991161742d 100644 --- a/lib/runtime/src/discovery/kube.rs +++ b/lib/runtime/src/discovery/kube.rs @@ -25,6 +25,14 @@ use std::collections::HashSet; use std::sync::Arc; use tokio::sync::RwLock; +fn validate_kubernetes_publisher_id(publisher_id: u64) -> Result<()> { + if i64::try_from(publisher_id).is_err() { + anyhow::bail!("Kubernetes discovery publisher ID {publisher_id} exceeds i64::MAX"); + } + + Ok(()) +} + /// Kubernetes-based discovery client #[derive(Clone)] pub struct KubeDiscoveryClient { @@ -115,6 +123,13 @@ impl Discovery for KubeDiscoveryClient { } async fn register_internal(&self, spec: DiscoverySpec) -> Result { + match &spec { + DiscoverySpec::EventChannel { publisher_id, .. } + | DiscoverySpec::EventSource { publisher_id, .. } => { + validate_kubernetes_publisher_id(*publisher_id)?; + } + _ => {} + } let instance = spec.into_instance(self.instance_id()); let instance_id = instance.instance_id(); @@ -503,3 +518,14 @@ impl Discovery for KubeDiscoveryClient { Ok(Box::pin(stream)) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn publisher_ids_must_fit_kubernetes_integer_range() { + assert!(validate_kubernetes_publisher_id(i64::MAX as u64).is_ok()); + assert!(validate_kubernetes_publisher_id((i64::MAX as u64) + 1).is_err()); + } +} diff --git a/lib/runtime/src/discovery/kube/crd.rs b/lib/runtime/src/discovery/kube/crd.rs index 1be5de2d98af..23d9f1426893 100644 --- a/lib/runtime/src/discovery/kube/crd.rs +++ b/lib/runtime/src/discovery/kube/crd.rs @@ -176,12 +176,19 @@ mod tests { endpoint: endpoint.clone(), }, topic: "kv-events".to_string(), - publisher_id: 205, + publisher_id: i64::MAX as u64, metadata: serde_json::json!({"worker_id": 7, "dp_rank": 0}), }; metadata.register_event_source(source.clone()).unwrap(); let cr = build_cr("test-pod", "test-pod", "pod-uid", &metadata).unwrap(); + let publisher_id = cr.spec.data["event_sources"] + .as_object() + .and_then(|sources| sources.values().next()) + .and_then(|source| source.get("publisher_id")) + .expect("serialized event source publisher ID"); + assert!(publisher_id.is_i64()); + let round_trip: DiscoveryMetadata = serde_json::from_value(cr.spec.data).unwrap(); assert_eq!( diff --git a/lib/runtime/src/transports/event_plane/mod.rs b/lib/runtime/src/transports/event_plane/mod.rs index 1b7fae8543ad..fcc2ac53da80 100644 --- a/lib/runtime/src/transports/event_plane/mod.rs +++ b/lib/runtime/src/transports/event_plane/mod.rs @@ -259,6 +259,11 @@ impl Stream for DeduplicatingStream { } } +/// Keep the shared wire, channel, and source publisher ID representable in Kubernetes metadata. +fn discovery_safe_publisher_id(random_id: u64) -> u64 { + random_id & (i64::MAX as u64) +} + /// Event publisher for a specific topic. pub struct EventPublisher { transport_kind: EventTransportKind, @@ -387,9 +392,11 @@ impl EventPublisher { // can host multiple publishers for the same scope/topic, each with its // own ZMQ endpoint and sequence space, so the process ID is not unique // enough here. - let publisher_id = rand::rngs::OsRng - .try_next_u64() - .map_err(|error| anyhow::anyhow!("failed to generate publisher ID: {error}"))?; + let publisher_id = discovery_safe_publisher_id( + rand::rngs::OsRng + .try_next_u64() + .map_err(|error| anyhow::anyhow!("failed to generate publisher ID: {error}"))?, + ); let discovery = Some(drt.discovery()); let runtime_handle = drt.runtime().secondary(); let subject = scope.subject(&topic); @@ -915,6 +922,12 @@ mod tests { use super::*; use crate::config::environment_names::zmq_broker as broker_env; + #[test] + fn publisher_ids_fit_kubernetes_discovery_integer_range() { + assert_eq!(discovery_safe_publisher_id(42), 42); + assert!(i64::try_from(discovery_safe_publisher_id(u64::MAX)).is_ok()); + } + #[test] fn direct_zmq_topology_selection_is_narrow() { let lookup = |url: Option<&str>, enabled: Option<&str>| {