diff --git a/rust/otap-dataflow/benchmarks/benches/exporter/main.rs b/rust/otap-dataflow/benchmarks/benches/exporter/main.rs index 24d16cbbe5..becadc63b7 100644 --- a/rust/otap-dataflow/benchmarks/benches/exporter/main.rs +++ b/rust/otap-dataflow/benchmarks/benches/exporter/main.rs @@ -60,6 +60,7 @@ use tonic::{Request, Response, Status}; use otap_df_config::node::NodeUserConfig; use otap_df_engine::context::ControllerContext; use otap_df_engine::control::{Controllable, NodeControlMsg, pipeline_ctrl_msg_channel}; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_otap::otap_exporter::OTAP_EXPORTER_URN; use otap_df_otap::otlp_grpc::OTLPData; use otap_df_otap::perf_exporter::exporter::OTAP_PERF_EXPORTER_URN; @@ -454,7 +455,7 @@ fn bench_exporter(c: &mut Criterion) { let local = LocalSet::new(); let _run_exporter_handle = local.spawn_local(async move { exporter - .start(node_req_tx, metrics_reporter) + .start(node_req_tx, metrics_reporter, ExtensionRegistry::empty()) .await .expect("Exporter event loop failed") }); @@ -520,7 +521,7 @@ fn bench_exporter(c: &mut Criterion) { let local = LocalSet::new(); let _run_exporter_handle = local.spawn_local(async move { exporter - .start(node_req_tx, metrics_reporter) + .start(node_req_tx, metrics_reporter, ExtensionRegistry::empty()) .await .expect("Exporter event loop failed") }); @@ -591,7 +592,7 @@ fn bench_exporter(c: &mut Criterion) { let local = LocalSet::new(); let _run_exporter_handle = local.spawn_local(async move { exporter - .start(node_req_tx, metrics_reporter) + .start(node_req_tx, metrics_reporter, ExtensionRegistry::empty()) .await .expect("Exporter event loop failed") }); diff --git a/rust/otap-dataflow/crates/config/src/node.rs b/rust/otap-dataflow/crates/config/src/node.rs index 3e96a83d80..d6bc3dbf5a 100644 --- a/rust/otap-dataflow/crates/config/src/node.rs +++ b/rust/otap-dataflow/crates/config/src/node.rs @@ -94,6 +94,8 @@ pub enum NodeKind { // Connector, /// A merged chain of consecutive processors (experimental). ProcessorChain, + /// A non-pipeline extension (e.g., auth provider, health check). + Extension, } impl From for Cow<'static, str> { @@ -103,6 +105,7 @@ impl From for Cow<'static, str> { NodeKind::Processor => "processor".into(), NodeKind::Exporter => "exporter".into(), NodeKind::ProcessorChain => "processor_chain".into(), + NodeKind::Extension => "extension".into(), } } } diff --git a/rust/otap-dataflow/crates/config/src/node_urn.rs b/rust/otap-dataflow/crates/config/src/node_urn.rs index ed39f870aa..aac39739f1 100644 --- a/rust/otap-dataflow/crates/config/src/node_urn.rs +++ b/rust/otap-dataflow/crates/config/src/node_urn.rs @@ -208,6 +208,7 @@ const fn kind_suffix(expected_kind: NodeKind) -> &'static str { NodeKind::Receiver => "receiver", NodeKind::Processor | NodeKind::ProcessorChain => "processor", NodeKind::Exporter => "exporter", + NodeKind::Extension => "extension", } } @@ -228,9 +229,12 @@ fn parse_kind(raw: &str, kind: &str) -> Result { "receiver" => Ok(NodeKind::Receiver), "processor" => Ok(NodeKind::Processor), "exporter" => Ok(NodeKind::Exporter), + "extension" => Ok(NodeKind::Extension), _ => Err(invalid_plugin_urn( raw, - format!("expected kind `receiver`, `processor`, or `exporter`, found `{kind}`"), + format!( + "expected kind `receiver`, `processor`, `exporter`, or `extension`, found `{kind}`" + ), )), } } diff --git a/rust/otap-dataflow/crates/config/src/pipeline.rs b/rust/otap-dataflow/crates/config/src/pipeline.rs index c6efedd3a9..9a6142eed3 100644 --- a/rust/otap-dataflow/crates/config/src/pipeline.rs +++ b/rust/otap-dataflow/crates/config/src/pipeline.rs @@ -526,6 +526,9 @@ impl PipelineConfig { !has_incoming || !has_outgoing } NodeKind::Exporter => !has_incoming, + // Extensions are standalone services; they never participate + // in the data-flow graph and must not be pruned. + NodeKind::Extension => false, }; if should_remove { diff --git a/rust/otap-dataflow/crates/contrib-nodes/Cargo.toml b/rust/otap-dataflow/crates/contrib-nodes/Cargo.toml index 5feadea6fe..e6d9dd3fda 100644 --- a/rust/otap-dataflow/crates/contrib-nodes/Cargo.toml +++ b/rust/otap-dataflow/crates/contrib-nodes/Cargo.toml @@ -84,6 +84,20 @@ recordset-kql-processor = [ ] resource-validator-processor = [] +contrib-extensions = [ + "bearer-auth-extension", + "azure-identity-auth-extension", +] +bearer-auth-extension = [ + "dep:http", +] +azure-identity-auth-extension = [ + "dep:azure_core", + "dep:azure_identity", + "dep:http", + "dep:rand", +] + [dev-dependencies] otap-df-engine = { path = "../engine", features = ["test-utils"] } otap-df-otap = { path = "../otap", features = ["test-utils"] } diff --git a/rust/otap-dataflow/crates/contrib-nodes/src/exporters/azure_monitor_exporter/exporter.rs b/rust/otap-dataflow/crates/contrib-nodes/src/exporters/azure_monitor_exporter/exporter.rs index 2d627fc814..2fe77b0c50 100644 --- a/rust/otap-dataflow/crates/contrib-nodes/src/exporters/azure_monitor_exporter/exporter.rs +++ b/rust/otap-dataflow/crates/contrib-nodes/src/exporters/azure_monitor_exporter/exporter.rs @@ -9,6 +9,7 @@ use otap_df_engine::ConsumerEffectHandlerExtension; use otap_df_engine::context::PipelineContext; use otap_df_engine::control::{AckMsg, NackMsg, NodeControlMsg}; use otap_df_engine::error::Error as EngineError; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::exporter::{EffectHandler, Exporter}; use otap_df_engine::message::{Message, MessageChannel}; use otap_df_engine::terminal_state::TerminalState; @@ -462,6 +463,7 @@ impl Exporter for AzureMonitorExporter { mut self: Box, mut msg_chan: MessageChannel, effect_handler: EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { effect_handler .info(&format!( diff --git a/rust/otap-dataflow/crates/contrib-nodes/src/exporters/geneva_exporter/mod.rs b/rust/otap-dataflow/crates/contrib-nodes/src/exporters/geneva_exporter/mod.rs index 87340ad8c7..44221678bd 100644 --- a/rust/otap-dataflow/crates/contrib-nodes/src/exporters/geneva_exporter/mod.rs +++ b/rust/otap-dataflow/crates/contrib-nodes/src/exporters/geneva_exporter/mod.rs @@ -36,6 +36,7 @@ use otap_df_engine::control::NodeControlMsg; use otap_df_engine::control::{AckMsg, NackMsg}; use otap_df_engine::error::Error; use otap_df_engine::exporter::ExporterWrapper; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::exporter::{EffectHandler, Exporter}; use otap_df_engine::message::{Message, MessageChannel}; use otap_df_engine::node::NodeId; @@ -501,6 +502,7 @@ impl Exporter for GenevaExporter { mut self: Box, mut msg_chan: MessageChannel, effect_handler: EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { otel_info!( "geneva_exporter.start", diff --git a/rust/otap-dataflow/crates/contrib-nodes/src/extensions/azure_identity_auth_extension/config.rs b/rust/otap-dataflow/crates/contrib-nodes/src/extensions/azure_identity_auth_extension/config.rs new file mode 100644 index 0000000000..17d5a46c56 --- /dev/null +++ b/rust/otap-dataflow/crates/contrib-nodes/src/extensions/azure_identity_auth_extension/config.rs @@ -0,0 +1,143 @@ +// Copyright The OpenTelemetry Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Configuration types for the Azure Identity Auth Extension. + +use serde::Deserialize; + +/// Authentication method for Azure. +#[derive(Debug, Deserialize, Clone, PartialEq, Default)] +#[serde(rename_all = "lowercase")] +pub enum AuthMethod { + /// Use Managed Identity (system or user-assigned with client_id). + #[serde(alias = "msi", alias = "managed_identity")] + #[default] + ManagedIdentity, + + /// Use developer tools (Azure CLI, Azure Developer CLI). + #[serde(alias = "dev", alias = "developer", alias = "cli")] + Development, +} + +impl std::fmt::Display for AuthMethod { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + AuthMethod::ManagedIdentity => write!(f, "managed_identity"), + AuthMethod::Development => write!(f, "development"), + } + } +} + +/// Configuration for the Azure Identity Auth Extension. +#[derive(Debug, Deserialize, Clone)] +#[serde(deny_unknown_fields)] +pub struct Config { + /// Authentication method to use. + #[serde(default)] + pub method: AuthMethod, + + /// Client ID for user-assigned managed identity (optional). + /// Only used when method is ManagedIdentity. + /// If not provided with ManagedIdentity, system-assigned identity will be used. + pub client_id: Option, + + /// OAuth scope for token acquisition. + /// Defaults to "https://management.azure.com/.default" for general Azure management. + #[serde(default = "default_scope")] + pub scope: String, +} + +impl Default for Config { + fn default() -> Self { + Self { + method: AuthMethod::default(), + client_id: None, + scope: default_scope(), + } + } +} + +impl Config { + /// Validate the configuration. + pub fn validate(&self) -> Result<(), super::error::Error> { + if self.scope.is_empty() { + return Err(super::error::Error::Config( + "OAuth scope cannot be empty".to_string(), + )); + } + Ok(()) + } +} + +fn default_scope() -> String { + "https://management.azure.com/.default".to_string() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn default_config() { + let config = Config::default(); + assert_eq!(config.method, AuthMethod::ManagedIdentity); + assert!(config.client_id.is_none()); + assert_eq!(config.scope, "https://management.azure.com/.default"); + } + + #[test] + fn auth_method_display() { + assert_eq!( + format!("{}", AuthMethod::ManagedIdentity), + "managed_identity" + ); + assert_eq!(format!("{}", AuthMethod::Development), "development"); + } + + #[test] + fn config_validation_empty_scope() { + let config = Config { + method: AuthMethod::ManagedIdentity, + client_id: None, + scope: String::new(), + }; + assert!(config.validate().is_err()); + } + + #[test] + fn config_validation_valid() { + let config = Config::default(); + assert!(config.validate().is_ok()); + } + + #[test] + fn config_deserialize_managed_identity() { + let json = serde_json::json!({ + "method": "managed_identity", + "scope": "https://monitor.azure.com/.default" + }); + let cfg: Config = serde_json::from_value(json).unwrap(); + assert_eq!(cfg.method, AuthMethod::ManagedIdentity); + assert_eq!(cfg.scope, "https://monitor.azure.com/.default"); + } + + #[test] + fn config_deserialize_development() { + let json = serde_json::json!({ + "method": "development", + "scope": "https://monitor.azure.com/.default" + }); + let cfg: Config = serde_json::from_value(json).unwrap(); + assert_eq!(cfg.method, AuthMethod::Development); + } + + #[test] + fn config_rejects_unknown_fields() { + let json = serde_json::json!({ + "method": "managed_identity", + "scope": "https://test.scope", + "unknown_field": true + }); + assert!(serde_json::from_value::(json).is_err()); + } +} diff --git a/rust/otap-dataflow/crates/contrib-nodes/src/extensions/azure_identity_auth_extension/error.rs b/rust/otap-dataflow/crates/contrib-nodes/src/extensions/azure_identity_auth_extension/error.rs new file mode 100644 index 0000000000..861c1bfd70 --- /dev/null +++ b/rust/otap-dataflow/crates/contrib-nodes/src/extensions/azure_identity_auth_extension/error.rs @@ -0,0 +1,68 @@ +// Copyright The OpenTelemetry Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Error types for the Azure Identity Auth Extension. + +use super::config::AuthMethod; + +/// Error definitions for Azure Identity Auth Extension. +#[derive(thiserror::Error, Debug)] +pub enum Error { + /// Error during configuration of a component. + #[error("Configuration error: {0}")] + Config(String), + + /// Authentication/authorization error. + #[error("Auth error ({kind})")] + Auth { + /// The kind of authentication error. + kind: AuthErrorKind, + /// The underlying Azure error, if any. + #[source] + source: Option, + }, +} + +/// Specific authentication error variants. +#[derive(Debug, Clone, PartialEq)] +pub enum AuthErrorKind { + /// Failed to create the credential provider. + CreateCredential { + /// The authentication method that failed. + method: AuthMethod, + }, + + /// Failed to acquire a token. + TokenAcquisition, +} + +impl std::fmt::Display for AuthErrorKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + AuthErrorKind::CreateCredential { method } => { + write!(f, "failed to create credential for method: {}", method) + } + AuthErrorKind::TokenAcquisition => write!(f, "failed to acquire token"), + } + } +} + +impl Error { + /// Creates a new credential creation error. + #[must_use] + pub fn create_credential(method: AuthMethod, source: azure_core::error::Error) -> Self { + Error::Auth { + kind: AuthErrorKind::CreateCredential { method }, + source: Some(source), + } + } + + /// Creates a new token acquisition error. + #[must_use] + pub fn token_acquisition(source: azure_core::error::Error) -> Self { + Error::Auth { + kind: AuthErrorKind::TokenAcquisition, + source: Some(source), + } + } +} diff --git a/rust/otap-dataflow/crates/contrib-nodes/src/extensions/azure_identity_auth_extension/extension.rs b/rust/otap-dataflow/crates/contrib-nodes/src/extensions/azure_identity_auth_extension/extension.rs new file mode 100644 index 0000000000..976a101eaf --- /dev/null +++ b/rust/otap-dataflow/crates/contrib-nodes/src/extensions/azure_identity_auth_extension/extension.rs @@ -0,0 +1,329 @@ +// Copyright The OpenTelemetry Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Azure Identity Auth Extension implementation. +//! +//! This extension provides Azure authentication services to the pipeline. +//! It manages Azure credentials and provides token acquisition capabilities +//! to consumers (e.g., exporters) via the [`BearerTokenProviderHandle`]. +//! +//! # Architecture +//! +//! The extension creates a [`BearerTokenProviderHandle`] at construction time, +//! which wraps an `AzureTokenProvider` implementing the [`BearerTokenProvider`] +//! trait. The extension background task refreshes tokens periodically and +//! broadcasts updates through a `watch` channel that subscribers receive via +//! the handle. +//! +//! Consumers retrieve the handle from the extension registry: +//! +//! ```rust,ignore +//! let handle = extension_registry +//! .get::("azure_auth")?; +//! let token = handle.get_token().await?; +//! ``` + +use async_trait::async_trait; +use azure_core::credentials::{AccessToken, TokenCredential}; +use azure_identity::{ + DeveloperToolsCredential, DeveloperToolsCredentialOptions, ManagedIdentityCredential, + ManagedIdentityCredentialOptions, UserAssignedId, +}; +use otap_df_engine::extensions::{ + BearerToken, BearerTokenError, BearerTokenProvider, +}; +use std::sync::{Arc, Mutex}; +use tokio::sync::watch; + +use otap_df_engine::error::Error as EngineError; +use otap_df_engine::local::extension as local; + +use super::config::{AuthMethod, Config}; +use super::error::Error; + +/// Buffer time before token expiry to trigger refresh (in seconds). +const TOKEN_EXPIRY_BUFFER_SECS: u64 = 299; +/// Minimum interval between token refresh attempts (in seconds). +const MIN_TOKEN_REFRESH_INTERVAL_SECS: u64 = 10; +/// Minimum delay between token refresh retry attempts in seconds. +const MIN_RETRY_DELAY_SECS: f64 = 5.0; +/// Maximum delay between token refresh retry attempts in seconds. +const MAX_RETRY_DELAY_SECS: f64 = 30.0; +/// Maximum jitter percentage (±10%) to add to retry delays. +const MAX_RETRY_JITTER_RATIO: f64 = 0.10; +/// Retry interval when token refresh fails (in seconds). +const TOKEN_REFRESH_RETRY_SECS: u64 = 10; + +/// The token provider implementation that backs the [`BearerTokenProviderHandle`]. +/// +/// Reads the latest token from a shared cache (updated by the extension's +/// background refresh loop) and exposes a watch channel for subscribers. +/// If the cache is empty or expired, falls back to fetching from Azure directly. +pub(crate) struct AzureTokenProvider { + /// The Azure credential provider. + credential: Arc, + /// The OAuth scope for token acquisition. + scope: String, + /// Sender for broadcasting token refresh events (used by `subscribe_token_refresh`). + token_sender: Arc>>, + /// Cached token updated by the extension's refresh loop (used by `get_token`). + token_cache_for_demonstration: Arc>>, +} + +#[async_trait] +impl BearerTokenProvider for AzureTokenProvider { + async fn get_token(&self) -> Result { + // Fast path: return cached token if it hasn't expired. + { + let cache = self.token_cache_for_demonstration.lock().expect("token_cache lock poisoned"); + if let Some(cached) = cache.as_ref() { + let now_secs = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs() as i64) + .unwrap_or(0); + if cached.expires_on > now_secs { + return Ok(cached.clone()); + } + } + } + + // Slow path: cache is empty or expired — fetch from Azure. + let access_token = self + .credential + .get_token( + &[&self.scope], + Some(azure_core::credentials::TokenRequestOptions::default()), + ) + .await + .map_err(|e| BearerTokenError { + message: e.to_string(), + })?; + + let token = BearerToken::new( + access_token.token.secret().to_string(), + access_token.expires_on.unix_timestamp(), + ); + + // Update cache so next call hits the fast path. + *self.token_cache_for_demonstration.lock().expect("token_cache lock poisoned") = Some(token.clone()); + + Ok(token) + } + + fn subscribe_token_refresh(&self) -> watch::Receiver> { + self.token_sender.subscribe() + } +} + +/// Azure Identity Auth Extension. +/// +/// Runs as a background task that periodically refreshes the Azure bearer +/// token and broadcasts updates through the token provider handle. +pub struct AzureIdentityAuthExtension { + /// The Azure credential provider (shared with the token provider). + credential: Arc, + /// Human-readable description of the credential type. + credential_type: &'static str, + /// The OAuth scope for token acquisition. + scope: String, + /// Sender for broadcasting token refresh events. + token_sender: Arc>>, + /// Shared cache updated by the refresh loop, read by the provider handle. + token_cache_for_demonstration: Arc>>, +} + +impl AzureIdentityAuthExtension { + /// Creates a new Azure Identity Auth Extension and its associated token provider. + /// + /// Returns both the extension (for the background task) and the provider + /// (to be wrapped in a `BearerTokenProviderHandle` and registered). + pub(crate) fn new(config: &Config) -> Result<(Self, AzureTokenProvider), Error> { + let (credential, credential_type) = Self::create_credential(config)?; + let (token_sender, _) = watch::channel(None); + let token_sender = Arc::new(token_sender); + let token_cache_for_demonstration = Arc::new(Mutex::new(None)); + + let provider = AzureTokenProvider { + credential: Arc::clone(&credential), + scope: config.scope.clone(), + token_sender: Arc::clone(&token_sender), + token_cache_for_demonstration: Arc::clone(&token_cache_for_demonstration), + }; + + let extension = Self { + credential, + credential_type, + scope: config.scope.clone(), + token_sender, + token_cache_for_demonstration, + }; + + Ok((extension, provider)) + } + + /// Creates a credential provider based on the configuration. + fn create_credential( + config: &Config, + ) -> Result<(Arc, &'static str), Error> { + match config.method { + AuthMethod::ManagedIdentity => { + let mut options = ManagedIdentityCredentialOptions::default(); + + let credential_type = if let Some(client_id) = &config.client_id { + options.user_assigned_id = Some(UserAssignedId::ClientId(client_id.clone())); + "user_assigned_managed_identity" + } else { + "system_assigned_managed_identity" + }; + + Ok(( + ManagedIdentityCredential::new(Some(options)) + .map_err(|e| Error::create_credential(AuthMethod::ManagedIdentity, e))?, + credential_type, + )) + } + AuthMethod::Development => Ok(( + DeveloperToolsCredential::new(Some(DeveloperToolsCredentialOptions::default())) + .map_err(|e| Error::create_credential(AuthMethod::Development, e))?, + "developer_tools", + )), + } + } + + /// Gets a token directly from the credential provider. + async fn get_token_internal(&self) -> Result { + self.credential + .get_token( + &[&self.scope], + Some(azure_core::credentials::TokenRequestOptions::default()), + ) + .await + .map_err(Error::token_acquisition) + } + + /// Gets a token with retry logic and exponential backoff. + async fn get_token_with_retry(&self) -> Result { + let mut attempt = 0_i32; + loop { + attempt += 1; + + match self.get_token_internal().await { + Ok(token) => return Ok(token), + Err(_e) if attempt < 10 => { + let base_delay_secs = MIN_RETRY_DELAY_SECS * 2.0_f64.powi(attempt - 1); + let capped_delay_secs = base_delay_secs.min(MAX_RETRY_DELAY_SECS); + + let jitter_range = capped_delay_secs * MAX_RETRY_JITTER_RATIO; + let jitter = if jitter_range > 0.0 { + let random_factor = rand::random::() * 2.0 - 1.0; + random_factor * jitter_range + } else { + 0.0 + }; + + let delay_secs = (capped_delay_secs + jitter).max(1.0); + tokio::time::sleep(tokio::time::Duration::from_secs_f64(delay_secs)).await; + } + Err(e) => return Err(e), + } + } + } + + /// Calculates when the next token refresh should occur. + fn get_next_token_refresh(token: &BearerToken) -> tokio::time::Instant { + let now_secs = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs() as i64) + .unwrap_or(0); + + let duration_remaining = if token.expires_on > now_secs { + std::time::Duration::from_secs((token.expires_on - now_secs) as u64) + } else { + std::time::Duration::ZERO + }; + + let token_valid_until = tokio::time::Instant::now() + duration_remaining; + let next_token_refresh = + token_valid_until - tokio::time::Duration::from_secs(TOKEN_EXPIRY_BUFFER_SECS); + std::cmp::max( + next_token_refresh, + tokio::time::Instant::now() + + tokio::time::Duration::from_secs(MIN_TOKEN_REFRESH_INTERVAL_SECS), + ) + } +} + +#[async_trait(?Send)] +impl local::Extension for AzureIdentityAuthExtension { + async fn start( + self: Box, + mut ctrl_chan: local::ControlChannel, + effect_handler: local::EffectHandler, + ) -> Result<(), EngineError> { + effect_handler + .info(&format!( + "[azure-identity-auth] starting (credential_type={}, scope={})", + self.credential_type, self.scope + )) + .await; + + // Fetch initial token immediately. + let mut next_token_refresh = tokio::time::Instant::now(); + + loop { + tokio::select! { + biased; + + // Proactive token refresh. + _ = tokio::time::sleep_until(next_token_refresh) => { + match self.get_token_with_retry().await { + Ok(access_token) => { + let bearer_token = BearerToken::new( + access_token.token.secret().to_string(), + access_token.expires_on.unix_timestamp(), + ); + + // Update shared cache so get_token() can return it directly. + *self.token_cache_for_demonstration.lock().expect("token_cache lock poisoned") = + Some(bearer_token.clone()); + + // Broadcast the new token to all subscribers. + let _ = self.token_sender.send(Some(bearer_token.clone())); + + next_token_refresh = + Self::get_next_token_refresh(&bearer_token); + + effect_handler + .info("[azure-identity-auth] token refreshed") + .await; + } + Err(_e) => { + next_token_refresh = tokio::time::Instant::now() + + tokio::time::Duration::from_secs(TOKEN_REFRESH_RETRY_SECS); + + effect_handler + .info("[azure-identity-auth] token refresh failed, retrying") + .await; + } + } + } + + // Handle control messages. + msg = ctrl_chan.recv() => { + match msg { + Ok(msg) if msg.is_shutdown() => { + effect_handler + .info("[azure-identity-auth] shutting down") + .await; + break; + } + Ok(_) => {} + Err(_) => break, + } + } + } + } + + Ok(()) + } +} diff --git a/rust/otap-dataflow/crates/contrib-nodes/src/extensions/azure_identity_auth_extension/mod.rs b/rust/otap-dataflow/crates/contrib-nodes/src/extensions/azure_identity_auth_extension/mod.rs new file mode 100644 index 0000000000..4954beceb1 --- /dev/null +++ b/rust/otap-dataflow/crates/contrib-nodes/src/extensions/azure_identity_auth_extension/mod.rs @@ -0,0 +1,115 @@ +// Copyright The OpenTelemetry Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Azure Identity Auth Extension for OTAP. +//! +//! Provides Azure authentication services to the pipeline using Azure Identity. +//! This extension manages token acquisition and refresh, making credentials +//! available to other components (e.g., exporters) via the +//! [`BearerTokenProviderHandle`](otap_df_engine::extensions::BearerTokenProviderHandle). +//! +//! # Usage +//! +//! Configure the extension in the pipeline configuration: +//! +//! ```yaml +//! extensions: +//! azure_auth: +//! type: "urn:microsoft:extension:azure_identity_auth" +//! config: +//! method: managed_identity +//! scope: "https://monitor.azure.com/.default" +//! ``` +//! +//! Consumers retrieve the handle from the extension registry: +//! +//! ```ignore +//! let handle = extension_registry +//! .get::("azure_auth")?; +//! let token = handle.get_token().await?; +//! ``` + +use linkme::distributed_slice; +use otap_df_config::node::NodeUserConfig; +use otap_df_engine::ExtensionFactory; +use otap_df_engine::config::ExtensionConfig; +use otap_df_engine::context::PipelineContext; +use otap_df_engine::extension::ExtensionWrapper; +use otap_df_engine::extensions::{BearerTokenProviderHandle, ExtensionHandles}; +use otap_df_engine::node::NodeId; +use otap_df_otap::OTAP_EXTENSION_FACTORIES; +use std::sync::Arc; + +pub mod config; +pub mod error; +mod extension; + +pub use config::{AuthMethod, Config}; +pub use error::Error; +pub use extension::AzureIdentityAuthExtension; + +/// URN identifying the Azure Identity Auth Extension in configuration pipelines. +pub const AZURE_IDENTITY_AUTH_EXTENSION_URN: &str = "urn:microsoft:extension:azure_identity_auth"; + +/// Register Azure Identity Auth Extension with the OTAP extension factory. +/// +/// Uses the `distributed_slice` macro for automatic discovery by the dataflow engine. +#[allow(unsafe_code)] +#[distributed_slice(OTAP_EXTENSION_FACTORIES)] +pub static AZURE_IDENTITY_AUTH_EXTENSION: ExtensionFactory = ExtensionFactory { + name: AZURE_IDENTITY_AUTH_EXTENSION_URN, + create: |_pipeline_ctx: PipelineContext, + node_id: NodeId, + node_config: Arc, + extension_config: &ExtensionConfig| { + let cfg: Config = serde_json::from_value(node_config.config.clone()).map_err(|e| { + otap_df_config::error::Error::InvalidUserConfig { + error: e.to_string(), + } + })?; + + cfg.validate() + .map_err(|e| otap_df_config::error::Error::InvalidUserConfig { + error: e.to_string(), + })?; + + let (extension, provider) = AzureIdentityAuthExtension::new(&cfg).map_err(|e| { + otap_df_config::error::Error::InvalidUserConfig { + error: e.to_string(), + } + })?; + + let mut handles = ExtensionHandles::new(); + handles.register(BearerTokenProviderHandle::new(provider)); + + Ok(ExtensionWrapper::local( + extension, + handles, + node_id, + node_config, + extension_config, + )) + }, + validate_config: otap_df_config::validation::validate_typed_config::, +}; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extension_urn() { + assert_eq!( + AZURE_IDENTITY_AUTH_EXTENSION_URN, + "urn:microsoft:extension:azure_identity_auth" + ); + } + + #[test] + fn factory_name_matches_urn() { + assert_eq!( + AZURE_IDENTITY_AUTH_EXTENSION.name, + AZURE_IDENTITY_AUTH_EXTENSION_URN + ); + } +} diff --git a/rust/otap-dataflow/crates/contrib-nodes/src/extensions/bearer_auth_extension/mod.rs b/rust/otap-dataflow/crates/contrib-nodes/src/extensions/bearer_auth_extension/mod.rs new file mode 100644 index 0000000000..adc0a3ee72 --- /dev/null +++ b/rust/otap-dataflow/crates/contrib-nodes/src/extensions/bearer_auth_extension/mod.rs @@ -0,0 +1,290 @@ +// Copyright The OpenTelemetry Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Bearer token authentication extension with periodic token refresh. +//! +//! This extension provides client-side (exporter) bearer token authentication. +//! The token is refreshed periodically in the extension's background loop, +//! and the `ClientAuthenticatorHandle` automatically sees updated values. +//! +//! # Configuration +//! +//! ```yaml +//! extensions: +//! my_auth: +//! type: "urn:otap:extension:auth/bearer" +//! config: +//! refresh_interval_secs: 300 +//! ``` +//! +//! # Usage +//! +//! Exporters can attach credentials to outgoing requests: +//! +//! ```rust,ignore +//! let auth = extension_registry +//! .get::("my_auth")?; +//! for (key, value) in auth.get_request_metadata()? { +//! request.headers_mut().insert(key, value); +//! } +//! ``` + +use async_trait::async_trait; +use linkme::distributed_slice; +use otap_df_config::node::NodeUserConfig; +use otap_df_engine::ExtensionFactory; +use otap_df_engine::config::ExtensionConfig; +use otap_df_engine::context::PipelineContext; +use otap_df_engine::error::Error; +use otap_df_engine::extension::ExtensionWrapper; +use otap_df_engine::extensions::{ + AuthError, ClientAuthenticator, ClientAuthenticatorHandle, ExtensionHandles, +}; +use otap_df_engine::local::extension as local; +use otap_df_engine::node::NodeId; +use otap_df_otap::OTAP_EXTENSION_FACTORIES; +use serde::Deserialize; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +/// A shared, mutable token that the extension updates and the authenticator reads. +type SharedToken = Arc>; + +/// Default interval between token refreshes. +const TOKEN_REFRESH_INTERVAL: Duration = Duration::from_secs(300); + +/// URN identifying the bearer auth extension in pipeline configuration. +pub const BEARER_AUTH_EXTENSION_URN: &str = "urn:otap:extension:auth/bearer"; + +/// Configuration for the bearer token auth extension. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Config { + /// Token refresh interval in seconds. Defaults to 300 (5 minutes). + #[serde(default = "default_refresh_interval_secs")] + pub refresh_interval_secs: u64, +} + +fn default_refresh_interval_secs() -> u64 { + TOKEN_REFRESH_INTERVAL.as_secs() +} + +/// Register the bearer auth extension with the OTAP extension factory. +/// +/// Uses the `distributed_slice` macro for automatic discovery by the dataflow engine. +#[allow(unsafe_code)] +#[distributed_slice(OTAP_EXTENSION_FACTORIES)] +pub static BEARER_AUTH_EXTENSION: ExtensionFactory = ExtensionFactory { + name: BEARER_AUTH_EXTENSION_URN, + create: |_pipeline_ctx: PipelineContext, + node_id: NodeId, + node_config: Arc, + extension_config: &ExtensionConfig| { + let cfg: Config = serde_json::from_value(node_config.config.clone()).map_err(|e| { + otap_df_config::error::Error::InvalidUserConfig { + error: e.to_string(), + } + })?; + + // Shared token — starts empty, populated by the first refresh in the start loop. + let shared_token: SharedToken = Arc::new(Mutex::new(String::new())); + + let mut handles = ExtensionHandles::new(); + handles.register(ClientAuthenticatorHandle::new(BearerClientAuth { + token: Arc::clone(&shared_token), + })); + + Ok(ExtensionWrapper::local( + BearerAuthExtension { + token: shared_token, + refresh_interval: Duration::from_secs(cfg.refresh_interval_secs), + }, + handles, + node_id, + node_config, + extension_config, + )) + }, + validate_config: otap_df_config::validation::validate_typed_config::, +}; + +/// The extension instance that runs as a background task. +/// +/// Periodically refreshes the shared bearer token. The new token is a +/// mock value generated from the current timestamp — in a real extension +/// this would call an identity provider or secret store. +struct BearerAuthExtension { + token: SharedToken, + refresh_interval: Duration, +} + +impl BearerAuthExtension { + /// Generates a mock token that looks like a GUID. + /// + /// Uses the current system time to produce a deterministic-looking but + /// unique value. A production implementation would fetch a real token + /// from an identity provider. + fn generate_token() -> String { + use std::time::SystemTime; + + let nanos = SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + + // Format as a GUID-like string from the timestamp bits. + format!( + "{:08x}-{:04x}-{:04x}-{:04x}-{:012x}", + (nanos >> 96) as u32, + (nanos >> 80) as u16, + (nanos >> 64) as u16, + (nanos >> 48) as u16, + nanos as u64 & 0xffff_ffff_ffff, + ) + } + + fn refresh_token(&self) { + let new_token = Self::generate_token(); + *self.token.lock().expect("token lock poisoned") = new_token; + } +} + +#[async_trait(?Send)] +impl local::Extension for BearerAuthExtension { + async fn start( + self: Box, + mut ctrl_chan: local::ControlChannel, + effect_handler: local::EffectHandler, + ) -> Result<(), Error> { + // Fetch the first token immediately. + self.refresh_token(); + effect_handler + .info("[bearer-auth] initial token acquired") + .await; + + let mut refresh_timer = tokio::time::interval(self.refresh_interval); + // The first tick completes immediately — skip it since we just refreshed. + let _ = refresh_timer.tick().await; + + loop { + tokio::select! { + msg = ctrl_chan.recv() => { + match msg { + Ok(msg) if msg.is_shutdown() => break, + Ok(_) => {} + Err(_) => break, + } + } + _ = refresh_timer.tick() => { + self.refresh_token(); + effect_handler.info("[bearer-auth] token refreshed").await; + } + } + } + Ok(()) + } +} + +/// Client-side authenticator that reads the current token from the shared +/// `Arc>` and produces an `Authorization: Bearer ` header. +struct BearerClientAuth { + token: SharedToken, +} + +impl ClientAuthenticator for BearerClientAuth { + fn get_request_metadata( + &self, + ) -> Result, AuthError> { + let current_token = self.token.lock().expect("token lock poisoned"); + if current_token.is_empty() { + return Err(AuthError { + message: "token not yet available".into(), + }); + } + Ok(vec![( + http::header::AUTHORIZATION, + http::HeaderValue::from_str(&format!("Bearer {}", *current_token)).map_err(|e| { + AuthError { + message: e.to_string(), + } + })?, + )]) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn shared_token(s: &str) -> SharedToken { + Arc::new(Mutex::new(s.to_owned())) + } + + #[test] + fn config_deserializes_defaults() { + let json = serde_json::json!({}); + let cfg: Config = serde_json::from_value(json).unwrap(); + assert_eq!(cfg.refresh_interval_secs, 300); + } + + #[test] + fn config_with_custom_interval() { + let json = serde_json::json!({ "refresh_interval_secs": 60 }); + let cfg: Config = serde_json::from_value(json).unwrap(); + assert_eq!(cfg.refresh_interval_secs, 60); + } + + #[test] + fn config_rejects_unknown_fields() { + let json = serde_json::json!({ "unknown": true }); + assert!(serde_json::from_value::(json).is_err()); + } + + #[test] + fn client_auth_produces_bearer_header() { + let auth = ClientAuthenticatorHandle::new(BearerClientAuth { + token: shared_token("outgoing-token"), + }); + let metadata = auth.get_request_metadata().unwrap(); + assert_eq!(metadata.len(), 1); + assert_eq!(metadata[0].0, http::header::AUTHORIZATION); + assert_eq!(metadata[0].1, "Bearer outgoing-token"); + } + + #[test] + fn client_auth_errors_when_token_empty() { + let auth = ClientAuthenticatorHandle::new(BearerClientAuth { + token: shared_token(""), + }); + let err = auth.get_request_metadata().unwrap_err(); + assert!(err.message.contains("not yet available")); + } + + #[test] + fn token_refresh_updates_client_handle() { + let token = shared_token("initial"); + let client = ClientAuthenticatorHandle::new(BearerClientAuth { + token: Arc::clone(&token), + }); + + // Simulate a token refresh. + *token.lock().unwrap() = "refreshed".to_owned(); + + let metadata = client.get_request_metadata().unwrap(); + assert_eq!(metadata[0].1, "Bearer refreshed"); + } + + #[test] + fn generate_token_produces_guid_format() { + let token = BearerAuthExtension::generate_token(); + let parts: Vec<&str> = token.split('-').collect(); + assert_eq!(parts.len(), 5, "token should have 5 dash-separated parts"); + } + + #[test] + fn handle_registers_in_extension_handles() { + let token = shared_token("tok"); + let mut handles = ExtensionHandles::new(); + handles.register(ClientAuthenticatorHandle::new(BearerClientAuth { token })); + } +} diff --git a/rust/otap-dataflow/crates/contrib-nodes/src/extensions/mod.rs b/rust/otap-dataflow/crates/contrib-nodes/src/extensions/mod.rs new file mode 100644 index 0000000000..49e780f281 --- /dev/null +++ b/rust/otap-dataflow/crates/contrib-nodes/src/extensions/mod.rs @@ -0,0 +1,10 @@ +// Copyright The OpenTelemetry Authors +// SPDX-License-Identifier: Apache-2.0 + +/// Azure Identity authentication extension. +#[cfg(feature = "azure-identity-auth-extension")] +pub mod azure_identity_auth_extension; + +/// Static bearer token authentication extension. +#[cfg(feature = "bearer-auth-extension")] +pub mod bearer_auth_extension; diff --git a/rust/otap-dataflow/crates/contrib-nodes/src/lib.rs b/rust/otap-dataflow/crates/contrib-nodes/src/lib.rs index 2b5cb59c14..3ace1a6c19 100644 --- a/rust/otap-dataflow/crates/contrib-nodes/src/lib.rs +++ b/rust/otap-dataflow/crates/contrib-nodes/src/lib.rs @@ -1,10 +1,13 @@ // Copyright The OpenTelemetry Authors // SPDX-License-Identifier: Apache-2.0 -//! Implementation of the Contrib nodes (receiver, exporter, processor). +//! Implementation of the Contrib nodes (receiver, exporter, processor, extension). /// Exporter implementations for contrib nodes. pub mod exporters; +/// Extension implementations for contrib nodes. +pub mod extensions; + /// Processor implementations for contrib nodes. pub mod processors; diff --git a/rust/otap-dataflow/crates/engine-macros/src/lib.rs b/rust/otap-dataflow/crates/engine-macros/src/lib.rs index dcad1a7ef3..ae87f18661 100644 --- a/rust/otap-dataflow/crates/engine-macros/src/lib.rs +++ b/rust/otap-dataflow/crates/engine-macros/src/lib.rs @@ -75,6 +75,7 @@ pub fn pipeline_factory(args: TokenStream, input: TokenStream) -> TokenStream { let receiver_factories_name = quote::format_ident!("{}_RECEIVER_FACTORIES", prefix); let processor_factories_name = quote::format_ident!("{}_PROCESSOR_FACTORIES", prefix); let exporter_factories_name = quote::format_ident!("{}_EXPORTER_FACTORIES", prefix); + let extension_factories_name = quote::format_ident!("{}_EXTENSION_FACTORIES", prefix); let get_receiver_factory_map_name = quote::format_ident!( "get_{}_receiver_factory_map", prefix.to_string().to_lowercase() @@ -87,6 +88,10 @@ pub fn pipeline_factory(args: TokenStream, input: TokenStream) -> TokenStream { "get_{}_exporter_factory_map", prefix.to_string().to_lowercase() ); + let get_extension_factory_map_name = quote::format_ident!( + "get_{}_extension_factory_map", + prefix.to_string().to_lowercase() + ); let output = quote! { /// A slice of receiver factories. @@ -101,6 +106,10 @@ pub fn pipeline_factory(args: TokenStream, input: TokenStream) -> TokenStream { #[::otap_df_engine::distributed_slice] pub static #exporter_factories_name: [::otap_df_engine::ExporterFactory<#pdata_type>] = [..]; + /// A slice of extension factories. + #[::otap_df_engine::distributed_slice] + pub static #extension_factories_name: [::otap_df_engine::ExtensionFactory] = [..]; + /// The factory registry instance. #registry_vis static #registry_name: std::sync::LazyLock> = std::sync::LazyLock::new(|| { // Reference build_registry to avoid unused import warning, even though we don't call it @@ -109,6 +118,7 @@ pub fn pipeline_factory(args: TokenStream, input: TokenStream) -> TokenStream { &#receiver_factories_name, &#processor_factories_name, &#exporter_factories_name, + &#extension_factories_name, ) }); @@ -126,6 +136,11 @@ pub fn pipeline_factory(args: TokenStream, input: TokenStream) -> TokenStream { pub fn #get_exporter_factory_map_name() -> &'static std::collections::HashMap<&'static str, ::otap_df_engine::ExporterFactory<#pdata_type>> { #registry_name.get_exporter_factory_map() } + + /// Gets the extension factory map, initializing it if necessary. + pub fn #get_extension_factory_map_name() -> &'static std::collections::HashMap<&'static str, ::otap_df_engine::ExtensionFactory> { + #registry_name.get_extension_factory_map() + } }; output.into() diff --git a/rust/otap-dataflow/crates/engine/Cargo.toml b/rust/otap-dataflow/crates/engine/Cargo.toml index 1aad1a0a77..3bb46f0fa5 100644 --- a/rust/otap-dataflow/crates/engine/Cargo.toml +++ b/rust/otap-dataflow/crates/engine/Cargo.toml @@ -45,6 +45,7 @@ data-encoding = { workspace = true } prost = { workspace = true } byte-unit = { workspace = true } cpu-time = { workspace = true } +http = { workspace = true } nix = { workspace = true, features = ["resource"] } [target.'cfg(not(windows))'.dependencies] diff --git a/rust/otap-dataflow/crates/engine/src/channel_mode.rs b/rust/otap-dataflow/crates/engine/src/channel_mode.rs index dc0d40aecd..70ea6143ce 100644 --- a/rust/otap-dataflow/crates/engine/src/channel_mode.rs +++ b/rust/otap-dataflow/crates/engine/src/channel_mode.rs @@ -20,7 +20,6 @@ use crate::channel_metrics::{ ChannelSenderMetrics, control_channel_id, }; use crate::context::PipelineContext; -use crate::control::NodeControlMsg; use crate::entity_context::current_node_telemetry_handle; use crate::local::message::{LocalReceiver, LocalSender}; use crate::shared::message::{SharedReceiver, SharedSender}; @@ -171,24 +170,21 @@ impl ChannelMode for SharedMode { } } -/// Generic helper used by receiver, processor, and exporter wrappers. +/// Generic helper used by receiver, processor, exporter, and extension wrappers. /// It keeps local and shared wiring identical while still emitting mode-specific code. /// /// The logic first attempts to unwrap the inner MPSC channel so metrics can be attached. /// If the channel is already wrapped, it preserves the existing wrapper to avoid double /// instrumentation. -pub(crate) fn wrap_control_channel_metrics( +pub(crate) fn wrap_control_channel_metrics( node_id: &crate::node::NodeId, pipeline_ctx: &PipelineContext, channel_metrics: &mut ChannelMetricsRegistry, channel_metrics_enabled: bool, capacity: u64, - control_sender: M::ControlSender>, - control_receiver: M::ControlReceiver>, -) -> ( - M::ControlSender>, - M::ControlReceiver>, -) + control_sender: M::ControlSender, + control_receiver: M::ControlReceiver, +) -> (M::ControlSender, M::ControlReceiver) where M: ChannelMode, { diff --git a/rust/otap-dataflow/crates/engine/src/config.rs b/rust/otap-dataflow/crates/engine/src/config.rs index 41097d9e58..e6e5445d6f 100644 --- a/rust/otap-dataflow/crates/engine/src/config.rs +++ b/rust/otap-dataflow/crates/engine/src/config.rs @@ -64,6 +64,18 @@ pub struct ExporterConfig { pub input_pdata_channel: PdataChannelConfig, } +/// Generic configuration for an extension. +/// +/// Extensions are non-pipeline components (e.g., auth providers, health checks) +/// that only receive control messages. They do not participate in pdata flow. +#[derive(Clone, Debug)] +pub struct ExtensionConfig { + /// Name of the extension. + pub name: NodeId, + /// Configuration for control channel. + pub control_channel: ControlChannelConfig, +} + impl ReceiverConfig { /// Creates a new receiver configuration with default channel capacities. pub fn new(name: T) -> Self @@ -172,3 +184,27 @@ impl ExporterConfig { } } } + +impl ExtensionConfig { + /// Creates a new extension configuration with default channel capacities. + pub fn new(name: T) -> Self + where + T: Into, + { + Self::with_channel_capacity(name, DEFAULT_CONTROL_CHANNEL_CAPACITY) + } + + /// Creates a new extension configuration with an explicit control channel capacity. + #[must_use] + pub fn with_channel_capacity(name: T, control_channel_capacity: usize) -> Self + where + T: Into, + { + ExtensionConfig { + name: name.into(), + control_channel: ControlChannelConfig { + capacity: control_channel_capacity, + }, + } + } +} diff --git a/rust/otap-dataflow/crates/engine/src/control.rs b/rust/otap-dataflow/crates/engine/src/control.rs index a943553091..4c518a1a61 100644 --- a/rust/otap-dataflow/crates/engine/src/control.rs +++ b/rust/otap-dataflow/crates/engine/src/control.rs @@ -183,6 +183,44 @@ pub enum NodeControlMsg { }, } +/// Control messages sent by the pipeline engine to **extensions**. +/// +/// This is a PData-free subset of [`NodeControlMsg`] — extensions never process +/// Ack/Nack/DelayedData, so they receive a simpler, non-generic enum. +#[derive(Debug, Clone)] +pub enum ExtensionControlMsg { + /// Notifies the extension of a configuration change. + Config { + /// The new configuration as a JSON value. + config: serde_json::Value, + }, + + /// Emitted when a scheduled timer expires. + TimerTick {}, + + /// Signal to collect/flush local telemetry metrics. + CollectTelemetry { + /// Metrics reporter used to collect telemetry metrics. + metrics_reporter: MetricsReporter, + }, + + /// Requests a graceful shutdown. + Shutdown { + /// Deadline for shutdown. + deadline: Instant, + /// Human-readable reason for the shutdown. + reason: String, + }, +} + +impl ExtensionControlMsg { + /// Returns `true` if this control message is a shutdown request. + #[must_use] + pub const fn is_shutdown(&self) -> bool { + matches!(self, ExtensionControlMsg::Shutdown { .. }) + } +} + /// Control messages sent by nodes to the pipeline engine to manage node-specific operations /// and control pipeline behavior. #[derive(Debug, Clone)] diff --git a/rust/otap-dataflow/crates/engine/src/error.rs b/rust/otap-dataflow/crates/engine/src/error.rs index 3bf6768241..752215d3f0 100644 --- a/rust/otap-dataflow/crates/engine/src/error.rs +++ b/rust/otap-dataflow/crates/engine/src/error.rs @@ -95,6 +95,31 @@ impl fmt::Display for ProcessorErrorKind { } } +/// High-level classification for extension failures to aid troubleshooting. +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +pub enum ExtensionErrorKind { + /// Errors caused by invalid or missing configuration. + Configuration, + /// Errors encountered during extension startup. + Startup, + /// Errors raised while shutting down an extension. + Shutdown, + /// Catch-all for extension failures that do not fit other categories. + Other, +} + +impl fmt::Display for ExtensionErrorKind { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let label = match self { + ExtensionErrorKind::Configuration => "configuration", + ExtensionErrorKind::Startup => "startup", + ExtensionErrorKind::Shutdown => "shutdown", + ExtensionErrorKind::Other => "other", + }; + write!(f, "{label}") + } +} + /// Formats the source chain of an error into a single display string. #[must_use] pub fn format_error_sources(error: &(dyn std::error::Error + 'static)) -> String { @@ -433,6 +458,54 @@ pub enum Error { /// The name of the unknown port. port: PortName, }, + + /// The specified extension already exists in the pipeline. + #[error("The extension `{extension}` already exists")] + ExtensionAlreadyExists { + /// The name of the extension that already exists. + extension: NodeId, + }, + + /// A wrapper for the extension errors. + #[error("An extension error occurred in node {extension} ({kind}): {error}{source_detail}")] + ExtensionError { + /// The name of the extension that encountered the error. + extension: NodeId, + + /// High-level classification for the extension failure. + kind: ExtensionErrorKind, + + /// The error that occurred. + error: String, + + /// Pre-formatted representation of the source chain used when rendering the error. + source_detail: String, + }, + + /// Unknown extension plugin. + #[error("Unknown extension plugin `{plugin_urn}`")] + UnknownExtension { + /// The name of the unknown extension plugin. + plugin_urn: NodeUrn, + }, + + /// An extension handle of the given type was already registered for the named extension. + #[error("Extension handle already registered for `{extension}` (type: {type_name})")] + ExtensionHandleAlreadyRegistered { + /// The name of the extension. + extension: String, + /// The type name of the handle. + type_name: String, + }, + + /// No extension handle of the requested type was found. + #[error("Extension handle not found for `{extension}` (type: {type_name})")] + ExtensionHandleNotFound { + /// The name of the extension. + extension: String, + /// The type name of the handle. + type_name: String, + }, } impl Error { @@ -463,6 +536,11 @@ impl Error { Error::ReceiverAlreadyExists { .. } => "ReceiverAlreadyExists", Error::ReceiverError { .. } => "ReceiverError", Error::SpmcSharedNotSupported { .. } => "SpmcSharedNotSupported", + Error::ExtensionAlreadyExists { .. } => "ExtensionAlreadyExists", + Error::ExtensionError { .. } => "ExtensionError", + Error::ExtensionHandleAlreadyRegistered { .. } => "ExtensionHandleAlreadyRegistered", + Error::ExtensionHandleNotFound { .. } => "ExtensionHandleNotFound", + Error::UnknownExtension { .. } => "UnknownExtension", Error::TooManyNodes {} => "TooManyNodes", Error::UnknownExporter { .. } => "UnknownExporter", Error::UnknownNode { .. } => "UnknownNode", @@ -516,6 +594,18 @@ pub fn error_summary_from(err: &Error) -> ErrorSummary { message: error.clone(), source: (!source_detail.is_empty()).then(|| source_detail.clone()), }, + Error::ExtensionError { + extension, + kind, + error, + source_detail, + } => ErrorSummary::Node { + node: extension.name.to_string(), + node_kind: NodeKind::Extension, + error_kind: kind.to_string(), + message: error.clone(), + source: (!source_detail.is_empty()).then(|| source_detail.clone()), + }, _ => ErrorSummary::Pipeline { error_kind: err.variant_name(), message: err.to_string(), diff --git a/rust/otap-dataflow/crates/engine/src/exporter.rs b/rust/otap-dataflow/crates/engine/src/exporter.rs index 5bc275516b..e07aadbcc7 100644 --- a/rust/otap-dataflow/crates/engine/src/exporter.rs +++ b/rust/otap-dataflow/crates/engine/src/exporter.rs @@ -14,6 +14,7 @@ use crate::context::PipelineContext; use crate::control::{Controllable, NodeControlMsg, PipelineCtrlMsgSender}; use crate::entity_context::NodeTelemetryGuard; use crate::error::{Error, ExporterErrorKind}; +use crate::extensions::ExtensionRegistry; use crate::local::exporter as local; use crate::local::message::{LocalReceiver, LocalSender}; use crate::message; @@ -208,7 +209,7 @@ impl ExporterWrapper { .. } => { let (control_sender, control_receiver) = - wrap_control_channel_metrics::( + wrap_control_channel_metrics::>( &node_id, pipeline_ctx, channel_metrics, @@ -241,7 +242,7 @@ impl ExporterWrapper { .. } => { let (control_sender, control_receiver) = - wrap_control_channel_metrics::( + wrap_control_channel_metrics::>( &node_id, pipeline_ctx, channel_metrics, @@ -270,6 +271,7 @@ impl ExporterWrapper { self, pipeline_ctrl_msg_tx: PipelineCtrlMsgSender, metrics_reporter: MetricsReporter, + extension_registry: ExtensionRegistry, ) -> Result { match (self, metrics_reporter) { ( @@ -294,7 +296,9 @@ impl ExporterWrapper { .set_pipeline_ctrl_msg_sender(pipeline_ctrl_msg_tx); let message_channel = message::MessageChannel::new(Receiver::Local(control_receiver), pdata_rx); - exporter.start(message_channel, effect_handler).await + exporter + .start(message_channel, effect_handler, extension_registry) + .await } ( ExporterWrapper::Shared { @@ -317,7 +321,9 @@ impl ExporterWrapper { .core .set_pipeline_ctrl_msg_sender(pipeline_ctrl_msg_tx); let message_channel = shared::MessageChannel::new(control_receiver, pdata_rx); - exporter.start(message_channel, effect_handler).await + exporter + .start(message_channel, effect_handler, extension_registry) + .await } } } @@ -394,6 +400,7 @@ mod tests { use crate::control::{AckMsg, NodeControlMsg}; use crate::error::ExporterErrorKind; use crate::exporter::{Error, ExporterWrapper}; + use crate::extensions::ExtensionRegistry; use crate::local::exporter as local; use crate::local::message::LocalReceiver; use crate::message; @@ -434,6 +441,7 @@ mod tests { self: Box, mut msg_chan: message::MessageChannel, effect_handler: local::EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { // Loop until a Shutdown event is received. loop { @@ -471,6 +479,7 @@ mod tests { self: Box, mut msg_chan: shared::MessageChannel, effect_handler: shared::EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { // Loop until a Shutdown event is received. loop { @@ -822,3 +831,137 @@ mod tests { assert!(matches!(chan.recv().await, Err(RecvError::Closed))); } } + +/// Tests verifying that an exporter can retrieve and use a [`ClientAuthenticatorHandle`] +/// from the [`ExtensionRegistry`] to attach credentials to outgoing requests. +#[cfg(test)] +mod auth_extension_tests { + use super::ExporterWrapper; + use crate::error::{Error, ExporterErrorKind}; + use crate::extensions::auth::{AuthError, ClientAuthenticator, ClientAuthenticatorHandle}; + use crate::extensions::{ExtensionHandles, ExtensionRegistry, ExtensionRegistryBuilder}; + use crate::local::exporter as local; + use crate::message; + use crate::message::Message; + use crate::terminal_state::TerminalState; + use crate::testing::exporter::TestRuntime; + use crate::testing::{TestMsg, test_node}; + use async_trait::async_trait; + use otap_df_config::node::NodeUserConfig; + use std::future::Future; + use std::pin::Pin; + use std::sync::Arc; + use std::time::{Duration, Instant}; + + /// A client authenticator that attaches a static bearer token. + struct StaticBearerAuth { + token: String, + } + + impl ClientAuthenticator for StaticBearerAuth { + fn get_request_metadata( + &self, + ) -> Result, AuthError> { + Ok(vec![( + http::header::AUTHORIZATION, + http::HeaderValue::from_str(&format!("Bearer {}", self.token)).map_err(|e| { + AuthError { + message: e.to_string(), + } + })?, + )]) + } + } + + /// A minimal exporter that retrieves a client auth handle from the + /// extension registry, uses it to produce outgoing headers, and + /// logs the result via the PData message channel. No network I/O — + /// auth is exercised directly in `start()`. + struct AuthExporter; + + #[async_trait(?Send)] + impl local::Exporter for AuthExporter { + async fn start( + self: Box, + mut msg_chan: message::MessageChannel, + effect_handler: local::EffectHandler, + extension_registry: ExtensionRegistry, + ) -> Result { + // Retrieve the client auth handle — this is what a real exporter does. + let auth = extension_registry + .get::("bearer_auth") + .expect("bearer_auth extension must be registered"); + + // Simulate attaching credentials to an outgoing request. + let headers = auth + .get_request_metadata() + .map_err(|e| Error::ExporterError { + exporter: effect_handler.exporter_id(), + kind: ExporterErrorKind::Other, + error: e.message, + source_detail: String::new(), + })?; + + assert_eq!(headers.len(), 1); + let (name, value) = &headers[0]; + assert_eq!(name, http::header::AUTHORIZATION); + assert_eq!(value, "Bearer test-token-42"); + + // Report success through the counter (increment_message). + // Wait for shutdown. + loop { + match msg_chan.recv().await? { + Message::Control(crate::control::NodeControlMsg::Shutdown { .. }) => break, + Message::PData(_) => {} + _ => {} + } + } + Ok(TerminalState::default()) + } + } + + fn build_auth_registry() -> ExtensionRegistry { + let auth = StaticBearerAuth { + token: "test-token-42".into(), + }; + let mut handles = ExtensionHandles::new(); + handles.register(ClientAuthenticatorHandle::new(auth)); + let mut builder = ExtensionRegistryBuilder::new(); + builder.merge("bearer_auth", handles).unwrap(); + builder.build() + } + + /// Exporter retrieves a client auth handle from the extension registry and + /// uses it to produce credentials for outgoing requests. + #[test] + fn test_exporter_with_auth_extension() { + let test_runtime = TestRuntime::new(); + let user_config = Arc::new(NodeUserConfig::new_exporter_config("auth_exporter")); + let exporter = ExporterWrapper::local( + AuthExporter, + test_node("auth_exp".to_string()), + user_config, + test_runtime.config(), + ); + + let registry = build_auth_registry(); + + test_runtime + .set_exporter(exporter) + .with_extension_registry(registry) + .run_test(|ctx| -> Pin>> { + Box::pin(async move { + ctx.send_shutdown(Instant::now() + Duration::from_millis(200), "Test") + .await + .expect("shutdown failed"); + }) + }) + .run_validation(|_ctx, result| -> Pin>> { + Box::pin(async move { + // The exporter should have completed successfully — any auth + // failure would have returned an error from `start()`. + result.expect("exporter start() should succeed with valid auth"); + }) + }); + } +} diff --git a/rust/otap-dataflow/crates/engine/src/extension.rs b/rust/otap-dataflow/crates/engine/src/extension.rs new file mode 100644 index 0000000000..28bdb136af --- /dev/null +++ b/rust/otap-dataflow/crates/engine/src/extension.rs @@ -0,0 +1,433 @@ +// Copyright The OpenTelemetry Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Extension wrapper used to provide a unified interface to the pipeline engine that abstracts over +//! the fact that extension implementations may be `!Send` or `Send`. +//! +//! Extensions are **not** generic over `PData` — they sit outside the data-flow graph +//! and never process pipeline data. They use the PData-free [`ExtensionControlMsg`] +//! for their control channels. +//! +//! For more details on the `!Send` implementation of an extension, see [`local::extension::Extension`]. +//! See [`shared::extension::Extension`] for the Send implementation. + +use crate::channel_metrics::ChannelMetricsRegistry; +use crate::channel_mode::{LocalMode, SharedMode, wrap_control_channel_metrics}; +use crate::config::ExtensionConfig; +use crate::context::PipelineContext; +use crate::control::ExtensionControlMsg; +use crate::entity_context::NodeTelemetryGuard; +use crate::error::Error; +use crate::extensions::ExtensionHandles; +use crate::local::extension as local; +use crate::local::message::{LocalReceiver, LocalSender}; +use crate::message::Sender; +use crate::node::NodeId; +use crate::shared::extension as shared; +use crate::shared::message::{SharedReceiver, SharedSender}; +use otap_df_channel::mpsc; +use otap_df_config::node::NodeUserConfig; +use otap_df_telemetry::reporter::MetricsReporter; +use std::sync::Arc; + +/// A wrapper for extensions that allows for both `Send` and `!Send` effect handlers. +/// +/// Unlike the pipeline node wrappers, `ExtensionWrapper` is **not** generic over `PData`. +pub enum ExtensionWrapper { + /// An extension with a `!Send` implementation. + Local { + /// Index identifier for the node. + node_id: NodeId, + /// The user configuration for the node. + user_config: Arc, + /// The runtime configuration for the extension. + runtime_config: ExtensionConfig, + /// The extension instance. + extension: Box, + /// A sender for control messages. + control_sender: LocalSender, + /// A receiver for control messages. + control_receiver: LocalReceiver, + /// Service handles produced by this extension, to be merged into the registry. + handles: Option, + /// Telemetry guard for node lifecycle cleanup. + telemetry: Option, + }, + /// An extension with a `Send` implementation. + Shared { + /// Index identifier for the node. + node_id: NodeId, + /// The user configuration for the node. + user_config: Arc, + /// The runtime configuration for the extension. + runtime_config: ExtensionConfig, + /// The extension instance. + extension: Box, + /// A sender for control messages. + control_sender: SharedSender, + /// A receiver for control messages. + control_receiver: SharedReceiver, + /// Service handles produced by this extension, to be merged into the registry. + handles: Option, + /// Telemetry guard for node lifecycle cleanup. + telemetry: Option, + }, +} + +impl ExtensionWrapper { + /// Creates a new local `ExtensionWrapper` with the given extension and configuration. + pub fn local( + extension: E, + handles: ExtensionHandles, + node_id: NodeId, + user_config: Arc, + config: &ExtensionConfig, + ) -> Self + where + E: local::Extension + 'static, + { + let (control_sender, control_receiver) = + mpsc::Channel::new(config.control_channel.capacity); + + ExtensionWrapper::Local { + node_id, + user_config, + runtime_config: config.clone(), + extension: Box::new(extension), + control_sender: LocalSender::mpsc(control_sender), + control_receiver: LocalReceiver::mpsc(control_receiver), + handles: Some(handles), + telemetry: None, + } + } + + /// Creates a new shared `ExtensionWrapper` with the given extension and configuration. + pub fn shared( + extension: E, + handles: ExtensionHandles, + node_id: NodeId, + user_config: Arc, + config: &ExtensionConfig, + ) -> Self + where + E: shared::Extension + 'static, + { + let (control_sender, control_receiver) = + tokio::sync::mpsc::channel(config.control_channel.capacity); + + ExtensionWrapper::Shared { + node_id, + user_config, + runtime_config: config.clone(), + extension: Box::new(extension), + control_sender: SharedSender::mpsc(control_sender), + control_receiver: SharedReceiver::mpsc(control_receiver), + handles: Some(handles), + telemetry: None, + } + } + + /// Takes the service handles out of this wrapper. + /// + /// Called during pipeline build to transfer handles into the + /// `ExtensionRegistryBuilder` (crate-internal). + /// Returns `None` if handles have already been taken. + pub fn take_handles(&mut self) -> Option { + match self { + ExtensionWrapper::Local { handles, .. } => handles.take(), + ExtensionWrapper::Shared { handles, .. } => handles.take(), + } + } + + /// Returns the node id of the extension. + #[must_use] + pub fn node_id(&self) -> NodeId { + match self { + ExtensionWrapper::Local { node_id, .. } => node_id.clone(), + ExtensionWrapper::Shared { node_id, .. } => node_id.clone(), + } + } + + /// Returns the extension name (from the user config URN). + #[must_use] + pub fn name(&self) -> &str { + match self { + ExtensionWrapper::Local { user_config, .. } => user_config.r#type.as_ref(), + ExtensionWrapper::Shared { user_config, .. } => user_config.r#type.as_ref(), + } + } + + pub(crate) fn with_node_telemetry_guard(self, guard: NodeTelemetryGuard) -> Self { + match self { + ExtensionWrapper::Local { + node_id, + user_config, + runtime_config, + extension, + control_sender, + control_receiver, + handles, + .. + } => ExtensionWrapper::Local { + node_id, + user_config, + runtime_config, + extension, + control_sender, + control_receiver, + handles, + telemetry: Some(guard), + }, + ExtensionWrapper::Shared { + node_id, + user_config, + runtime_config, + extension, + control_sender, + control_receiver, + handles, + .. + } => ExtensionWrapper::Shared { + node_id, + user_config, + runtime_config, + extension, + control_sender, + control_receiver, + handles, + telemetry: Some(guard), + }, + } + } + + pub(crate) const fn take_telemetry_guard(&mut self) -> Option { + match self { + ExtensionWrapper::Local { telemetry, .. } => telemetry.take(), + ExtensionWrapper::Shared { telemetry, .. } => telemetry.take(), + } + } + + pub(crate) fn with_control_channel_metrics( + self, + pipeline_ctx: &PipelineContext, + channel_metrics: &mut ChannelMetricsRegistry, + channel_metrics_enabled: bool, + ) -> Self { + match self { + ExtensionWrapper::Local { + node_id, + runtime_config, + control_sender, + control_receiver, + user_config, + extension, + handles, + telemetry, + } => { + let (control_sender, control_receiver) = + wrap_control_channel_metrics::( + &node_id, + pipeline_ctx, + channel_metrics, + channel_metrics_enabled, + runtime_config.control_channel.capacity as u64, + control_sender, + control_receiver, + ); + + ExtensionWrapper::Local { + node_id, + user_config, + runtime_config, + extension, + control_sender, + control_receiver, + handles, + telemetry, + } + } + ExtensionWrapper::Shared { + node_id, + runtime_config, + control_sender, + control_receiver, + user_config, + extension, + handles, + telemetry, + } => { + let (control_sender, control_receiver) = + wrap_control_channel_metrics::( + &node_id, + pipeline_ctx, + channel_metrics, + channel_metrics_enabled, + runtime_config.control_channel.capacity as u64, + control_sender, + control_receiver, + ); + + ExtensionWrapper::Shared { + node_id, + user_config, + runtime_config, + extension, + control_sender, + control_receiver, + handles, + telemetry, + } + } + } + } + + /// Returns a clone of the control message sender for this extension. + /// + /// This must be called **before** [`start`](Self::start), which consumes `self`. + /// The returned sender is held by the [`PipelineCtrlMsgManager`](crate::pipeline_ctrl::PipelineCtrlMsgManager) + /// so the extension can receive shutdown and other control messages. + pub fn control_sender(&self) -> Sender { + match self { + ExtensionWrapper::Local { control_sender, .. } => Sender::Local(control_sender.clone()), + ExtensionWrapper::Shared { control_sender, .. } => { + Sender::Shared(control_sender.clone()) + } + } + } + + /// Starts the extension's background work. + /// + /// The extension task runs independently, processing only control messages. + /// Unlike pipeline node wrappers, extensions do not receive a `PipelineCtrlMsgSender` + /// — they manage their own timers via `tokio::time` if needed. + pub async fn start(self, metrics_reporter: MetricsReporter) -> Result<(), Error> { + match self { + ExtensionWrapper::Local { + node_id, + extension, + control_receiver, + .. + } => { + let effect_handler = local::EffectHandler::new(node_id, metrics_reporter); + let ctrl_chan = local::ControlChannel::new(control_receiver); + extension.start(ctrl_chan, effect_handler).await + } + ExtensionWrapper::Shared { + node_id, + extension, + control_receiver, + .. + } => { + let effect_handler = shared::EffectHandler::new(node_id, metrics_reporter); + let ctrl_chan = shared::ControlChannel::new(control_receiver); + extension.start(ctrl_chan, effect_handler).await + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::ExtensionConfig; + use crate::control::ExtensionControlMsg; + use crate::extensions::ExtensionHandles; + use crate::local::extension as local; + use crate::testing::test_node; + use async_trait::async_trait; + use otap_df_config::node::NodeUserConfig; + use otap_df_telemetry::reporter::MetricsReporter; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::time::{Duration, Instant}; + + /// A minimal extension that loops on its control channel and sets a flag + /// when it receives a Shutdown message. + struct ShutdownTracker { + received: Arc, + } + + #[async_trait(?Send)] + impl local::Extension for ShutdownTracker { + async fn start( + self: Box, + mut ctrl_chan: local::ControlChannel, + _effect_handler: local::EffectHandler, + ) -> Result<(), Error> { + loop { + match ctrl_chan.recv().await { + Ok(ExtensionControlMsg::Shutdown { .. }) => { + self.received.store(true, Ordering::SeqCst); + break; + } + Ok(_) => {} // ignore other messages + Err(_) => break, + } + } + Ok(()) + } + } + + /// Verifies that a Shutdown message sent through the cloned control_sender + /// is received by the extension's start() loop. + #[test] + fn test_extension_receives_shutdown_via_control_sender() { + let shutdown_received = Arc::new(AtomicBool::new(false)); + let tracker = ShutdownTracker { + received: shutdown_received.clone(), + }; + + let config = ExtensionConfig::new("shutdown_ext"); + let user_config = Arc::new(NodeUserConfig::new_receiver_config("test_ext")); + let ext = ExtensionWrapper::local( + tracker, + ExtensionHandles::new(), + test_node("shutdown_ext"), + user_config, + &config, + ); + + // Clone the sender BEFORE start() consumes the wrapper. + let sender = ext.control_sender(); + + let (_metrics_rx, metrics_reporter) = MetricsReporter::create_new_and_receiver(1); + + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + rt.block_on(async { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let handle = tokio::task::spawn_local(async move { + ext.start(metrics_reporter).await.expect("extension failed"); + }); + + // Give the extension a moment to start its recv loop. + tokio::time::sleep(Duration::from_millis(10)).await; + + // Send shutdown through the cloned sender. + sender + .send(ExtensionControlMsg::Shutdown { + deadline: Instant::now(), + reason: "test shutdown".to_owned(), + }) + .await + .expect("send failed"); + + tokio::time::timeout(Duration::from_secs(2), handle) + .await + .expect("extension did not shut down in time") + .expect("join error"); + }) + .await; + }); + + assert!( + shutdown_received.load(Ordering::SeqCst), + "extension should have received the Shutdown message" + ); + } +} diff --git a/rust/otap-dataflow/crates/engine/src/extensions/auth.rs b/rust/otap-dataflow/crates/engine/src/extensions/auth.rs new file mode 100644 index 0000000000..cdceecd111 --- /dev/null +++ b/rust/otap-dataflow/crates/engine/src/extensions/auth.rs @@ -0,0 +1,505 @@ +// Copyright The OpenTelemetry Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Trait-based authentication handles for extensions. +//! +//! This module defines a generic, pluggable authentication contract: +//! +//! - [`ServerAuthenticator`] + [`ServerAuthenticatorHandle`] — for **receivers** +//! that need to validate credentials on incoming requests. +//! - [`ClientAuthenticator`] + [`ClientAuthenticatorHandle`] — for **exporters** +//! that need to attach credentials to outgoing requests. +//! +//! These are the only auth types that receivers and exporters need to know about. +//! Concrete auth strategies (bearer tokens, API keys, OIDC, etc.) implement +//! the traits and are selected purely through configuration. +//! +//! # Examples +//! +//! ## Extension factory — registering both handles +//! +//! ```rust,ignore +//! let auth = MyAuthImpl { /* ... */ }; +//! +//! let mut handles = ExtensionHandles::new(); +//! handles.register(ServerAuthenticatorHandle::new(auth.clone())); +//! handles.register(ClientAuthenticatorHandle::new(auth)); +//! ``` +//! +//! ## Receiver — validating incoming requests +//! +//! ```rust,ignore +//! // Look up by config node name (the key in the YAML `nodes:` map), +//! // not the plugin URN. +//! let auth = extension_registry +//! .get::("my_auth")?; +//! +//! // In the gRPC/HTTP handler: +//! auth.authenticate(request.headers())?; +//! ``` +//! +//! ## Exporter — attaching outgoing credentials +//! +//! ```rust,ignore +//! let auth = extension_registry +//! .get::("my_auth")?; +//! +//! for (key, value) in auth.get_request_metadata()? { +//! request.headers_mut().insert(key, value); +//! } +//! ``` + +use std::fmt; +use std::sync::{Arc, Mutex}; + +/// An error returned by authenticator operations. +#[derive(Debug, Clone)] +pub struct AuthError { + /// A human-readable description of the authentication failure. + pub message: String, +} + +impl fmt::Display for AuthError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "authentication error: {}", self.message) + } +} + +impl std::error::Error for AuthError {} + +// ─── Server side (receivers validate incoming requests) ──────────────────── + +/// Trait for validating credentials on incoming requests. +/// +/// Implement this trait in an auth extension to provide server-side +/// authentication. Different strategies (bearer token validation, API key +/// allow-lists, OIDC token verification) all implement this same interface. +/// +/// Receivers call [`ServerAuthenticatorHandle::authenticate`] without knowing +/// which concrete strategy is behind it — swapping auth is a config change. +pub trait ServerAuthenticator: Send { + /// Validates the request headers. + /// + /// Returns `Ok(())` if the request is authenticated, or an [`AuthError`] + /// describing why authentication failed. + fn authenticate(&self, headers: &http::HeaderMap) -> Result<(), AuthError>; +} + +/// A cloneable handle that receivers use to authenticate incoming requests. +/// +/// This wraps any [`ServerAuthenticator`] behind an `Arc>` so that +/// each receiver gets its own clone. The `Mutex` makes the handle `Sync` +/// (required by tonic services) without requiring `Sync` on the trait itself. +/// The lock is never contended because the engine uses a thread-per-core +/// architecture in both local and shared modes. +#[derive(Clone)] +pub struct ServerAuthenticatorHandle { + inner: Arc>>, +} + +impl ServerAuthenticatorHandle { + /// Creates a new handle wrapping the given authenticator implementation. + pub fn new(auth: impl ServerAuthenticator + 'static) -> Self { + Self { + inner: Arc::new(Mutex::new(Box::new(auth))), + } + } + + /// Validates the request headers. + /// + /// Delegates to the underlying [`ServerAuthenticator`] implementation. + pub fn authenticate(&self, headers: &http::HeaderMap) -> Result<(), AuthError> { + self.inner + .lock() + .expect("ServerAuthenticator lock poisoned") + .authenticate(headers) + } +} + +impl fmt::Debug for ServerAuthenticatorHandle { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ServerAuthenticatorHandle").finish() + } +} + +// ─── Client side (exporters attach outgoing credentials) ─────────────────── + +/// Trait for producing credentials to attach to outgoing requests. +/// +/// Implement this trait in an auth extension to provide client-side +/// authentication. The extension decides what headers to attach +/// (e.g., `Authorization: Bearer `, custom API key headers). +pub trait ClientAuthenticator: Send { + /// Returns the headers to attach to an outgoing request. + /// + /// Each entry is a `(header_name, header_value)` pair. The exporter + /// inserts them into the request's header map before sending. + /// + /// # Errors + /// + /// Returns an [`AuthError`] if credentials are unavailable + /// (e.g., token not yet refreshed, provider unreachable). + fn get_request_metadata(&self) + -> Result, AuthError>; +} + +/// A cloneable handle that exporters use to attach credentials to outgoing requests. +/// +/// This wraps any [`ClientAuthenticator`] behind an `Arc>` so that +/// each exporter gets its own clone. The `Mutex` makes the handle `Sync` +/// (required by tonic services) without requiring `Sync` on the trait itself. +/// The lock is never contended because the engine uses a thread-per-core +/// architecture in both local and shared modes. +#[derive(Clone)] +pub struct ClientAuthenticatorHandle { + inner: Arc>>, +} + +impl ClientAuthenticatorHandle { + /// Creates a new handle wrapping the given authenticator implementation. + pub fn new(auth: impl ClientAuthenticator + 'static) -> Self { + Self { + inner: Arc::new(Mutex::new(Box::new(auth))), + } + } + + /// Returns the headers to attach to an outgoing request. + /// + /// Delegates to the underlying [`ClientAuthenticator`] implementation. + pub fn get_request_metadata( + &self, + ) -> Result, AuthError> { + self.inner + .lock() + .expect("ClientAuthenticator lock poisoned") + .get_request_metadata() + } +} + +impl fmt::Debug for ClientAuthenticatorHandle { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ClientAuthenticatorHandle").finish() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A trivial static bearer token authenticator that implements both traits. + struct StaticBearerAuth { + token: String, + } + + impl ServerAuthenticator for StaticBearerAuth { + fn authenticate(&self, headers: &http::HeaderMap) -> Result<(), AuthError> { + let auth_value = headers + .get(http::header::AUTHORIZATION) + .and_then(|v| v.to_str().ok()) + .ok_or_else(|| AuthError { + message: "missing Authorization header".into(), + })?; + + let expected = format!("Bearer {}", self.token); + if auth_value != expected { + return Err(AuthError { + message: "invalid bearer token".into(), + }); + } + Ok(()) + } + } + + impl ClientAuthenticator for StaticBearerAuth { + fn get_request_metadata( + &self, + ) -> Result, AuthError> { + Ok(vec![( + http::header::AUTHORIZATION, + http::HeaderValue::from_str(&format!("Bearer {}", self.token)).map_err(|e| { + AuthError { + message: e.to_string(), + } + })?, + )]) + } + } + + #[test] + fn server_auth_valid_token() { + let auth = ServerAuthenticatorHandle::new(StaticBearerAuth { + token: "secret123".into(), + }); + + let mut headers = http::HeaderMap::new(); + let _ = headers.insert( + http::header::AUTHORIZATION, + "Bearer secret123".parse().unwrap(), + ); + assert!(auth.authenticate(&headers).is_ok()); + } + + #[test] + fn server_auth_invalid_token() { + let auth = ServerAuthenticatorHandle::new(StaticBearerAuth { + token: "secret123".into(), + }); + + let mut headers = http::HeaderMap::new(); + let _ = headers.insert(http::header::AUTHORIZATION, "Bearer wrong".parse().unwrap()); + let err = auth.authenticate(&headers).unwrap_err(); + assert!(err.message.contains("invalid")); + } + + #[test] + fn server_auth_missing_header() { + let auth = ServerAuthenticatorHandle::new(StaticBearerAuth { + token: "secret123".into(), + }); + + let headers = http::HeaderMap::new(); + let err = auth.authenticate(&headers).unwrap_err(); + assert!(err.message.contains("missing")); + } + + #[test] + fn client_auth_produces_metadata() { + let auth = ClientAuthenticatorHandle::new(StaticBearerAuth { + token: "mytoken".into(), + }); + + let metadata = auth.get_request_metadata().unwrap(); + assert_eq!(metadata.len(), 1); + assert_eq!(metadata[0].0, http::header::AUTHORIZATION); + assert_eq!(metadata[0].1, "Bearer mytoken"); + } + + #[test] + fn separate_handles_from_same_type() { + let server = ServerAuthenticatorHandle::new(StaticBearerAuth { + token: "shared".into(), + }); + let client = ClientAuthenticatorHandle::new(StaticBearerAuth { + token: "shared".into(), + }); + + // Server side validates + let mut headers = http::HeaderMap::new(); + let _ = headers.insert( + http::header::AUTHORIZATION, + "Bearer shared".parse().unwrap(), + ); + assert!(server.authenticate(&headers).is_ok()); + + // Client side produces the same token + let metadata = client.get_request_metadata().unwrap(); + assert_eq!(metadata[0].1, "Bearer shared"); + } +} + +/// End-to-end scenario tests demonstrating realistic auth extension +/// patterns: a receiver-side header allow-list and an exporter-side +/// token refresher backed by a shared `Arc>`. +#[cfg(test)] +mod scenario_tests { + use super::*; + use crate::extensions::{ExtensionHandles, ExtensionRegistryBuilder}; + + // ─── Scenario 1: Header allow-list (receiver-side) ──────────── + + /// A server authenticator that checks a specific header is present + /// and its value belongs to a known allow-list. + struct HeaderAllowListAuth { + header_name: http::HeaderName, + allowed_values: Vec, + } + + impl ServerAuthenticator for HeaderAllowListAuth { + fn authenticate(&self, headers: &http::HeaderMap) -> Result<(), AuthError> { + let value = headers + .get(&self.header_name) + .and_then(|v| v.to_str().ok()) + .ok_or_else(|| AuthError { + message: format!("missing required header: {}", self.header_name), + })?; + + if !self.allowed_values.iter().any(|allowed| allowed == value) { + return Err(AuthError { + message: format!("header value '{}' is not in the allow-list", value), + }); + } + Ok(()) + } + } + + #[test] + fn header_allowlist_valid_value() { + let auth = ServerAuthenticatorHandle::new(HeaderAllowListAuth { + header_name: http::HeaderName::from_static("x-tenant-id"), + allowed_values: vec!["tenant-a".into(), "tenant-b".into()], + }); + + let mut headers = http::HeaderMap::new(); + let _ = headers.insert("x-tenant-id", "tenant-a".parse().unwrap()); + assert!(auth.authenticate(&headers).is_ok()); + } + + #[test] + fn header_allowlist_invalid_value() { + let auth = ServerAuthenticatorHandle::new(HeaderAllowListAuth { + header_name: http::HeaderName::from_static("x-tenant-id"), + allowed_values: vec!["tenant-a".into(), "tenant-b".into()], + }); + + let mut headers = http::HeaderMap::new(); + let _ = headers.insert("x-tenant-id", "tenant-unknown".parse().unwrap()); + let err = auth.authenticate(&headers).unwrap_err(); + assert!(err.message.contains("not in the allow-list")); + } + + #[test] + fn header_allowlist_missing_header() { + let auth = ServerAuthenticatorHandle::new(HeaderAllowListAuth { + header_name: http::HeaderName::from_static("x-tenant-id"), + allowed_values: vec!["tenant-a".into()], + }); + + let headers = http::HeaderMap::new(); + let err = auth.authenticate(&headers).unwrap_err(); + assert!(err.message.contains("missing required header")); + } + + #[test] + fn header_allowlist_via_registry() { + let auth = HeaderAllowListAuth { + header_name: http::HeaderName::from_static("x-tenant-id"), + allowed_values: vec!["tenant-a".into(), "tenant-b".into()], + }; + + // Extension factory registers the handle + let mut handles = ExtensionHandles::new(); + handles.register(ServerAuthenticatorHandle::new(auth)); + + let mut builder = ExtensionRegistryBuilder::new(); + builder.merge("header_allowlist", handles).unwrap(); + let registry = builder.build(); + + // Receiver retrieves it by name + type at startup + let handle = registry + .get::("header_allowlist") + .unwrap(); + + let mut headers = http::HeaderMap::new(); + let _ = headers.insert("x-tenant-id", "tenant-b".parse().unwrap()); + assert!(handle.authenticate(&headers).is_ok()); + + let _ = headers.insert("x-tenant-id", "tenant-c".parse().unwrap()); + assert!(handle.authenticate(&headers).is_err()); + } + + // ─── Scenario 2: Shared-state bearer token (exporter-side) ────── + + struct SharedTokenAuth { + token: Arc>, + } + + impl ClientAuthenticator for SharedTokenAuth { + fn get_request_metadata( + &self, + ) -> Result, AuthError> { + let token = self.token.lock().expect("token lock poisoned").clone(); + if token.is_empty() { + return Err(AuthError { + message: "token not yet available".into(), + }); + } + Ok(vec![( + http::header::AUTHORIZATION, + http::HeaderValue::from_str(&format!("Bearer {}", token)).map_err(|e| { + AuthError { + message: e.to_string(), + } + })?, + )]) + } + } + + #[test] + fn shared_token_initial_empty_token_fails() { + let token = Arc::new(Mutex::new(String::new())); + let auth = ClientAuthenticatorHandle::new(SharedTokenAuth { token }); + + let err = auth.get_request_metadata().unwrap_err(); + assert!(err.message.contains("not yet available")); + } + + #[test] + fn shared_token_returns_current_token() { + let token = Arc::new(Mutex::new("initial-token".to_string())); + let auth = ClientAuthenticatorHandle::new(SharedTokenAuth { + token: token.clone(), + }); + + let metadata = auth.get_request_metadata().unwrap(); + assert_eq!(metadata[0].1, "Bearer initial-token"); + + // Simulate token refresh by the extension's start() task + *token.lock().unwrap() = "refreshed-token".to_string(); + + let metadata = auth.get_request_metadata().unwrap(); + assert_eq!(metadata[0].1, "Bearer refreshed-token"); + } + + #[test] + fn shared_token_cloned_handle_sees_updates() { + let token = Arc::new(Mutex::new("v1".to_string())); + let auth = ClientAuthenticatorHandle::new(SharedTokenAuth { + token: token.clone(), + }); + + // Clone the handle (simulating the registry cloning for multiple exporters) + let auth2 = auth.clone(); + + let m1 = auth.get_request_metadata().unwrap(); + let m2 = auth2.get_request_metadata().unwrap(); + assert_eq!(m1[0].1, "Bearer v1"); + assert_eq!(m2[0].1, "Bearer v1"); + + // Refresh token — both clones see the update + *token.lock().unwrap() = "v2".to_string(); + + let m1 = auth.get_request_metadata().unwrap(); + let m2 = auth2.get_request_metadata().unwrap(); + assert_eq!(m1[0].1, "Bearer v2"); + assert_eq!(m2[0].1, "Bearer v2"); + } + + #[test] + fn shared_token_via_registry() { + let token = Arc::new(Mutex::new("tok-abc".to_string())); + let auth = SharedTokenAuth { + token: token.clone(), + }; + + // Extension factory registers the handle + let mut handles = ExtensionHandles::new(); + handles.register(ClientAuthenticatorHandle::new(auth)); + + let mut builder = ExtensionRegistryBuilder::new(); + builder.merge("token_refresher", handles).unwrap(); + let registry = builder.build(); + + // Exporter retrieves it by name + type at startup + let handle = registry + .get::("token_refresher") + .unwrap(); + + let metadata = handle.get_request_metadata().unwrap(); + assert_eq!(metadata[0].1, "Bearer tok-abc"); + + // Token refresh propagates through the registry-retrieved handle + *token.lock().unwrap() = "tok-xyz".to_string(); + let metadata = handle.get_request_metadata().unwrap(); + assert_eq!(metadata[0].1, "Bearer tok-xyz"); + } +} diff --git a/rust/otap-dataflow/crates/engine/src/extensions/bearer_token.rs b/rust/otap-dataflow/crates/engine/src/extensions/bearer_token.rs new file mode 100644 index 0000000000..d2d8f44213 --- /dev/null +++ b/rust/otap-dataflow/crates/engine/src/extensions/bearer_token.rs @@ -0,0 +1,316 @@ +// Copyright The OpenTelemetry Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Bearer token provider handle for extensions. +//! +//! This module defines a token provider contract for extensions that manage +//! bearer authentication tokens (e.g., Azure Managed Identity, OAuth2 flows): +//! +//! - [`BearerTokenProvider`] — trait for components that acquire and refresh tokens. +//! - [`BearerTokenProviderHandle`] — a cloneable handle that consumers use to +//! obtain tokens and subscribe to refresh events. +//! +//! # Examples +//! +//! ## Extension factory — registering the handle +//! +//! ```rust,ignore +//! let provider = MyTokenProvider { /* ... */ }; +//! +//! let mut handles = ExtensionHandles::new(); +//! handles.register(BearerTokenProviderHandle::new(provider)); +//! ``` +//! +//! ## Exporter — obtaining a token +//! +//! ```rust,ignore +//! let token_handle = extension_registry +//! .get::("my_auth")?; +//! +//! let token = token_handle.get_token().await?; +//! request.headers_mut().insert( +//! http::header::AUTHORIZATION, +//! format!("Bearer {}", token.token.secret()).parse().unwrap(), +//! ); +//! ``` +//! +//! ## Subscribing to token refresh +//! +//! ```rust,ignore +//! let token_handle = extension_registry +//! .get::("my_auth")?; +//! +//! let mut token_rx = token_handle.subscribe_token_refresh().await; +//! loop { +//! tokio::select! { +//! _ = token_rx.changed() => { +//! if let Some(token) = token_rx.borrow().as_ref() { +//! // Update headers, etc. +//! } +//! } +//! } +//! } +//! ``` + +use async_trait::async_trait; +use std::borrow::Cow; +use std::fmt; +use std::sync::Arc; + +// ─── Secret ──────────────────────────────────────────────────────────────── + +/// Represents a secret value that should not be exposed in logs or debug output. +/// +/// The [`Debug`] implementation will not print the actual secret value. +#[derive(Clone, Eq)] +pub struct Secret(Cow<'static, str>); + +impl Secret { + /// Creates a new `Secret`. + #[must_use] + pub fn new(value: T) -> Self + where + T: Into>, + { + Self(value.into()) + } + + /// Returns the secret value. + #[must_use] + pub fn secret(&self) -> &str { + &self.0 + } +} + +impl PartialEq for Secret { + fn eq(&self, other: &Self) -> bool { + self.secret() == other.secret() + } +} + +impl From for Secret { + fn from(value: String) -> Self { + Self::new(value) + } +} + +impl From<&'static str> for Secret { + fn from(value: &'static str) -> Self { + Self::new(value) + } +} + +impl fmt::Debug for Secret { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("Secret") + } +} + +// ─── BearerToken ─────────────────────────────────────────────────────────── + +/// Represents a bearer token with its expiration time. +/// +/// The token value is wrapped in [`Secret`] to prevent accidental exposure +/// in logs or debug output. +#[derive(Debug, Clone)] +pub struct BearerToken { + /// The token value. + pub token: Secret, + /// The expiration time as a UNIX timestamp (seconds since epoch). + pub expires_on: i64, +} + +impl BearerToken { + /// Creates a new bearer token. + #[must_use] + pub fn new(token: T, expires_on: i64) -> Self + where + T: Into, + { + Self { + token: token.into(), + expires_on, + } + } +} + +// ─── Error ───────────────────────────────────────────────────────────────── + +/// An error returned by bearer token provider operations. +#[derive(Debug, Clone)] +pub struct BearerTokenError { + /// A human-readable description of the failure. + pub message: String, +} + +impl fmt::Display for BearerTokenError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "bearer token error: {}", self.message) + } +} + +impl std::error::Error for BearerTokenError {} + +// ─── Trait ───────────────────────────────────────────────────────────────── + +/// A trait for components that can provide bearer authentication tokens. +/// +/// Extensions implementing this trait can be looked up by other components +/// (e.g., exporters) to obtain tokens for authentication. +/// +/// The extension background task handles periodic token refresh. Consumers +/// can either call [`get_token`](BearerTokenProvider::get_token) on demand +/// or subscribe to refresh notifications via +/// [`subscribe_token_refresh`](BearerTokenProvider::subscribe_token_refresh). +#[async_trait] +pub trait BearerTokenProvider: Send { + /// Returns an authentication token. + /// + /// # Errors + /// + /// Returns a [`BearerTokenError`] if the token cannot be obtained. + async fn get_token(&self) -> Result; + + /// Subscribes to token refresh events. + /// + /// Returns a new receiver that will be notified whenever the token + /// is refreshed. Each call creates an independent subscription. + /// The receiver always contains the latest token value (or `None` + /// if no token has been acquired yet). + fn subscribe_token_refresh(&self) -> tokio::sync::watch::Receiver>; +} + +// ─── Handle ──────────────────────────────────────────────────────────────── + +/// A cloneable handle that consumers use to obtain bearer tokens. +/// +/// This wraps any [`BearerTokenProvider`] behind an `Arc>` +/// because [`get_token`](BearerTokenProvider::get_token) is async — a +/// `std::sync::Mutex` cannot be held across `.await` points. Each consumer +/// gets its own clone. The `tokio::Mutex` makes the handle `Sync` without +/// requiring `Sync` on the trait itself. +#[derive(Clone)] +pub struct BearerTokenProviderHandle { + inner: Arc>>, +} + +impl BearerTokenProviderHandle { + /// Creates a new handle wrapping the given provider implementation. + pub fn new(provider: impl BearerTokenProvider + 'static) -> Self { + Self { + inner: Arc::new(tokio::sync::Mutex::new(Box::new(provider))), + } + } + + /// Returns an authentication token. + /// + /// Acquires the internal lock, then delegates to the underlying + /// [`BearerTokenProvider`] implementation. + pub async fn get_token(&self) -> Result { + self.inner.lock().await.get_token().await + } + + /// Subscribes to token refresh events. + /// + /// Acquires the internal lock, then delegates to the underlying + /// [`BearerTokenProvider`] implementation. + pub async fn subscribe_token_refresh( + &self, + ) -> tokio::sync::watch::Receiver> { + self.inner.lock().await.subscribe_token_refresh() + } +} + +impl fmt::Debug for BearerTokenProviderHandle { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BearerTokenProviderHandle").finish() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tokio::sync::watch; + + /// A trivial in-memory token provider for testing. + struct StaticTokenProvider { + token: String, + expires_on: i64, + sender: Arc>>, + } + + #[async_trait] + impl BearerTokenProvider for StaticTokenProvider { + async fn get_token(&self) -> Result { + Ok(BearerToken::new(self.token.clone(), self.expires_on)) + } + + fn subscribe_token_refresh(&self) -> watch::Receiver> { + self.sender.subscribe() + } + } + + fn make_static_provider(token: &str, expires_on: i64) -> StaticTokenProvider { + let (sender, _) = watch::channel(None); + StaticTokenProvider { + token: token.to_owned(), + expires_on, + sender: Arc::new(sender), + } + } + + #[tokio::test] + async fn handle_get_token() { + let handle = + BearerTokenProviderHandle::new(make_static_provider("test-token", 1_700_000_000)); + + let token = handle.get_token().await.unwrap(); + assert_eq!(token.token.secret(), "test-token"); + assert_eq!(token.expires_on, 1_700_000_000); + } + + #[tokio::test] + async fn handle_subscribe_receives_updates() { + let (sender, _) = watch::channel(None); + let sender = Arc::new(sender); + + let provider = StaticTokenProvider { + token: "initial".to_owned(), + expires_on: 100, + sender: Arc::clone(&sender), + }; + + let handle = BearerTokenProviderHandle::new(provider); + let mut rx = handle.subscribe_token_refresh().await; + + // Simulate a token refresh from the extension background task. + let _ = sender.send(Some(BearerToken::new("refreshed", 200))); + + rx.changed().await.unwrap(); + let refreshed = rx.borrow().clone().unwrap(); + assert_eq!(refreshed.token.secret(), "refreshed"); + assert_eq!(refreshed.expires_on, 200); + } + + #[test] + fn secret_debug_does_not_leak() { + let s = Secret::new("super-secret-value"); + assert_eq!(format!("{:?}", s), "Secret"); + } + + #[test] + fn secret_equality() { + let a = Secret::new("same"); + let b = Secret::new("same"); + let c = Secret::new("different"); + assert_eq!(a, b); + assert_ne!(a, c); + } + + #[test] + fn bearer_token_from_string() { + let token = BearerToken::new("my-token".to_string(), 42); + assert_eq!(token.token.secret(), "my-token"); + assert_eq!(token.expires_on, 42); + } +} diff --git a/rust/otap-dataflow/crates/engine/src/extensions/mod.rs b/rust/otap-dataflow/crates/engine/src/extensions/mod.rs new file mode 100644 index 0000000000..32ba68e384 --- /dev/null +++ b/rust/otap-dataflow/crates/engine/src/extensions/mod.rs @@ -0,0 +1,43 @@ +// Copyright The OpenTelemetry Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Extension registry and built-in handle types. +//! +//! Extensions are non-pipeline components that provide cross-cutting capabilities +//! (e.g., authentication, health checks) to pipeline nodes. An extension produces +//! one or more **service handles** — lightweight, cloneable values that pipeline +//! components use to interact with the extension at runtime. +//! +//! # Design Principles +//! +//! * **No `Sync` bounds** — handles are `Clone + Send` so each component owns its +//! own copy. There is no shared mutable state between threads. +//! * **No `unsafe` code** — the registry stores handles as `Box` +//! and retrieves them via standard `Any::downcast_ref` + `Clone`. +//! * **Channel-based communication** — handles typically wrap a `tokio::sync::watch` +//! receiver or similar primitive. The extension task owns the sender end. +//! * **Lifecycle ordering** — extension tasks start before pipeline components and +//! shut down after them, guaranteeing that handles remain valid for the duration +//! of the pipeline. +//! +//! # Adding a new handle type +//! +//! 1. Define a concrete struct implementing `Clone + Send + 'static`. +//! 2. In the extension factory, create both the extension task and the handle, +//! then register the handle via [`ExtensionHandles::register`]. +//! 3. Pipeline components retrieve the handle at start-up via +//! `effect_handler.get_extension_handle::("extension_name")`. + +pub mod auth; +pub mod bearer_token; +pub mod registry; + +pub use auth::{ + AuthError, ClientAuthenticator, ClientAuthenticatorHandle, ServerAuthenticator, + ServerAuthenticatorHandle, +}; +pub use bearer_token::{ + BearerToken, BearerTokenError, BearerTokenProvider, BearerTokenProviderHandle, Secret, +}; +pub(crate) use registry::ExtensionRegistryBuilder; +pub use registry::{ExtensionHandles, ExtensionRegistry}; diff --git a/rust/otap-dataflow/crates/engine/src/extensions/registry.rs b/rust/otap-dataflow/crates/engine/src/extensions/registry.rs new file mode 100644 index 0000000000..062242c2d1 --- /dev/null +++ b/rust/otap-dataflow/crates/engine/src/extensions/registry.rs @@ -0,0 +1,344 @@ +// Copyright The OpenTelemetry Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Type-safe extension registry. +//! +//! The registry stores extension service handles as type-erased, cloneable values. +//! Handles are looked up by `(extension_name, TypeId)` and returned via +//! `downcast_ref::()` followed by `.clone()`, so every consumer gets its own +//! independent copy. +//! +//! `ExtensionRegistry` is `Clone + Send` — each pipeline component receives +//! its own clone at startup, consistent with the thread-per-core, shared-nothing +//! architecture. No `Arc` is needed. +//! +//! Handle types must be `Clone + Send + 'static`. The registry intentionally +//! does **not** require `Sync` — it is a startup-time resource consumed once +//! at the top of each component's `start()` method. Handles that enter +//! tonic/gRPC services (which require `Sync`) are naturally `Sync` because +//! they wrap `Arc`, but that is a property of the handle, not a +//! registry requirement. + +use crate::error::Error; +use std::any::{Any, TypeId}; +use std::collections::HashMap; + +/// A type-erased, cloneable handle entry. +/// +/// Stores the handle as `Box` alongside a function pointer +/// that knows how to clone it back into another `Box`. +/// +/// ## Why `clone_fn`? +/// +/// `Clone` is not object-safe, so it cannot be used as a supertrait of `Any`. +/// This means `Box` does not implement `Clone`, even when the +/// concrete type inside it does. To work around this, we capture a function +/// pointer at construction time (in [`ErasedHandle::new`]) that is monomorphized +/// for the concrete type `T`. When cloning, the function pointer downcasts the +/// `&dyn Any` back to `&T`, calls `T::clone()`, and re-boxes the result. This +/// is a standard Rust pattern for type-erased cloning. +pub(crate) struct ErasedHandle { + /// The type-erased handle value. + value: Box, + /// A function pointer, monomorphized for `T` at construction time, that + /// downcasts `&dyn Any` → `&T`, clones it, and boxes the clone. + clone_fn: fn(&dyn Any) -> Box, +} + +impl ErasedHandle { + /// Creates a new erased handle, capturing the concrete clone function. + fn new(handle: T) -> Self { + Self { + value: Box::new(handle), + clone_fn: |any| { + let val = any + .downcast_ref::() + .expect("TypeId mismatch in ErasedHandle clone — this is a bug"); + Box::new(val.clone()) + }, + } + } +} + +impl Clone for ErasedHandle { + fn clone(&self) -> Self { + Self { + value: (self.clone_fn)(&*self.value), + clone_fn: self.clone_fn, + } + } +} + +/// A single extension's set of typed service handles. +/// +/// Created by the extension factory and later merged into the +/// [`ExtensionRegistryBuilder`]. +pub struct ExtensionHandles { + /// (TypeId → handle) for this extension. + handles: Vec<(TypeId, ErasedHandle)>, +} + +impl ExtensionHandles { + /// Creates a new, empty set of handles. + #[must_use] + pub fn new() -> Self { + Self { + handles: Vec::new(), + } + } + + /// Registers a typed handle. + /// + /// The same concrete type should only be registered once per extension. + /// Registering the same type twice will result in a duplicate entry; the + /// builder will keep the last one. + pub fn register(&mut self, handle: T) { + self.handles + .push((TypeId::of::(), ErasedHandle::new(handle))); + } + + /// Consumes self and returns the inner handle list. + #[must_use] + pub(crate) fn into_inner(self) -> Vec<(TypeId, ErasedHandle)> { + self.handles + } +} + +impl Default for ExtensionHandles { + fn default() -> Self { + Self::new() + } +} + +/// Builder for [`ExtensionRegistry`]. +/// +/// Collects handles from all extensions during pipeline build, then freezes +/// into an immutable registry that is cloned to each component. +pub(crate) struct ExtensionRegistryBuilder { + /// (extension_name, TypeId) → handle + entries: HashMap<(String, TypeId), ErasedHandle>, +} + +impl ExtensionRegistryBuilder { + /// Creates a new, empty builder. + #[must_use] + pub(crate) fn new() -> Self { + Self { + entries: HashMap::new(), + } + } + + /// Merges all handles produced by one extension into the builder. + /// + /// # Errors + /// + /// Returns [`Error::ExtensionHandleAlreadyRegistered`] if a `(name, TypeId)` pair + /// is already present. + pub(crate) fn merge( + &mut self, + extension_name: &str, + handles: ExtensionHandles, + ) -> Result<(), Error> { + for (type_id, handle) in handles.into_inner() { + let key = (extension_name.to_owned(), type_id); + if self.entries.contains_key(&key) { + return Err(Error::ExtensionHandleAlreadyRegistered { + extension: extension_name.to_owned(), + type_name: format!("{type_id:?}"), + }); + } + let _ = self.entries.insert(key, handle); + } + Ok(()) + } + + /// Freezes the builder into an immutable [`ExtensionRegistry`]. + #[must_use] + pub(crate) fn build(self) -> ExtensionRegistry { + ExtensionRegistry { + entries: self.entries, + } + } +} + +impl Default for ExtensionRegistryBuilder { + fn default() -> Self { + Self::new() + } +} + +/// An immutable, cloneable registry of extension service handles. +/// +/// Created once during pipeline build, then cloned to each component's +/// `start()` call. Handles are retrieved by extension name + concrete type. +/// +/// `Clone + Send` — no `Arc` needed. Each component owns its own copy, +/// consistent with the shared-nothing, thread-per-core model. `Sync` is +/// intentionally not required — the registry is consumed at startup before +/// any handles enter tonic services. +#[derive(Clone)] +pub struct ExtensionRegistry { + /// (extension_name, TypeId) → type-erased handle. + entries: HashMap<(String, TypeId), ErasedHandle>, +} + +impl ExtensionRegistry { + /// Creates an empty registry (useful for tests or pipelines with no extensions). + #[must_use] + pub fn empty() -> Self { + Self { + entries: HashMap::new(), + } + } + + /// Looks up a handle by extension name and concrete type, returning a clone. + /// + /// # Errors + /// + /// Returns [`Error::ExtensionHandleNotFound`] when no handle of the requested + /// type is registered for the given extension name. + pub fn get(&self, extension_name: &str) -> Result { + let key = (extension_name.to_owned(), TypeId::of::()); + let entry = self + .entries + .get(&key) + .ok_or_else(|| Error::ExtensionHandleNotFound { + extension: extension_name.to_owned(), + type_name: std::any::type_name::().to_owned(), + })?; + // Safety: we stored a `T` under `TypeId::of::()`, so downcast always succeeds. + let handle = entry + .value + .downcast_ref::() + .expect("TypeId mismatch in extension registry — this is a bug"); + Ok(handle.clone()) + } + + /// Returns `true` if the registry contains no handles. + #[must_use] + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + /// Returns the number of registered handles. + #[must_use] + pub fn len(&self) -> usize { + self.entries.len() + } +} + +impl std::fmt::Debug for ExtensionRegistry { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ExtensionRegistry") + .field("num_handles", &self.entries.len()) + .finish() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[derive(Clone, Debug, PartialEq)] + struct TestHandle { + value: String, + } + + #[derive(Clone, Debug, PartialEq)] + struct AnotherHandle { + count: u64, + } + + #[test] + fn register_and_retrieve_handle() { + let mut handles = ExtensionHandles::new(); + handles.register(TestHandle { + value: "hello".into(), + }); + + let mut builder = ExtensionRegistryBuilder::new(); + builder.merge("my_ext", handles).unwrap(); + let registry = builder.build(); + + let h: TestHandle = registry.get("my_ext").unwrap(); + assert_eq!(h.value, "hello"); + } + + #[test] + fn retrieve_wrong_type_fails() { + let mut handles = ExtensionHandles::new(); + handles.register(TestHandle { + value: "hello".into(), + }); + + let mut builder = ExtensionRegistryBuilder::new(); + builder.merge("my_ext", handles).unwrap(); + let registry = builder.build(); + + let result = registry.get::("my_ext"); + assert!(result.is_err()); + } + + #[test] + fn retrieve_wrong_name_fails() { + let mut handles = ExtensionHandles::new(); + handles.register(TestHandle { + value: "hello".into(), + }); + + let mut builder = ExtensionRegistryBuilder::new(); + builder.merge("my_ext", handles).unwrap(); + let registry = builder.build(); + + let result = registry.get::("other_ext"); + assert!(result.is_err()); + } + + #[test] + fn multiple_types_from_same_extension() { + let mut handles = ExtensionHandles::new(); + handles.register(TestHandle { + value: "token".into(), + }); + handles.register(AnotherHandle { count: 42 }); + + let mut builder = ExtensionRegistryBuilder::new(); + builder.merge("auth", handles).unwrap(); + let registry = builder.build(); + + assert_eq!( + registry.get::("auth").unwrap(), + TestHandle { + value: "token".into() + } + ); + assert_eq!( + registry.get::("auth").unwrap(), + AnotherHandle { count: 42 } + ); + } + + #[test] + fn duplicate_registration_fails() { + let mut handles_a = ExtensionHandles::new(); + handles_a.register(TestHandle { + value: "first".into(), + }); + let mut handles_b = ExtensionHandles::new(); + handles_b.register(TestHandle { + value: "second".into(), + }); + + let mut builder = ExtensionRegistryBuilder::new(); + builder.merge("ext", handles_a).unwrap(); + let err = builder.merge("ext", handles_b); + assert!(err.is_err()); + } + + #[test] + fn empty_registry() { + let registry = ExtensionRegistry::empty(); + assert!(registry.is_empty()); + assert_eq!(registry.len(), 0); + } +} diff --git a/rust/otap-dataflow/crates/engine/src/lib.rs b/rust/otap-dataflow/crates/engine/src/lib.rs index 6e2fab22ed..8f6a03003f 100644 --- a/rust/otap-dataflow/crates/engine/src/lib.rs +++ b/rust/otap-dataflow/crates/engine/src/lib.rs @@ -9,12 +9,14 @@ use crate::{ CHANNEL_MODE_LOCAL, CHANNEL_MODE_SHARED, CHANNEL_TYPE_MPMC, CHANNEL_TYPE_MPSC, ChannelMetricsRegistry, ChannelReceiverMetrics, ChannelSenderMetrics, }, - config::{ExporterConfig, ProcessorConfig, ReceiverConfig}, + config::{ExporterConfig, ExtensionConfig, ProcessorConfig, ReceiverConfig}, control::{AckMsg, CallData, NackMsg}, effect_handler::SourceTagging, entity_context::{NodeTelemetryGuard, NodeTelemetryHandle, with_node_telemetry_handle}, error::{Error, TypedError}, exporter::ExporterWrapper, + extension::ExtensionWrapper, + extensions::ExtensionRegistryBuilder, local::message::{LocalReceiver, LocalSender}, message::{Receiver, Sender}, node::{Node, NodeDefs, NodeId, NodeName, NodeType}, @@ -47,6 +49,8 @@ use std::{ pub mod error; pub mod exporter; +pub mod extension; +pub mod extensions; pub mod message; pub mod processor; pub mod receiver; @@ -195,6 +199,41 @@ impl NamedFactory for ExporterFactory { } } +/// A factory for creating extensions. +/// +/// Unlike the pipeline node factories, `ExtensionFactory` is **not** generic over +/// `PData` — extensions do not participate in the data-flow graph. +pub struct ExtensionFactory { + /// The name of the extension. + pub name: &'static str, + /// A function that creates a new extension instance. + pub create: fn( + pipeline_ctx: PipelineContext, + node: NodeId, + node_config: Arc, + extension_config: &ExtensionConfig, + ) -> Result, + /// Validates the node-specific config statically, without creating the component. + pub validate_config: fn(config: &serde_json::Value) -> Result<(), otap_df_config::error::Error>, +} + +// Note: We don't use `#[derive(Clone)]` here to keep consistency with other factories. +impl Clone for ExtensionFactory { + fn clone(&self) -> Self { + ExtensionFactory { + name: self.name, + create: self.create, + validate_config: self.validate_config, + } + } +} + +impl NamedFactory for ExtensionFactory { + fn name(&self) -> &'static str { + self.name + } +} + /// Returns a map of factory names to factory instances. pub fn get_factory_map( factory_map: &'static OnceLock>, @@ -325,9 +364,11 @@ pub struct PipelineFactory { receiver_factory_map: OnceLock>>, processor_factory_map: OnceLock>>, exporter_factory_map: OnceLock>>, + extension_factory_map: OnceLock>, receiver_factories: &'static [ReceiverFactory], processor_factories: &'static [ProcessorFactory], exporter_factories: &'static [ExporterFactory], + extension_factories: &'static [ExtensionFactory], } impl PipelineFactory { @@ -337,14 +378,17 @@ impl PipelineFactory { receiver_factories: &'static [ReceiverFactory], processor_factories: &'static [ProcessorFactory], exporter_factories: &'static [ExporterFactory], + extension_factories: &'static [ExtensionFactory], ) -> Self { Self { receiver_factory_map: OnceLock::new(), processor_factory_map: OnceLock::new(), exporter_factory_map: OnceLock::new(), + extension_factory_map: OnceLock::new(), receiver_factories, processor_factories, exporter_factories, + extension_factories, } } @@ -378,6 +422,16 @@ impl PipelineFactory { }) } + /// Gets the extension factory map, initializing it if necessary. + pub fn get_extension_factory_map(&self) -> &HashMap<&'static str, ExtensionFactory> { + self.extension_factory_map.get_or_init(|| { + self.extension_factories + .iter() + .map(|f| (f.name(), f.clone())) + .collect::>() + }) + } + /// Builds a runtime pipeline from the given pipeline configuration. /// /// Main phases: @@ -457,7 +511,9 @@ impl PipelineFactory { let mut receiver_count = 0usize; let mut processor_count = 0usize; let mut exporter_count = 0usize; + let mut extension_count = 0usize; let mut node_ids: HashMap = HashMap::new(); + let mut extension_node_names: Vec = Vec::new(); for (name, node_config) in config.node_iter() { let (node_type, pipe_node) = match node_config.kind() { @@ -476,6 +532,12 @@ impl PipelineFactory { exporter_count += 1; (NodeType::Exporter, pn) } + otap_df_config::node::NodeKind::Extension => { + let pn = PipeNode::new(extension_count); + extension_count += 1; + extension_node_names.push(name.clone()); + (NodeType::Extension, pn) + } otap_df_config::node::NodeKind::ProcessorChain => { return Err(Error::UnsupportedNodeKind { kind: "ProcessorChain".into(), @@ -494,11 +556,65 @@ impl PipelineFactory { ); pipeline_ctx.set_node_names(node_names); + // Extension creation pass: create extensions before other nodes so that handles + // are available in the registry when pipeline components start. + let mut extensions = Vec::new(); + let mut extension_registry_builder = ExtensionRegistryBuilder::new(); + for ext_name in &extension_node_names { + let node_config = config + .nodes() + .get(ext_name) + .expect("extension must exist in config") + .clone(); + let node_id = node_ids + .get(ext_name) + .expect("allocated in first pass") + .clone(); + let base_ctx = pipeline_ctx.with_node_context( + ext_name.clone(), + node_config.r#type.clone(), + otap_df_config::node::NodeKind::Extension, + node_config.identity_attributes(), + ); + let control_channel_capacity = channel_capacity_policy.control.node; + let extension_wrapper = self.build_node_wrapper( + &mut build_state, + &base_ctx, + NodeType::Extension, + node_id, + channel_metrics_enabled, + || { + self.create_extension( + &pipeline_ctx, + ext_name.clone(), + node_config.clone(), + control_channel_capacity, + ) + }, + )?; + extensions.push(extension_wrapper); + } + // Extract handles from each extension and merge into the registry builder. + // Use the config node name (from node_id) as the registry key, not the + // plugin URN — this allows multiple extensions of the same type with + // different configurations (e.g., "auth_team_a" and "auth_team_b" both + // using "extension:bearer_auth"). + for ext in &mut extensions { + if let Some(handles) = ext.take_handles() { + extension_registry_builder.merge(ext.node_id().name.as_ref(), handles)?; + } + } + let extension_registry = extension_registry_builder.build(); + // Second pass: create runtime nodes. Node IDs were pre-assigned above, // so we look them up from `node_ids` instead of calling `next_node_id`. + // Extensions were already created above and are skipped here. // ToDo(LQ): Collect all errors instead of failing fast to provide better feedback. for (name, node_config) in config.node_iter() { let node_kind = node_config.kind(); + if matches!(node_kind, otap_df_config::node::NodeKind::Extension) { + continue; + } let node_id = node_ids.get(name).expect("allocated in first pass").clone(); let base_ctx = pipeline_ctx.with_node_context( name.clone(), @@ -578,6 +694,10 @@ impl PipelineFactory { // ToDo(LQ): Implement processor chain optimization to eliminate intermediary channels. unreachable!("rejected in first pass"); } + otap_df_config::node::NodeKind::Extension => { + // Extensions are handled in the extension creation pass above. + unreachable!("skipped by continue above"); + } } } @@ -592,6 +712,8 @@ impl PipelineFactory { receivers, processors, exporters, + extensions, + extension_registry, nodes, telemetry_policy, ); @@ -670,6 +792,10 @@ impl PipelineFactory { kind: "ProcessorChain".into(), }); } + otap_df_config::node::NodeKind::Extension => { + // Extensions don't participate in data-flow wiring. + continue; + } }; _ = contracts_by_node.insert(node_name.as_ref().to_string().into(), contract); @@ -1381,6 +1507,63 @@ impl PipelineFactory { Ok(exporter) } + + /// Creates an extension instance from configuration. + fn create_extension( + &self, + pipeline_ctx: &PipelineContext, + name: NodeName, + node_config: Arc, + control_channel_capacity: usize, + ) -> Result { + let pipeline_group_id = pipeline_ctx.pipeline_group_id(); + let pipeline_id = pipeline_ctx.pipeline_id(); + let core_id = pipeline_ctx.core_id(); + + otel_debug!( + "extension.create.start", + pipeline_group_id = pipeline_group_id.as_ref(), + pipeline_id = pipeline_id.as_ref(), + core_id = core_id, + node_id = name.as_ref(), + ); + + // Validate plugin URN structure during registration + let normalized = otap_df_config::node_urn::validate_plugin_urn( + node_config.r#type.as_ref(), + otap_df_config::node::NodeKind::Extension, + ) + .map_err(|e| Error::ConfigError(Box::new(e)))?; + + let factory = self + .get_extension_factory_map() + .get(normalized.as_str()) + .ok_or(Error::UnknownExtension { + plugin_urn: normalized, + })?; + let extension_config = + ExtensionConfig::with_channel_capacity(name.clone(), control_channel_capacity); + let create = factory.create; + + let node_id = NodeId::build(usize::MAX, name.clone()); + let extension = create( + (*pipeline_ctx).clone(), + node_id, + node_config, + &extension_config, + ) + .map_err(|e| Error::ConfigError(Box::new(e)))?; + + otel_debug!( + "extension.create.complete", + pipeline_group_id = pipeline_group_id.as_ref(), + pipeline_id = pipeline_id.as_ref(), + core_id = core_id, + node_id = name.as_ref(), + ); + + Ok(extension) + } } trait TelemetryWrapped: Sized { @@ -1453,6 +1636,26 @@ impl TelemetryWrapped for ExporterWrapper { } } +impl TelemetryWrapped for ExtensionWrapper { + fn with_control_channel_metrics( + self, + pipeline_ctx: &PipelineContext, + channel_metrics: &mut ChannelMetricsRegistry, + channel_metrics_enabled: bool, + ) -> Self { + ExtensionWrapper::with_control_channel_metrics( + self, + pipeline_ctx, + channel_metrics, + channel_metrics_enabled, + ) + } + + fn with_node_telemetry_guard(self, guard: NodeTelemetryGuard) -> Self { + ExtensionWrapper::with_node_telemetry_guard(self, guard) + } +} + struct NodeRegistration { node_id: NodeId, node_type: NodeType, @@ -1496,6 +1699,7 @@ impl BuildState { NodeType::Receiver => Error::ReceiverAlreadyExists { receiver: node_id }, NodeType::Processor => Error::ProcessorAlreadyExists { processor: node_id }, NodeType::Exporter => Error::ExporterAlreadyExists { exporter: node_id }, + NodeType::Extension => Error::ExtensionAlreadyExists { extension: node_id }, }); } @@ -1529,7 +1733,9 @@ impl BuildState { let registration = self.registration(name)?; match registration.node_type { NodeType::Processor | NodeType::Exporter => Ok(registration.node_id.clone()), - NodeType::Receiver => Err(Error::UnknownNode { node: name.clone() }), + NodeType::Receiver | NodeType::Extension => { + Err(Error::UnknownNode { node: name.clone() }) + } } } } diff --git a/rust/otap-dataflow/crates/engine/src/local/exporter.rs b/rust/otap-dataflow/crates/engine/src/local/exporter.rs index f85a303435..167fed9e9e 100644 --- a/rust/otap-dataflow/crates/engine/src/local/exporter.rs +++ b/rust/otap-dataflow/crates/engine/src/local/exporter.rs @@ -36,6 +36,7 @@ use crate::control::{AckMsg, NackMsg}; use crate::effect_handler::{EffectHandlerCore, TelemetryTimerCancelHandle, TimerCancelHandle}; use crate::error::Error; +use crate::extensions::ExtensionRegistry; use crate::message::MessageChannel; use crate::node::NodeId; use crate::terminal_state::TerminalState; @@ -87,6 +88,7 @@ pub trait Exporter { self: Box, msg_chan: MessageChannel, effect_handler: EffectHandler, + extension_registry: ExtensionRegistry, ) -> Result; } diff --git a/rust/otap-dataflow/crates/engine/src/local/extension.rs b/rust/otap-dataflow/crates/engine/src/local/extension.rs new file mode 100644 index 0000000000..5bc9907971 --- /dev/null +++ b/rust/otap-dataflow/crates/engine/src/local/extension.rs @@ -0,0 +1,126 @@ +// Copyright The OpenTelemetry Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Trait and structures for local (!Send) extensions. +//! +//! Extensions are non-pipeline components that provide cross-cutting capabilities +//! such as authentication, health checking, or service discovery. Unlike receivers, +//! processors, and exporters, extensions do not participate in the pdata flow. +//! +//! # Lifecycle +//! +//! 1. The extension is instantiated and configured by its factory. +//! 2. The factory creates **service handles** — lightweight, cloneable values — +//! that are placed in the [`ExtensionRegistry`](crate::extensions::ExtensionRegistry). +//! 3. The `start` method is called, beginning the extension's background operation. +//! 4. Pipeline components retrieve handles from the registry (via their effect +//! handler) and use them to interact with the extension. +//! 5. The extension shuts down when it receives a `Shutdown` control message or +//! encounters a fatal error. +//! +//! # Thread Safety +//! +//! This implementation is designed for a single-threaded environment. +//! The `Extension` trait does not require the `Send` bound. +//! +//! # PData Independence +//! +//! Extension types are deliberately **not** generic over `PData`. Extensions sit +//! outside the data-flow graph and never touch pipeline data, so they use the +//! PData-free [`ExtensionControlMsg`] instead of `NodeControlMsg`. + +use crate::control::ExtensionControlMsg; +use crate::error::Error; +use crate::local::message::LocalReceiver; +use crate::node::NodeId; +use async_trait::async_trait; +use otap_df_channel::error::RecvError; +use otap_df_telemetry::reporter::MetricsReporter; + +/// A trait for local (!Send) extensions. +/// +/// Extensions run as independent tasks alongside pipeline components. They +/// receive only control messages (no pdata) and are started **before** the +/// rest of the pipeline, ensuring their service handles are ready for +/// consumers. +#[async_trait(?Send)] +pub trait Extension { + /// Starts the extension's background work. + /// + /// The pipeline engine calls this once at startup. The extension should + /// process control messages from `ctrl_chan` and perform its background + /// duties until a `Shutdown` message is received. + /// + /// # Parameters + /// + /// - `ctrl_chan`: Channel for receiving control messages (Shutdown, TimerTick, etc.). + /// - `effect_handler`: Provides node identity and logging. + /// + /// # Errors + /// + /// Returns an [`Error`] if an unrecoverable failure occurs. + async fn start( + self: Box, + ctrl_chan: ControlChannel, + effect_handler: EffectHandler, + ) -> Result<(), Error>; +} + +/// A channel for receiving control messages in a `!Send` extension. +pub struct ControlChannel { + ctrl_rx: LocalReceiver, +} + +impl ControlChannel { + /// Creates a new control channel. + #[must_use] + pub const fn new(ctrl_rx: LocalReceiver) -> Self { + Self { ctrl_rx } + } + + /// Receives the next control message, waiting if none is available. + /// + /// # Errors + /// + /// Returns [`RecvError`] if the channel is closed. + pub async fn recv(&mut self) -> Result { + self.ctrl_rx.recv().await + } +} + +/// A `!Send` effect handler for extensions. +/// +/// Provides a minimal set of capabilities — primarily node identity and logging. +/// Extensions that need periodic timers should use `tokio::time::interval` directly. +#[derive(Clone)] +pub struct EffectHandler { + node_id: NodeId, + #[allow(dead_code)] + metrics_reporter: MetricsReporter, +} + +impl EffectHandler { + /// Creates a new local extension effect handler. + #[must_use] + pub const fn new(node_id: NodeId, metrics_reporter: MetricsReporter) -> Self { + EffectHandler { + node_id, + metrics_reporter, + } + } + + /// Returns the id of the extension associated with this handler. + #[must_use] + pub fn extension_id(&self) -> NodeId { + self.node_id.clone() + } + + /// Print an info message to stdout. + pub async fn info(&self, message: &str) { + use tokio::io::{AsyncWriteExt, stdout}; + let mut out = stdout(); + let _ = out.write_all(message.as_bytes()).await; + let _ = out.write_all(b"\n").await; + let _ = out.flush().await; + } +} diff --git a/rust/otap-dataflow/crates/engine/src/local/mod.rs b/rust/otap-dataflow/crates/engine/src/local/mod.rs index b4bb3e477d..02b2059a14 100644 --- a/rust/otap-dataflow/crates/engine/src/local/mod.rs +++ b/rust/otap-dataflow/crates/engine/src/local/mod.rs @@ -4,6 +4,7 @@ //! Traits and structs defining the local (!Send) version of receivers, processors, and exporters. pub mod exporter; +pub mod extension; pub mod message; pub mod processor; pub mod receiver; diff --git a/rust/otap-dataflow/crates/engine/src/local/receiver.rs b/rust/otap-dataflow/crates/engine/src/local/receiver.rs index b86b3c7649..4e7ce77e82 100644 --- a/rust/otap-dataflow/crates/engine/src/local/receiver.rs +++ b/rust/otap-dataflow/crates/engine/src/local/receiver.rs @@ -37,6 +37,7 @@ use crate::effect_handler::{ EffectHandlerCore, SourceTagging, TelemetryTimerCancelHandle, TimerCancelHandle, }; use crate::error::{Error, TypedError}; +use crate::extensions::ExtensionRegistry; use crate::message::Sender; use crate::node::NodeId; use crate::terminal_state::TerminalState; @@ -80,6 +81,7 @@ pub trait Receiver { /// /// - `ctrl_chan`: A channel to receive control messages. /// - `effect_handler`: A handler to perform side effects such as opening a listener. + /// - `extension_registry`: A registry of extension service handles, consumed at startup. /// /// Each of these parameters is **NOT** [`Send`]. /// @@ -94,6 +96,7 @@ pub trait Receiver { self: Box, ctrl_chan: ControlChannel, effect_handler: EffectHandler, + extension_registry: ExtensionRegistry, ) -> Result; } @@ -125,7 +128,7 @@ impl ControlChannel { /// A `!Send` implementation of the EffectHandler. #[derive(Clone)] pub struct EffectHandler { - core: EffectHandlerCore, + pub(crate) core: EffectHandlerCore, /// A sender used to forward messages from the receiver. /// Supports multiple named output ports. diff --git a/rust/otap-dataflow/crates/engine/src/node.rs b/rust/otap-dataflow/crates/engine/src/node.rs index ff9f460b04..692d4c28a5 100644 --- a/rust/otap-dataflow/crates/engine/src/node.rs +++ b/rust/otap-dataflow/crates/engine/src/node.rs @@ -59,6 +59,8 @@ pub enum NodeType { Processor, /// Represents a node that exports data to an external destination. Exporter, + /// Represents a non-pipeline extension (e.g., auth provider, health check). + Extension, } /// Trait for nodes that can send pdata to a specific port. diff --git a/rust/otap-dataflow/crates/engine/src/pipeline_ctrl.rs b/rust/otap-dataflow/crates/engine/src/pipeline_ctrl.rs index 2ddc20cf1d..3765d8be35 100644 --- a/rust/otap-dataflow/crates/engine/src/pipeline_ctrl.rs +++ b/rust/otap-dataflow/crates/engine/src/pipeline_ctrl.rs @@ -12,8 +12,12 @@ //! are supported. use crate::context::PipelineContext; -use crate::control::{ControlSenders, NodeControlMsg, PipelineControlMsg, PipelineCtrlMsgReceiver}; +use crate::control::{ + ControlSenders, ExtensionControlMsg, NodeControlMsg, PipelineControlMsg, + PipelineCtrlMsgReceiver, +}; use crate::error::Error; +use crate::message::Sender; use crate::pipeline_metrics::PipelineMetricsMonitor; use otap_df_config::DeployedPipelineKey; use otap_df_config::policy::TelemetryPolicy; @@ -183,6 +187,8 @@ pub struct PipelineCtrlMsgManager { pipeline_ctrl_msg_receiver: PipelineCtrlMsgReceiver, /// Allows sending control messages back to nodes. control_senders: ControlSenders, + /// Control message senders for extensions (shutdown after pipeline draining). + extension_senders: Vec>, /// Repeating timers for generic TimerTick. tick_timers: TimerSet, /// Repeating timers for telemetry collection (CollectTelemetry). @@ -208,6 +214,7 @@ impl PipelineCtrlMsgManager { pipeline_context: PipelineContext, pipeline_ctrl_msg_receiver: PipelineCtrlMsgReceiver, control_senders: ControlSenders, + extension_senders: Vec>, event_reporter: ObservedEventReporter, metrics_reporter: MetricsReporter, telemetry_policy: TelemetryPolicy, @@ -218,6 +225,7 @@ impl PipelineCtrlMsgManager { pipeline_context, pipeline_ctrl_msg_receiver, control_senders, + extension_senders, tick_timers: TimerSet::new(), telemetry_timers: TimerSet::new(), delayed_data: BinaryHeap::new(), @@ -461,6 +469,17 @@ impl PipelineCtrlMsgManager { } } } + + // Shut down extensions after the data-plane has drained. + for sender in &self.extension_senders { + let _ = sender + .send(ExtensionControlMsg::Shutdown { + deadline: Instant::now(), + reason: "Pipeline draining complete".to_owned(), + }) + .await; + } + Ok(()) } @@ -597,6 +616,7 @@ mod tests { pipeline_context, pipeline_rx, control_senders, + Vec::new(), observed_state_store.reporter(SendPolicy::default()), metrics_reporter, TelemetryPolicy::default(), @@ -1041,6 +1061,7 @@ mod tests { pipeline_context, pipeline_rx, ControlSenders::new(), + Vec::new(), observed_state_store.reporter(SendPolicy::default()), metrics_reporter, TelemetryPolicy::default(), diff --git a/rust/otap-dataflow/crates/engine/src/processor.rs b/rust/otap-dataflow/crates/engine/src/processor.rs index 4ddf30c161..44b68418e5 100644 --- a/rust/otap-dataflow/crates/engine/src/processor.rs +++ b/rust/otap-dataflow/crates/engine/src/processor.rs @@ -249,7 +249,7 @@ impl ProcessorWrapper { .. } => { let (control_sender, control_receiver) = - wrap_control_channel_metrics::( + wrap_control_channel_metrics::>( &node_id, pipeline_ctx, channel_metrics, @@ -286,7 +286,7 @@ impl ProcessorWrapper { .. } => { let (control_sender, control_receiver) = - wrap_control_channel_metrics::( + wrap_control_channel_metrics::>( &node_id, pipeline_ctx, channel_metrics, diff --git a/rust/otap-dataflow/crates/engine/src/receiver.rs b/rust/otap-dataflow/crates/engine/src/receiver.rs index e36a74a533..3636e3738b 100644 --- a/rust/otap-dataflow/crates/engine/src/receiver.rs +++ b/rust/otap-dataflow/crates/engine/src/receiver.rs @@ -15,6 +15,7 @@ use crate::control::{Controllable, NodeControlMsg, PipelineCtrlMsgSender}; use crate::effect_handler::SourceTagging; use crate::entity_context::NodeTelemetryGuard; use crate::error::{Error, ReceiverErrorKind}; +use crate::extensions::ExtensionRegistry; use crate::local::message::{LocalReceiver, LocalSender}; use crate::local::receiver as local; use crate::message::{Receiver, Sender}; @@ -233,7 +234,7 @@ impl ReceiverWrapper { .. } => { let (control_sender, control_receiver) = - wrap_control_channel_metrics::( + wrap_control_channel_metrics::>( &node_id, pipeline_ctx, channel_metrics, @@ -270,7 +271,7 @@ impl ReceiverWrapper { .. } => { let (control_sender, control_receiver) = - wrap_control_channel_metrics::( + wrap_control_channel_metrics::>( &node_id, pipeline_ctx, channel_metrics, @@ -301,6 +302,7 @@ impl ReceiverWrapper { self, pipeline_ctrl_msg_tx: PipelineCtrlMsgSender, metrics_reporter: MetricsReporter, + extension_registry: ExtensionRegistry, ) -> Result { match (self, metrics_reporter) { ( @@ -335,7 +337,9 @@ impl ReceiverWrapper { metrics_reporter, ); effect_handler.set_source_tagging(source_tag); - receiver.start(ctrl_msg_chan, effect_handler).await + receiver + .start(ctrl_msg_chan, effect_handler, extension_registry) + .await } ( ReceiverWrapper::Shared { @@ -369,7 +373,9 @@ impl ReceiverWrapper { metrics_reporter, ); effect_handler.set_source_tagging(source_tag); - receiver.start(ctrl_msg_chan, effect_handler).await + receiver + .start(ctrl_msg_chan, effect_handler, extension_registry) + .await } } } @@ -464,6 +470,7 @@ impl NodeWithPDataSender for ReceiverWrapper { #[cfg(test)] mod tests { use super::ReceiverWrapper; + use crate::extensions::ExtensionRegistry; use crate::local::receiver as local; use crate::receiver::Error; use crate::shared::receiver as shared; @@ -510,6 +517,7 @@ mod tests { self: Box, mut ctrl_msg_recv: local::ControlChannel, effect_handler: local::EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { // Bind to an ephemeral port. let addr: SocketAddr = "127.0.0.1:0".parse().unwrap(); @@ -589,6 +597,7 @@ mod tests { self: Box, mut ctrl_msg_recv: shared::ControlChannel, effect_handler: shared::EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { // Bind to an ephemeral port. let addr: SocketAddr = "127.0.0.1:0".parse().unwrap(); @@ -772,3 +781,139 @@ mod tests { .run_validation(validation_procedure()); } } + +/// Tests that a receiver can retrieve an auth handle from the extension +/// registry (passed to `start()`) and use it to validate requests. +#[cfg(test)] +mod auth_extension_tests { + use super::ReceiverWrapper; + use crate::error::Error; + use crate::extensions::ExtensionRegistry; + use crate::extensions::auth::{AuthError, ServerAuthenticator, ServerAuthenticatorHandle}; + use crate::extensions::{ExtensionHandles, ExtensionRegistryBuilder}; + use crate::local::receiver as local; + use crate::terminal_state::TerminalState; + use crate::testing::receiver::TestRuntime; + use crate::testing::{TestMsg, test_node}; + use async_trait::async_trait; + use otap_df_config::node::NodeUserConfig; + use std::future::Future; + use std::pin::Pin; + use std::sync::Arc; + use std::time::{Duration, Instant}; + use tokio::time::timeout; + + /// A server authenticator that checks `x-api-key` against an allow list. + struct ApiKeyAuth { + allowed_keys: Vec, + } + + impl ServerAuthenticator for ApiKeyAuth { + fn authenticate(&self, headers: &http::HeaderMap) -> Result<(), AuthError> { + let key = headers + .get("x-api-key") + .and_then(|v| v.to_str().ok()) + .ok_or_else(|| AuthError { + message: "missing x-api-key header".into(), + })?; + if !self.allowed_keys.iter().any(|k| k == key) { + return Err(AuthError { + message: format!("invalid API key: {key}"), + }); + } + Ok(()) + } + } + + /// A minimal receiver that retrieves an auth handle from the extension + /// registry, validates a set of headers, and forwards authenticated + /// payloads as messages. No TCP — auth is exercised directly in `start()`. + struct AuthReceiver; + + #[async_trait(?Send)] + impl local::Receiver for AuthReceiver { + async fn start( + self: Box, + mut ctrl: local::ControlChannel, + effect_handler: local::EffectHandler, + extension_registry: ExtensionRegistry, + ) -> Result { + // Retrieve the auth handle — this is what a real receiver does. + let auth = extension_registry + .get::("api_key_auth") + .expect("api_key_auth extension must be registered"); + + // Valid key → forward the payload. + let mut headers = http::HeaderMap::new(); + let _ = headers.insert("x-api-key", "valid-key".parse().unwrap()); + assert!(auth.authenticate(&headers).is_ok()); + effect_handler + .send_message(TestMsg("authenticated".into())) + .await + .unwrap(); + + // Invalid key → rejected, nothing forwarded. + let _ = headers.insert("x-api-key", "bad-key".parse().unwrap()); + assert!(auth.authenticate(&headers).is_err()); + + // Wait for shutdown. + loop { + let msg = ctrl.recv().await?; + if msg.is_shutdown() { + break; + } + } + Ok(TerminalState::default()) + } + } + + fn build_auth_registry() -> ExtensionRegistry { + let auth = ApiKeyAuth { + allowed_keys: vec!["valid-key".into()], + }; + let mut handles = ExtensionHandles::new(); + handles.register(ServerAuthenticatorHandle::new(auth)); + let mut builder = ExtensionRegistryBuilder::new(); + builder.merge("api_key_auth", handles).unwrap(); + builder.build() + } + + /// Receiver retrieves auth handle from the extension registry and + /// uses it to accept/reject requests. + #[test] + fn test_receiver_with_auth_extension() { + let test_runtime = TestRuntime::new(); + let user_config = Arc::new(NodeUserConfig::new_receiver_config("auth_receiver")); + let receiver = ReceiverWrapper::local( + AuthReceiver, + test_node("auth_recv"), + user_config, + test_runtime.config(), + ); + + let registry = build_auth_registry(); + + test_runtime + .set_receiver(receiver) + .with_extension_registry(registry) + .run_test(|ctx| -> Pin>> { + Box::pin(async move { + ctx.send_shutdown(Instant::now() + Duration::from_millis(200), "Test") + .await + .expect("shutdown failed"); + }) + }) + .run_validation(|mut ctx| -> Pin>> { + Box::pin(async move { + let msg = timeout(Duration::from_secs(3), ctx.recv()) + .await + .expect("timed out") + .expect("no message"); + assert!( + matches!(&msg, TestMsg(payload) if payload == "authenticated"), + "expected authenticated payload, got {msg:?}" + ); + }) + }); + } +} diff --git a/rust/otap-dataflow/crates/engine/src/runtime_pipeline.rs b/rust/otap-dataflow/crates/engine/src/runtime_pipeline.rs index 373fb39298..a26b4c9fdb 100644 --- a/rust/otap-dataflow/crates/engine/src/runtime_pipeline.rs +++ b/rust/otap-dataflow/crates/engine/src/runtime_pipeline.rs @@ -10,6 +10,8 @@ use crate::control::{ }; use crate::entity_context::{NodeTaskContext, instrument_with_node_context}; use crate::error::{Error, TypedError}; +use crate::extension::ExtensionWrapper; +use crate::extensions::ExtensionRegistry; use crate::node::{Node, NodeDefs, NodeId, NodeType, NodeWithPDataReceiver, NodeWithPDataSender}; use crate::pipeline_ctrl::PipelineCtrlMsgManager; use crate::terminal_state::TerminalState; @@ -36,6 +38,10 @@ pub struct RuntimePipeline { processors: Vec>, /// A map node id to exporter runtime node. exporters: Vec>, + /// Extension instances (not part of the data-flow graph). + extensions: Vec, + /// Immutable registry of typed handles produced by extensions. + extension_registry: ExtensionRegistry, /// A precomputed map of all node IDs to their Node trait objects (? @@@) for efficient access /// Indexed by NodeIndex @@ -73,6 +79,8 @@ impl RuntimePipeline { receivers: Vec>, processors: Vec>, exporters: Vec>, + extensions: Vec, + extension_registry: ExtensionRegistry, nodes: NodeDefs, telemetry_policy: TelemetryPolicy, ) -> Self { @@ -81,6 +89,8 @@ impl RuntimePipeline { receivers, processors, exporters, + extensions, + extension_registry, nodes, channel_metrics: Default::default(), telemetry_policy, @@ -121,6 +131,8 @@ impl RuntimePipeline { receivers, processors, exporters, + extensions, + extension_registry, nodes: _nodes, channel_metrics, telemetry_policy, @@ -136,6 +148,36 @@ impl RuntimePipeline { let mut futures = FuturesUnordered::new(); let mut control_senders = ControlSenders::default(); + // Spawn extension tasks first so services are available before pipeline nodes start. + // Clone each extension's control sender so the PipelineCtrlMsgManager can + // deliver shutdown (and future control) messages after pipeline draining. + let mut extension_senders = Vec::with_capacity(extensions.len()); + for extension in extensions { + let mut extension = extension; + extension_senders.push(extension.control_sender()); + let metrics_reporter = metrics_reporter.clone(); + let telemetry_guard = extension.take_telemetry_guard(); + let node_entity_key = telemetry_guard.as_ref().map(|t| t.entity_key()); + let telemetry_handle = telemetry_guard.as_ref().map(|t| t.handle()); + let fut = async move { + let result = extension.start(metrics_reporter).await; + drop(telemetry_guard); + result + }; + if let Some(handle) = telemetry_handle { + let input_key = handle.input_channel_key(); + let output_keys = handle.output_channel_keys(); + let node_ctx = + NodeTaskContext::new(node_entity_key, Some(handle), input_key, output_keys); + futures.push(local_tasks.spawn_local(instrument_with_node_context(node_ctx, fut))); + } else if let Some(key) = node_entity_key { + let node_ctx = NodeTaskContext::new(Some(key), None, None, Vec::new()); + futures.push(local_tasks.spawn_local(instrument_with_node_context(node_ctx, fut))); + } else { + futures.push(local_tasks.spawn_local(fut)); + } + } + // Spawn node tasks and register their control senders, scoping telemetry where available. for exporter in exporters { let mut exporter = exporter; @@ -151,9 +193,14 @@ impl RuntimePipeline { let pipeline_ctrl_msg_tx = pipeline_ctrl_msg_tx.clone(); let effect_metrics_reporter = metrics_reporter.clone(); let final_metrics_reporter = metrics_reporter.clone(); + let extension_registry = extension_registry.clone(); let fut = async move { let result = exporter - .start(pipeline_ctrl_msg_tx, effect_metrics_reporter) + .start( + pipeline_ctrl_msg_tx, + effect_metrics_reporter, + extension_registry, + ) .await .map(|terminal_state| { report_terminal_metrics(&final_metrics_reporter, terminal_state); @@ -221,9 +268,14 @@ impl RuntimePipeline { let pipeline_ctrl_msg_tx = pipeline_ctrl_msg_tx.clone(); let effect_metrics_reporter = metrics_reporter.clone(); let final_metrics_reporter = metrics_reporter.clone(); + let extension_registry = extension_registry.clone(); let fut = async move { let result = receiver - .start(pipeline_ctrl_msg_tx, effect_metrics_reporter) + .start( + pipeline_ctrl_msg_tx, + effect_metrics_reporter, + extension_registry, + ) .await .map(|terminal_state| { report_terminal_metrics(&final_metrics_reporter, terminal_state); @@ -252,6 +304,7 @@ impl RuntimePipeline { pipeline_context, pipeline_ctrl_msg_rx, control_senders, + extension_senders, event_reporter, metrics_reporter, telemetry_policy, @@ -311,6 +364,8 @@ impl RuntimePipeline { .exporters .get(ndef.inner.index) .map(|e| e as &dyn Node), + // Extensions are not tracked in NodeDefs. + NodeType::Extension => None, } } @@ -332,6 +387,8 @@ impl RuntimePipeline { .get_mut(ndef.inner.index) .map(|p| p as &mut dyn NodeWithPDataSender), NodeType::Exporter => None, + // Extensions are not tracked in NodeDefs. + NodeType::Extension => None, } } @@ -353,6 +410,8 @@ impl RuntimePipeline { .exporters .get_mut(ndef.inner.index) .map(|e| e as &mut dyn NodeWithPDataReceiver), + // Extensions are not tracked in NodeDefs. + NodeType::Extension => None, } } @@ -385,6 +444,12 @@ impl RuntimePipeline { .send_control_msg(ctrl_msg) .await } + // Extensions are not part of the node graph. + NodeType::Extension => { + return Err(TypedError::Error(Error::InternalError { + message: format!("extension node {node_id:?} cannot receive control messages via node graph"), + })); + } } .map_err(|e| TypedError::NodeControlMsgSendError { node_id: node_id.index, diff --git a/rust/otap-dataflow/crates/engine/src/shared/exporter.rs b/rust/otap-dataflow/crates/engine/src/shared/exporter.rs index db9474f60d..8debc16e6a 100644 --- a/rust/otap-dataflow/crates/engine/src/shared/exporter.rs +++ b/rust/otap-dataflow/crates/engine/src/shared/exporter.rs @@ -35,6 +35,7 @@ use crate::control::{AckMsg, NackMsg, NodeControlMsg}; use crate::effect_handler::{EffectHandlerCore, TelemetryTimerCancelHandle, TimerCancelHandle}; use crate::error::Error; +use crate::extensions::ExtensionRegistry; use crate::message::Message; use crate::node::NodeId; use crate::shared::message::SharedReceiver; @@ -57,6 +58,7 @@ pub trait Exporter { self: Box, msg_chan: MessageChannel, effect_handler: EffectHandler, + extension_registry: ExtensionRegistry, ) -> Result; } diff --git a/rust/otap-dataflow/crates/engine/src/shared/extension.rs b/rust/otap-dataflow/crates/engine/src/shared/extension.rs new file mode 100644 index 0000000000..320010c98b --- /dev/null +++ b/rust/otap-dataflow/crates/engine/src/shared/extension.rs @@ -0,0 +1,94 @@ +// Copyright The OpenTelemetry Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Trait and structures for shared (Send) extensions. +//! +//! This is the `Send`-bound counterpart of [`super::super::local::extension`]. +//! Extensions using this module can be moved across thread boundaries, which +//! may be useful for integrations that require `tokio::spawn` (as opposed to +//! `spawn_local`). +//! +//! See the local extension module documentation for full lifecycle details. +//! +//! # PData Independence +//! +//! Like the local variant, shared extension types are **not** generic over +//! `PData`. They use the PData-free [`ExtensionControlMsg`]. + +use crate::control::ExtensionControlMsg; +use crate::error::Error; +use crate::node::NodeId; +use crate::shared::message::SharedReceiver; +use async_trait::async_trait; +use otap_df_channel::error::RecvError; +use otap_df_telemetry::reporter::MetricsReporter; + +/// A trait for shared (Send) extensions. +#[async_trait] +pub trait Extension: Send { + /// Starts the extension's background work (Send-compatible). + async fn start( + self: Box, + ctrl_chan: ControlChannel, + effect_handler: EffectHandler, + ) -> Result<(), Error>; +} + +/// A channel for receiving control messages in a `Send` extension. +pub struct ControlChannel { + rx: SharedReceiver, +} + +impl ControlChannel { + /// Creates a new control channel. + #[must_use] + pub const fn new(rx: SharedReceiver) -> Self { + Self { rx } + } + + /// Receives the next control message, waiting if none is available. + /// + /// # Errors + /// + /// Returns [`RecvError`] if the channel is closed. + pub async fn recv(&mut self) -> Result { + self.rx.recv().await + } +} + +/// A `Send` effect handler for extensions. +/// +/// Provides a minimal set of capabilities — primarily node identity and logging. +/// Extensions that need periodic timers should use `tokio::time::interval` directly. +#[derive(Clone)] +pub struct EffectHandler { + node_id: NodeId, + #[allow(dead_code)] + metrics_reporter: MetricsReporter, +} + +impl EffectHandler { + /// Creates a new shared extension effect handler. + #[must_use] + pub const fn new(node_id: NodeId, metrics_reporter: MetricsReporter) -> Self { + EffectHandler { + node_id, + metrics_reporter, + } + } + + /// Returns the id of the extension associated with this handler. + #[must_use] + pub fn extension_id(&self) -> NodeId { + self.node_id.clone() + } + + /// Print an info message to stdout. + pub async fn info(&self, message: &str) { + use tokio::io::{AsyncWriteExt, stdout}; + let mut out = stdout(); + let _ = out.write_all(message.as_bytes()).await; + let _ = out.write_all(b"\n").await; + let _ = out.flush().await; + } +} diff --git a/rust/otap-dataflow/crates/engine/src/shared/mod.rs b/rust/otap-dataflow/crates/engine/src/shared/mod.rs index 3f6f246f76..bd55d8cb26 100644 --- a/rust/otap-dataflow/crates/engine/src/shared/mod.rs +++ b/rust/otap-dataflow/crates/engine/src/shared/mod.rs @@ -4,6 +4,7 @@ //! Traits and structs defining the shared (Send) version of receivers, processors, and exporters. pub mod exporter; +pub mod extension; pub mod message; pub mod processor; pub mod receiver; diff --git a/rust/otap-dataflow/crates/engine/src/shared/receiver.rs b/rust/otap-dataflow/crates/engine/src/shared/receiver.rs index 967c9b0a6a..62ccbabca5 100644 --- a/rust/otap-dataflow/crates/engine/src/shared/receiver.rs +++ b/rust/otap-dataflow/crates/engine/src/shared/receiver.rs @@ -37,6 +37,7 @@ use crate::effect_handler::{ EffectHandlerCore, SourceTagging, TelemetryTimerCancelHandle, TimerCancelHandle, }; use crate::error::{Error, TypedError}; +use crate::extensions::ExtensionRegistry; use crate::node::NodeId; use crate::shared::message::{SharedReceiver, SharedSender}; use crate::terminal_state::TerminalState; @@ -62,6 +63,7 @@ pub trait Receiver { self: Box, ctrl_chan: ControlChannel, effect_handler: EffectHandler, + extension_registry: ExtensionRegistry, ) -> Result; } @@ -93,7 +95,7 @@ impl ControlChannel { /// A `Send` implementation of the EffectHandlerTrait. #[derive(Clone)] pub struct EffectHandler { - core: EffectHandlerCore, + pub(crate) core: EffectHandlerCore, /// A sender used to forward messages from the receiver. /// Supports multiple named output ports. diff --git a/rust/otap-dataflow/crates/engine/src/testing/exporter.rs b/rust/otap-dataflow/crates/engine/src/testing/exporter.rs index 8a8c508cc8..e9be9a8377 100644 --- a/rust/otap-dataflow/crates/engine/src/testing/exporter.rs +++ b/rust/otap-dataflow/crates/engine/src/testing/exporter.rs @@ -10,10 +10,12 @@ use crate::ExporterFactory; use crate::config::ExporterConfig; use crate::context::{ControllerContext, PipelineContext}; use crate::control::{ - Controllable, NodeControlMsg, PipelineCtrlMsgReceiver, pipeline_ctrl_msg_channel, + Controllable, NodeControlMsg, PipelineCtrlMsgReceiver, PipelineCtrlMsgSender, + pipeline_ctrl_msg_channel, }; use crate::error::Error; use crate::exporter::ExporterWrapper; +use crate::extensions::{ExtensionRegistry, ExtensionRegistryBuilder}; use crate::local::message::{LocalReceiver, LocalSender}; use crate::message::{Receiver, Sender}; use crate::node::NodeWithPDataReceiver; @@ -170,10 +172,15 @@ pub struct TestPhase { control_sender: Sender>, pdata_sender: Sender, - /// Join handle for the starting the exporter task - run_exporter_handle: tokio::task::JoinHandle>, + /// The exporter to start (deferred until `run_test`). + exporter: ExporterWrapper, + pipeline_ctrl_msg_sender: PipelineCtrlMsgSender, pipeline_ctrl_msg_receiver: PipelineCtrlMsgReceiver, + + metrics_system: InternalTelemetrySystem, + + extension_registry: Option, } /// Data and operations for the validation phase of an exporter. @@ -229,6 +236,10 @@ impl TestRuntime { } /// Sets the exporter for the test runtime and returns the test phase. + /// + /// The exporter is not started until [`TestPhase::run_test`] is called, which + /// allows injecting a custom [`ExtensionRegistry`] via + /// [`TestPhase::with_extension_registry`] before the exporter begins. pub fn set_exporter(self, mut exporter: ExporterWrapper) -> TestPhase { let control_sender = exporter.control_sender(); let (pdata_tx, pdata_rx) = match &exporter { @@ -254,28 +265,18 @@ impl TestRuntime { exporter .set_pdata_receiver(test_node(self.config.name.clone()), pdata_rx) .expect("Failed to set PData receiver"); - let metrics_reporter_start = self.metrics_reporter(); - let metrics_reporter_terminal = self.metrics_reporter(); - let metrics_collector = self.metrics_system.collector(); - let run_exporter_handle = self.local_tasks.spawn_local(async move { - exporter - .start(pipeline_ctrl_msg_tx, metrics_reporter_start) - .await - .map(|terminal_state| { - for snapshot in terminal_state.into_metrics() { - let _ = metrics_reporter_terminal.try_report_snapshot(snapshot); - } - metrics_collector.collect_pending(); // Collect after sending all the - }) - }); + TestPhase { rt: self.rt, local_tasks: self.local_tasks, counters: self.counter.clone(), control_sender, pdata_sender: pdata_tx, - run_exporter_handle, + exporter, + pipeline_ctrl_msg_sender: pipeline_ctrl_msg_tx, pipeline_ctrl_msg_receiver: pipeline_ctrl_msg_rx, + metrics_system: self.metrics_system, + extension_registry: None, } } } @@ -287,14 +288,49 @@ impl Default for TestRuntime { } impl TestPhase { + /// Sets a custom extension registry to be passed to the exporter's `start()` method. + /// + /// When not set, the exporter will receive an empty registry. + pub fn with_extension_registry(mut self, registry: ExtensionRegistry) -> Self { + self.extension_registry = Some(registry); + self + } + /// Starts the test scenario by executing the provided function with the test context. - pub fn run_test(self, f: F) -> ValidationPhase + pub fn run_test(mut self, f: F) -> ValidationPhase where F: FnOnce(TestContext) -> Fut + 'static, Fut: Future + 'static, { + let metrics_reporter_start = self.metrics_system.reporter(); + let metrics_reporter_terminal = self.metrics_system.reporter(); + let metrics_collector = self.metrics_system.collector(); + + let extension_registry = self + .extension_registry + .take() + .unwrap_or_else(|| ExtensionRegistryBuilder::new().build()); + let mut context = self.create_context(); let ctx_test = context.clone(); + + let pipeline_ctrl_msg_tx = self.pipeline_ctrl_msg_sender; + let exporter = self.exporter; + let run_exporter_handle = self.local_tasks.spawn_local(async move { + exporter + .start( + pipeline_ctrl_msg_tx, + metrics_reporter_start, + extension_registry, + ) + .await + .map(|terminal_state| { + for snapshot in terminal_state.into_metrics() { + let _ = metrics_reporter_terminal.try_report_snapshot(snapshot); + } + metrics_collector.collect_pending(); + }) + }); _ = self.local_tasks.spawn_local(f(ctx_test)); context.pipeline_ctrl_msg_receiver = Some(self.pipeline_ctrl_msg_receiver); @@ -303,7 +339,7 @@ impl TestPhase { rt: self.rt, local_tasks: self.local_tasks, context, - run_exporter_handle: self.run_exporter_handle, + run_exporter_handle, } } diff --git a/rust/otap-dataflow/crates/engine/src/testing/receiver.rs b/rust/otap-dataflow/crates/engine/src/testing/receiver.rs index d99db1b0d4..23b1f47c34 100644 --- a/rust/otap-dataflow/crates/engine/src/testing/receiver.rs +++ b/rust/otap-dataflow/crates/engine/src/testing/receiver.rs @@ -11,6 +11,7 @@ use crate::control::{ Controllable, NodeControlMsg, PipelineCtrlMsgReceiver, pipeline_ctrl_msg_channel, }; use crate::error::Error; +use crate::extensions::{ExtensionRegistry, ExtensionRegistryBuilder}; use crate::local::message::{LocalReceiver, LocalSender}; use crate::message::{Receiver, Sender}; use crate::node::NodeWithPDataSender; @@ -158,6 +159,7 @@ pub struct TestPhase { control_sender: Sender>, receiver: ReceiverWrapper, counters: CtrlMsgCounters, + extension_registry: Option, } /// Data and operations for the validation phase of a receiver. @@ -221,6 +223,7 @@ impl TestRuntime { receiver, control_sender, counters: self.counter, + extension_registry: None, } } } @@ -232,6 +235,14 @@ impl Default for TestRuntime { } impl TestPhase { + /// Sets a custom extension registry to be passed to the receiver's `start()` method. + /// + /// When not set, the receiver will receive an empty registry. + pub fn with_extension_registry(mut self, registry: ExtensionRegistry) -> Self { + self.extension_registry = Some(registry); + self + } + /// Starts the test scenario by executing the provided function with the test context. pub fn run_test(mut self, f: F) -> ValidationPhase where @@ -278,10 +289,14 @@ impl TestPhase { let (_metrics_rx, metrics_reporter) = MetricsReporter::create_new_and_receiver(1); let final_metrics_reporter = metrics_reporter.clone(); + let extension_registry = self + .extension_registry + .unwrap_or_else(|| ExtensionRegistryBuilder::new().build()); + let run_receiver_handle = self.local_tasks.spawn_local(async move { let terminal_state = self .receiver - .start(pipeline_ctrl_msg_tx, metrics_reporter) + .start(pipeline_ctrl_msg_tx, metrics_reporter, extension_registry) .await .expect("Receiver event loop failed"); diff --git a/rust/otap-dataflow/crates/otap/src/console_exporter.rs b/rust/otap-dataflow/crates/otap/src/console_exporter.rs index 921e6fdb60..df9579feaf 100644 --- a/rust/otap-dataflow/crates/otap/src/console_exporter.rs +++ b/rust/otap-dataflow/crates/otap/src/console_exporter.rs @@ -15,6 +15,7 @@ use otap_df_engine::context::PipelineContext; use otap_df_engine::control::{AckMsg, NodeControlMsg}; use otap_df_engine::error::Error; use otap_df_engine::exporter::ExporterWrapper; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::exporter::{EffectHandler, Exporter}; use otap_df_engine::message::{Message, MessageChannel}; use otap_df_engine::node::NodeId; @@ -101,6 +102,7 @@ impl Exporter for ConsoleExporter { self: Box, mut msg_chan: MessageChannel, effect_handler: EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { loop { match msg_chan.recv().await? { diff --git a/rust/otap-dataflow/crates/otap/src/error_exporter.rs b/rust/otap-dataflow/crates/otap/src/error_exporter.rs index 5fe16cc46a..f7215d14bf 100644 --- a/rust/otap-dataflow/crates/otap/src/error_exporter.rs +++ b/rust/otap-dataflow/crates/otap/src/error_exporter.rs @@ -11,6 +11,7 @@ use otap_df_engine::context::PipelineContext; use otap_df_engine::control::{NackMsg, NodeControlMsg}; use otap_df_engine::error::Error; use otap_df_engine::exporter::ExporterWrapper; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::exporter::{EffectHandler, Exporter}; use otap_df_engine::message::{Message, MessageChannel}; use otap_df_engine::node::NodeId; @@ -78,6 +79,7 @@ impl Exporter for ErrorExporter { self: Box, mut msg_chan: MessageChannel, effect_handler: EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { loop { match msg_chan.recv().await? { diff --git a/rust/otap-dataflow/crates/otap/src/fake_data_generator.rs b/rust/otap-dataflow/crates/otap/src/fake_data_generator.rs index f46e6fde3b..5f54c774a9 100644 --- a/rust/otap-dataflow/crates/otap/src/fake_data_generator.rs +++ b/rust/otap-dataflow/crates/otap/src/fake_data_generator.rs @@ -17,6 +17,7 @@ use otap_df_engine::MessageSourceLocalEffectHandlerExtension; use otap_df_engine::config::ReceiverConfig; use otap_df_engine::context::PipelineContext; use otap_df_engine::error::{Error, ReceiverErrorKind, format_error_sources}; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::receiver as local; use otap_df_engine::node::NodeId; use otap_df_engine::receiver::ReceiverWrapper; @@ -228,6 +229,7 @@ impl local::Receiver for FakeGeneratorReceiver { mut self: Box, mut ctrl_msg_recv: local::ControlChannel, effect_handler: local::EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { //start event loop let traffic_config = self.config.get_traffic_config(); diff --git a/rust/otap-dataflow/crates/otap/src/internal_telemetry_receiver.rs b/rust/otap-dataflow/crates/otap/src/internal_telemetry_receiver.rs index 24d7096f33..bd84dc99d7 100644 --- a/rust/otap-dataflow/crates/otap/src/internal_telemetry_receiver.rs +++ b/rust/otap-dataflow/crates/otap/src/internal_telemetry_receiver.rs @@ -17,6 +17,7 @@ use otap_df_engine::config::ReceiverConfig; use otap_df_engine::context::PipelineContext; use otap_df_engine::control::NodeControlMsg; use otap_df_engine::error::Error; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::receiver as local; use otap_df_engine::node::NodeId; use otap_df_engine::receiver::ReceiverWrapper; @@ -106,6 +107,7 @@ impl local::Receiver for InternalTelemetryReceiver { mut self: Box, mut ctrl_msg_recv: local::ControlChannel, effect_handler: local::EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { let internal = self.internal_telemetry.clone(); let logs_receiver = internal.logs_receiver; diff --git a/rust/otap-dataflow/crates/otap/src/noop_exporter.rs b/rust/otap-dataflow/crates/otap/src/noop_exporter.rs index a0c0e8d85b..cce2b38100 100644 --- a/rust/otap-dataflow/crates/otap/src/noop_exporter.rs +++ b/rust/otap-dataflow/crates/otap/src/noop_exporter.rs @@ -11,6 +11,7 @@ use otap_df_engine::context::PipelineContext; use otap_df_engine::control::{AckMsg, NodeControlMsg}; use otap_df_engine::error::Error; use otap_df_engine::exporter::ExporterWrapper; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::exporter::{EffectHandler, Exporter}; use otap_df_engine::message::{Message, MessageChannel}; use otap_df_engine::node::NodeId; @@ -50,6 +51,7 @@ impl Exporter for NoopExporter { self: Box, mut msg_chan: MessageChannel, effect_handler: EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { loop { match msg_chan.recv().await? { diff --git a/rust/otap-dataflow/crates/otap/src/otap_exporter.rs b/rust/otap-dataflow/crates/otap/src/otap_exporter.rs index 3b730f043d..0db4d2c0a0 100644 --- a/rust/otap-dataflow/crates/otap/src/otap_exporter.rs +++ b/rust/otap-dataflow/crates/otap/src/otap_exporter.rs @@ -21,6 +21,7 @@ use otap_df_engine::context::PipelineContext; use otap_df_engine::control::NodeControlMsg; use otap_df_engine::error::{Error, ExporterErrorKind, format_error_sources}; use otap_df_engine::exporter::ExporterWrapper; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::exporter as local; use otap_df_engine::message::{Message, MessageChannel}; use otap_df_engine::node::NodeId; @@ -110,6 +111,7 @@ impl local::Exporter for OTAPExporter { mut self: Box, mut msg_chan: MessageChannel, effect_handler: local::EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { otel_info!( "exporter.start", @@ -475,6 +477,7 @@ mod tests { use otap_df_engine::control::pipeline_ctrl_msg_channel; use otap_df_engine::error::Error; use otap_df_engine::exporter::ExporterWrapper; + use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::message::LocalReceiver; use otap_df_engine::local::message::LocalSender; use otap_df_engine::message::Receiver; @@ -845,7 +848,10 @@ mod tests { pipeline_ctrl_msg_tx: PipelineCtrlMsgSender, metrics_reporter: MetricsReporter, ) -> Result<(), Error> { - _ = exporter.start(pipeline_ctrl_msg_tx, metrics_reporter).await; + let extension_registry = ExtensionRegistry::empty(); + _ = exporter + .start(pipeline_ctrl_msg_tx, metrics_reporter, extension_registry) + .await; Ok(()) } diff --git a/rust/otap-dataflow/crates/otap/src/otap_receiver.rs b/rust/otap-dataflow/crates/otap/src/otap_receiver.rs index 7a57708aaa..1a8be92ec8 100644 --- a/rust/otap-dataflow/crates/otap/src/otap_receiver.rs +++ b/rust/otap-dataflow/crates/otap/src/otap_receiver.rs @@ -31,6 +31,7 @@ use otap_df_engine::config::ReceiverConfig; use otap_df_engine::context::PipelineContext; use otap_df_engine::control::{AckMsg, NackMsg, NodeControlMsg}; use otap_df_engine::error::{Error, ReceiverErrorKind, format_error_sources}; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::node::NodeId; use otap_df_engine::receiver::ReceiverWrapper; use otap_df_engine::shared::receiver as shared; @@ -235,6 +236,7 @@ impl shared::Receiver for OTAPReceiver { mut self: Box, mut ctrl_msg_recv: shared::ControlChannel, effect_handler: shared::EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { otap_df_telemetry::otel_info!( "receiver.start", diff --git a/rust/otap-dataflow/crates/otap/src/otlp_exporter.rs b/rust/otap-dataflow/crates/otap/src/otlp_exporter.rs index 8ba8293bde..41385ded04 100644 --- a/rust/otap-dataflow/crates/otap/src/otlp_exporter.rs +++ b/rust/otap-dataflow/crates/otap/src/otlp_exporter.rs @@ -28,6 +28,7 @@ use otap_df_engine::context::PipelineContext; use otap_df_engine::control::{AckMsg, NackMsg, NodeControlMsg}; use otap_df_engine::error::{Error, ExporterErrorKind, format_error_sources}; use otap_df_engine::exporter::ExporterWrapper; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::exporter::{EffectHandler, Exporter}; use otap_df_engine::message::{Message, MessageChannel}; use otap_df_engine::node::NodeId; @@ -119,6 +120,7 @@ impl Exporter for OTLPExporter { mut self: Box, mut msg_chan: MessageChannel, effect_handler: EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { otel_info!( "otlp.exporter.grpc.start", @@ -781,6 +783,7 @@ mod tests { #[cfg(not(windows))] use { otap_df_engine::control::{Controllable, PipelineCtrlMsgSender, pipeline_ctrl_msg_channel}, + otap_df_engine::extensions::ExtensionRegistry, otap_df_engine::local::message::{LocalReceiver, LocalSender}, otap_df_engine::message::{Receiver, Sender}, otap_df_engine::node::NodeWithPDataReceiver, @@ -1092,8 +1095,9 @@ mod tests { pipeline_ctrl_msg_tx: PipelineCtrlMsgSender, metrics_reporter: MetricsReporter, ) -> Result<(), Error> { + let extension_registry = ExtensionRegistry::empty(); exporter - .start(pipeline_ctrl_msg_tx, metrics_reporter) + .start(pipeline_ctrl_msg_tx, metrics_reporter, extension_registry) .await .map(|_| ()) } diff --git a/rust/otap-dataflow/crates/otap/src/otlp_http_exporter/mod.rs b/rust/otap-dataflow/crates/otap/src/otlp_http_exporter/mod.rs index ed5200b94b..d57ef669a9 100644 --- a/rust/otap-dataflow/crates/otap/src/otlp_http_exporter/mod.rs +++ b/rust/otap-dataflow/crates/otap/src/otlp_http_exporter/mod.rs @@ -31,6 +31,7 @@ use otap_df_engine::context::PipelineContext; use otap_df_engine::control::{AckMsg, NackMsg, NodeControlMsg}; use otap_df_engine::error::{Error as EngineError, ExporterErrorKind}; use otap_df_engine::exporter::ExporterWrapper; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::exporter::{EffectHandler, Exporter}; use otap_df_engine::message::{Message, MessageChannel}; use otap_df_engine::node::NodeId; @@ -204,6 +205,7 @@ impl Exporter for OtlpHttpExporter { mut self: Box, mut msg_chan: MessageChannel, effect_handler: EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { let logs_endpoint = Rc::new( self.config diff --git a/rust/otap-dataflow/crates/otap/src/otlp_receiver.rs b/rust/otap-dataflow/crates/otap/src/otlp_receiver.rs index 125af4576d..94bf138d7d 100644 --- a/rust/otap-dataflow/crates/otap/src/otlp_receiver.rs +++ b/rust/otap-dataflow/crates/otap/src/otlp_receiver.rs @@ -38,6 +38,7 @@ use otap_df_engine::config::ReceiverConfig; use otap_df_engine::context::PipelineContext; use otap_df_engine::control::{AckMsg, NackMsg, NodeControlMsg}; use otap_df_engine::error::{Error, ReceiverErrorKind, format_error_sources}; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::node::NodeId; use otap_df_engine::receiver::ReceiverWrapper; use otap_df_engine::shared::receiver as shared; @@ -484,6 +485,7 @@ impl shared::Receiver for OTLPReceiver { mut self: Box, mut ctrl_msg_recv: shared::ControlChannel, effect_handler: shared::EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { let grpc_enabled = self.config.protocols.grpc.is_some(); let both_enabled = self.config.protocols.has_both(); diff --git a/rust/otap-dataflow/crates/otap/src/parquet_exporter.rs b/rust/otap-dataflow/crates/otap/src/parquet_exporter.rs index affa1d2ef5..ce40e3d2d2 100644 --- a/rust/otap-dataflow/crates/otap/src/parquet_exporter.rs +++ b/rust/otap-dataflow/crates/otap/src/parquet_exporter.rs @@ -42,6 +42,7 @@ use otap_df_engine::context::PipelineContext; use otap_df_engine::control::NodeControlMsg; use otap_df_engine::error::{Error, ExporterErrorKind, format_error_sources}; use otap_df_engine::exporter::ExporterWrapper; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::exporter::{EffectHandler, Exporter}; use otap_df_engine::message::{Message, MessageChannel}; use otap_df_engine::node::NodeId; @@ -153,6 +154,7 @@ impl Exporter for ParquetExporter { mut self: Box, mut msg_chan: MessageChannel, effect_handler: EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { let exporter_id = effect_handler.exporter_id(); let object_store = @@ -432,6 +434,7 @@ mod test { pipeline_ctrl_msg_channel, }; use otap_df_engine::exporter::ExporterWrapper; + use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::message::{LocalReceiver, LocalSender}; use otap_df_engine::message::{Receiver, Sender}; use otap_df_engine::node::NodeWithPDataReceiver; @@ -919,8 +922,9 @@ mod test { ) -> Result<(), Error> { let (_metrics_rx, metrics_reporter) = otap_df_telemetry::reporter::MetricsReporter::create_new_and_receiver(1); + let extension_registry = ExtensionRegistry::empty(); exporter - .start(pipeline_ctrl_msg_tx, metrics_reporter) + .start(pipeline_ctrl_msg_tx, metrics_reporter, extension_registry) .await .map(|_| ()) } @@ -1069,8 +1073,9 @@ mod test { ) -> Result<(), Error> { let (_metrics_rx, metrics_reporter) = otap_df_telemetry::reporter::MetricsReporter::create_new_and_receiver(1); + let extension_registry = ExtensionRegistry::empty(); exporter - .start(pipeline_ctrl_msg_tx, metrics_reporter) + .start(pipeline_ctrl_msg_tx, metrics_reporter, extension_registry) .await .map(|_| ()) } @@ -1217,8 +1222,9 @@ mod test { ) -> Result<(), Error> { let (_metrics_rx, metrics_reporter) = otap_df_telemetry::reporter::MetricsReporter::create_new_and_receiver(1); + let extension_registry = ExtensionRegistry::empty(); exporter - .start(pipeline_ctrl_msg_tx, metrics_reporter) + .start(pipeline_ctrl_msg_tx, metrics_reporter, extension_registry) .await .map(|_| ()) } @@ -1460,8 +1466,9 @@ mod test { pipeline_ctrl_msg_tx: PipelineCtrlMsgSender, metrics_reporter: otap_df_telemetry::reporter::MetricsReporter, ) -> Result<(), Error> { + let extension_registry = ExtensionRegistry::empty(); exporter - .start(pipeline_ctrl_msg_tx, metrics_reporter) + .start(pipeline_ctrl_msg_tx, metrics_reporter, extension_registry) .await .map(|_| ()) } diff --git a/rust/otap-dataflow/crates/otap/src/perf_exporter/exporter.rs b/rust/otap-dataflow/crates/otap/src/perf_exporter/exporter.rs index a29d51b3de..3ca26c1b01 100644 --- a/rust/otap-dataflow/crates/otap/src/perf_exporter/exporter.rs +++ b/rust/otap-dataflow/crates/otap/src/perf_exporter/exporter.rs @@ -32,6 +32,7 @@ use otap_df_engine::context::PipelineContext; use otap_df_engine::control::{AckMsg, NodeControlMsg}; use otap_df_engine::error::{Error, ExporterErrorKind}; use otap_df_engine::exporter::ExporterWrapper; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::exporter as local; use otap_df_engine::message::{Message, MessageChannel}; use otap_df_engine::node::NodeId; @@ -128,6 +129,7 @@ impl local::Exporter for PerfExporter { mut self: Box, mut msg_chan: MessageChannel, effect_handler: local::EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { // init variables for tracking // let mut average_pipeline_latency: f64 = 0.0; diff --git a/rust/otap-dataflow/crates/otap/src/syslog_cef_receiver.rs b/rust/otap-dataflow/crates/otap/src/syslog_cef_receiver.rs index f0912d3f14..0e93c43406 100644 --- a/rust/otap-dataflow/crates/otap/src/syslog_cef_receiver.rs +++ b/rust/otap-dataflow/crates/otap/src/syslog_cef_receiver.rs @@ -10,6 +10,7 @@ use otap_df_config::node::NodeUserConfig; use otap_df_engine::config::ReceiverConfig; use otap_df_engine::context::PipelineContext; use otap_df_engine::control::NodeControlMsg; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::node::NodeId; use otap_df_engine::receiver::ReceiverWrapper; use otap_df_engine::terminal_state::TerminalState; @@ -222,6 +223,7 @@ impl local::Receiver for SyslogCefReceiver { self: Box, mut ctrl_chan: local::ControlChannel, effect_handler: local::EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { // Start periodic telemetry collection (1s), similar to other nodes let timer_cancel_handle = effect_handler @@ -1619,7 +1621,9 @@ mod telemetry_tests { // Start receiver let handle = tokio::task::spawn_local(async move { - let _ = Box::new(receiver).start(ctrl_chan, eh).await; + let _ = Box::new(receiver) + .start(ctrl_chan, eh, ExtensionRegistry::empty()) + .await; }); // Send one valid and one invalid UDP datagram @@ -1711,7 +1715,9 @@ mod telemetry_tests { // Start receiver let handle = tokio::task::spawn_local(async move { - let _ = Box::new(receiver).start(ctrl_chan, eh).await; + let _ = Box::new(receiver) + .start(ctrl_chan, eh, ExtensionRegistry::empty()) + .await; }); // Allow bind tokio::time::sleep(Duration::from_millis(50)).await; diff --git a/rust/otap-dataflow/crates/otap/src/topic_exporter.rs b/rust/otap-dataflow/crates/otap/src/topic_exporter.rs index f0e597534f..f9b71f93a4 100644 --- a/rust/otap-dataflow/crates/otap/src/topic_exporter.rs +++ b/rust/otap-dataflow/crates/otap/src/topic_exporter.rs @@ -18,6 +18,7 @@ use otap_df_engine::context::PipelineContext; use otap_df_engine::control::{AckMsg, NodeControlMsg}; use otap_df_engine::error::Error; use otap_df_engine::exporter::ExporterWrapper; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::exporter::{EffectHandler, Exporter}; use otap_df_engine::message::{Message, MessageChannel}; use otap_df_engine::node::NodeId; @@ -84,6 +85,7 @@ impl Exporter for TopicExporter { self: Box, mut msg_chan: MessageChannel, effect_handler: EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { loop { match msg_chan.recv().await? { diff --git a/rust/otap-dataflow/crates/otap/src/topic_receiver.rs b/rust/otap-dataflow/crates/otap/src/topic_receiver.rs index 8948d14b6f..e3091d3181 100644 --- a/rust/otap-dataflow/crates/otap/src/topic_receiver.rs +++ b/rust/otap-dataflow/crates/otap/src/topic_receiver.rs @@ -18,6 +18,7 @@ use otap_df_engine::config::ReceiverConfig; use otap_df_engine::context::PipelineContext; use otap_df_engine::control::NodeControlMsg; use otap_df_engine::error::Error; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::receiver as local; use otap_df_engine::node::NodeId; use otap_df_engine::receiver::ReceiverWrapper; @@ -101,6 +102,7 @@ impl local::Receiver for TopicReceiver { self: Box, mut ctrl_msg_recv: local::ControlChannel, _effect_handler: local::EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { loop { match ctrl_msg_recv.recv().await { diff --git a/rust/otap-dataflow/crates/otap/tests/common/counting_exporter.rs b/rust/otap-dataflow/crates/otap/tests/common/counting_exporter.rs index 8357049b49..77749474ae 100644 --- a/rust/otap-dataflow/crates/otap/tests/common/counting_exporter.rs +++ b/rust/otap-dataflow/crates/otap/tests/common/counting_exporter.rs @@ -17,6 +17,7 @@ use otap_df_engine::context::PipelineContext; use otap_df_engine::control::{AckMsg, NodeControlMsg}; use otap_df_engine::error::Error; use otap_df_engine::exporter::ExporterWrapper; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::exporter::{EffectHandler, Exporter}; use otap_df_engine::message::{Message, MessageChannel}; use otap_df_engine::node::NodeId; @@ -90,6 +91,7 @@ impl Exporter for CountingExporter { self: Box, mut msg_chan: MessageChannel, effect_handler: EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { loop { match msg_chan.recv().await? { diff --git a/rust/otap-dataflow/crates/otap/tests/common/flaky_exporter.rs b/rust/otap-dataflow/crates/otap/tests/common/flaky_exporter.rs index df575a8b01..87b9bd8b86 100644 --- a/rust/otap-dataflow/crates/otap/tests/common/flaky_exporter.rs +++ b/rust/otap-dataflow/crates/otap/tests/common/flaky_exporter.rs @@ -22,6 +22,7 @@ use otap_df_engine::context::PipelineContext; use otap_df_engine::control::{AckMsg, NackMsg, NodeControlMsg}; use otap_df_engine::error::Error; use otap_df_engine::exporter::ExporterWrapper; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::exporter::{EffectHandler, Exporter}; use otap_df_engine::message::{Message, MessageChannel}; use otap_df_engine::node::NodeId; @@ -136,6 +137,7 @@ impl Exporter for FlakyExporter { self: Box, mut msg_chan: MessageChannel, effect_handler: EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { loop { match msg_chan.recv().await? { diff --git a/rust/otap-dataflow/crates/validation/src/validation_exporter.rs b/rust/otap-dataflow/crates/validation/src/validation_exporter.rs index fc358c6db6..8c2d091642 100644 --- a/rust/otap-dataflow/crates/validation/src/validation_exporter.rs +++ b/rust/otap-dataflow/crates/validation/src/validation_exporter.rs @@ -16,6 +16,7 @@ use otap_df_engine::context::PipelineContext; use otap_df_engine::control::NodeControlMsg; use otap_df_engine::error::Error as EngineError; use otap_df_engine::exporter::ExporterWrapper; +use otap_df_engine::extensions::ExtensionRegistry; use otap_df_engine::local::exporter::{EffectHandler, Exporter}; use otap_df_engine::message::{Message, MessageChannel}; use otap_df_engine::node::NodeId; @@ -158,6 +159,7 @@ impl Exporter for ValidationExporter { mut self: Box, mut msg_chan: MessageChannel, effect_handler: EffectHandler, + _extension_registry: ExtensionRegistry, ) -> Result { let _ = effect_handler .start_periodic_telemetry(Duration::from_secs(1)) diff --git a/rust/otap-dataflow/src/main.rs b/rust/otap-dataflow/src/main.rs index 95f107da9e..add8841dcc 100644 --- a/rust/otap-dataflow/src/main.rs +++ b/rust/otap-dataflow/src/main.rs @@ -195,6 +195,10 @@ fn validate_pipeline_components( .get_exporter_factory_map() .get(urn_str) .map(|f| f.validate_config), + NodeKind::Extension => OTAP_PIPELINE_FACTORY + .get_extension_factory_map() + .get(urn_str) + .map(|f| f.validate_config), }; match validate_config_fn { @@ -203,6 +207,7 @@ fn validate_pipeline_components( NodeKind::Receiver => "receiver", NodeKind::Processor | NodeKind::ProcessorChain => "processor", NodeKind::Exporter => "exporter", + NodeKind::Extension => "extension", }; return Err(std::io::Error::other(format!( "Unknown {} component `{}` in pipeline_group={} pipeline={} node={}",