otel-arrow-dfe-engine 0.61.0

// Copyright The OpenTelemetry Authors
// SPDX-License-Identifier: Apache-2.0

//! Extension wrapper, control channel, and bundle types.
//!
//! Extensions are PData-free -- they never process PData, only control
//! messages. This module defines [`EffectHandler`], [`ExtensionWrapper`], and
//! [`ExtensionBundle`]. The typestate builder that constructs bundles lives
//! in [`super::builder`].
//!
//! For the local (!Send) and shared (Send) Extension traits, see
//! [`local::extension`](crate::local::extension) and
//! [`shared::extension`](crate::shared::extension).

use crate::capability::factory::{LocalInstanceFactory, SharedInstanceFactory};
use crate::channel_metrics::ChannelMetricsRegistry;
use crate::channel_mode::{LocalMode, SharedMode, wrap_extension_control_channel_metrics};
use crate::config::ExtensionConfig;
use crate::context::ExtensionContext;
use crate::control::{ExtensionControlMsg, ShutdownPayload};
use crate::entity_context::{EntityTelemetryGuard, EntityTelemetryHandle};
use crate::error::Error;
use crate::local::extension as local_ext;
use crate::local::message::LocalReceiver;
use crate::message::Sender;
use crate::shared::extension as shared_ext;
use crate::shared::message::SharedReceiver;
use crate::terminal_state::TerminalState;
use otel_arrow_dfe_channel::error::RecvError;
use otel_arrow_dfe_config::ExtensionId;
use otel_arrow_dfe_config::extension::ExtensionUserConfig;
use otel_arrow_dfe_telemetry::otel_debug;
use otel_arrow_dfe_telemetry::reporter::MetricsReporter;
use std::future::Future;
use std::sync::Arc;

// -- ControlChannel ----------------------------------------------------------

/// A shutdown-aware channel for receiving extension control messages.
///
/// Generic over the receiver type `R` to support both local (!Send) and
/// shared (Send) channels.
///
/// `Shutdown` is delivered immediately via a dedicated priority
/// oneshot channel (`shutdown_rx`) that is checked ahead of the regular
/// control channel via a biased select. After `Shutdown` is returned
/// the regular channel is closed and subsequent [`recv`](Self::recv)
/// calls yield [`RecvError::Closed`] so the event loop exits cleanly.
#[doc(hidden)]
pub struct ControlChannel<R> {
    control_rx: Option<R>,
    shutdown_rx: Option<tokio::sync::oneshot::Receiver<ShutdownPayload>>,
}

/// Trait abstracting over local and shared receivers for control messages.
#[doc(hidden)]
pub trait ControlReceiver {
    /// Receives the next message, or returns an error if the channel is closed.
    fn recv(&mut self) -> impl Future<Output = Result<ExtensionControlMsg, RecvError>>;
}

impl ControlReceiver for LocalReceiver<ExtensionControlMsg> {
    async fn recv(&mut self) -> Result<ExtensionControlMsg, RecvError> {
        LocalReceiver::recv(self).await
    }
}

impl ControlReceiver for SharedReceiver<ExtensionControlMsg> {
    async fn recv(&mut self) -> Result<ExtensionControlMsg, RecvError> {
        SharedReceiver::recv(self).await
    }
}

impl<R: ControlReceiver + Unpin> ControlChannel<R> {
    /// Creates a new `ControlChannel`. The `shutdown_rx` is a dedicated
    /// priority oneshot path; if it fires it preempts any pending
    /// messages on `control_rx`.
    #[must_use]
    pub fn new(
        control_rx: R,
        shutdown_rx: tokio::sync::oneshot::Receiver<ShutdownPayload>,
    ) -> Self {
        ControlChannel {
            control_rx: Some(control_rx),
            shutdown_rx: Some(shutdown_rx),
        }
    }

    /// Asynchronously receives the next control message. The priority
    /// shutdown channel is checked ahead of the regular control channel
    /// (biased select). After `Shutdown` is delivered the regular
    /// channel is closed and subsequent calls return [`RecvError::Closed`].
    ///
    /// # Errors
    ///
    /// Returns a [`RecvError`] if the channel is closed.
    pub async fn recv(&mut self) -> Result<ExtensionControlMsg, RecvError> {
        loop {
            let control_rx = self.control_rx.as_mut().ok_or(RecvError::Closed)?;
            match self.shutdown_rx.as_mut() {
                Some(shutdown_rx) => {
                    tokio::select! {
                        biased;
                        res = shutdown_rx => {
                            // Receiver is consumed regardless of outcome.
                            self.shutdown_rx = None;
                            match res {
                                Ok(payload) => {
                                    // Close the regular channel; post-Shutdown semantics.
                                    let _ = self.control_rx.take();
                                    return Ok(ExtensionControlMsg::Shutdown {
                                        deadline: payload.deadline,
                                        reason: payload.reason,
                                    });
                                }
                                Err(_) => {
                                    // Sender was dropped without sending; no
                                    // shutdown will ever arrive. Fall through
                                    // to control-only path on the next loop.
                                    continue;
                                }
                            }
                        }
                        msg = control_rx.recv() => return msg,
                    }
                }
                None => return control_rx.recv().await,
            }
        }
    }
}

// -- EffectHandler -----------------------------------------------------------

/// The effect handler for extensions.
///
/// Provides extensions with identity, basic I/O, and engine-supplied
/// signalling capabilities. Extensions manage their own timers directly
/// via `tokio::time` rather than through the engine's timer
/// infrastructure, keeping the extension system fully PData-free.
#[derive(Clone)]
pub struct EffectHandler {
    name: ExtensionId,
    #[allow(dead_code)]
    metrics_reporter: MetricsReporter,
    readiness: Option<super::readiness::ReadinessSignaller>,
}

impl EffectHandler {
    /// Creates a new `EffectHandler` for the given extension.
    #[must_use]
    pub(crate) fn new(
        name: ExtensionId,
        metrics_reporter: MetricsReporter,
        readiness: Option<super::readiness::ReadinessSignaller>,
    ) -> Self {
        EffectHandler {
            name,
            metrics_reporter,
            readiness,
        }
    }

    /// Returns the name of the extension associated with this handler.
    #[must_use]
    pub fn extension_id(&self) -> ExtensionId {
        self.name.clone()
    }

    /// Signal that this variant is ready. No-op when not opted in;
    /// idempotent across clones when opted in.
    pub fn signal_ready(&self) {
        if let Some(sig) = &self.readiness {
            sig.ready();
        }
    }
}

// -- Variant kind -------------------------------------------------------------

/// Identifies a physical extension variant within an [`ExtensionBundle`].
/// A single `ExtensionId` may register both `Local` and `Shared` variants.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ExtensionVariant {
    /// `!Send` variant scheduled onto the host `LocalSet`.
    Local,
    /// `Send` variant runnable on multi-threaded executors.
    /// Whether an extension is shared between pipelines or not
    /// is a separate concern that will be handled by the scope
    /// that the extension is declared in. Currently only pipeline
    /// scope is supported.
    Shared,
}

impl ExtensionVariant {
    /// Attribute string used in telemetry (`extension.variant`).
    #[must_use]
    pub const fn as_str(self) -> &'static str {
        match self {
            ExtensionVariant::Local => "local",
            ExtensionVariant::Shared => "shared",
        }
    }
}

// -- Lifecycle ----------------------------------------------------------------

/// The lifecycle state of an extension variant (local or shared).
///
/// Generic over the extension trait object type `E`, which is
/// `Rc<dyn local::Extension>` for local variants and
/// `Box<dyn shared::Extension>` for shared variants.
pub enum ExtensionLifecycle<E, R> {
    /// Active extension with an event loop and control channel.
    Active {
        /// The extension trait object with start method.
        extension: E,
        /// Control channel sender.
        control_sender: Sender<ExtensionControlMsg>,
        /// Control channel receiver.
        control_receiver: R,
        /// Owned sender for the dedicated priority shutdown oneshot.
        /// `Some` until [`ExtensionWrapper::take_shutdown_sender`] takes
        /// it out at lifecycle-spawn time.
        shutdown_sender: Option<tokio::sync::oneshot::Sender<ShutdownPayload>>,
        /// Owned receiver for the dedicated priority shutdown oneshot.
        /// `Some` until [`ExtensionWrapper::start`] takes it out and
        /// hands it to the extension's [`ControlChannel`].
        shutdown_receiver: Option<tokio::sync::oneshot::Receiver<ShutdownPayload>>,
        /// Engine-side readiness probe; `None` when not opted in.
        readiness_probe: Option<super::readiness::ReadinessProbe>,
        /// Send-side handle paired with `readiness_probe`.
        readiness_signaller: Option<super::readiness::ReadinessSignaller>,
    },
    /// Passive extension -- capabilities only, no task spawned.
    Passive,
}

// -- ExtensionWrapper --------------------------------------------------------

/// Wrapper for a single extension variant in the engine.
///
/// Extensions are NOT generic over PData -- they operate exclusively on
/// [`ExtensionControlMsg`], keeping the extension system entirely decoupled
/// from the data-plane type.
///
/// The first layer is local vs shared (enum variant). The second layer is
/// active vs passive ([`ExtensionLifecycle`]). This two-level structure makes
/// impossible states unrepresentable while keeping local and shared
/// concerns cleanly separated.
pub enum ExtensionWrapper {
    /// A local (!Send) extension variant.
    Local {
        /// Extension identifier.
        name: ExtensionId,
        /// User-provided extension configuration.
        user_config: Arc<ExtensionUserConfig>,
        /// Engine runtime configuration for this extension.
        runtime_config: ExtensionConfig,
        /// Telemetry guard for this extension entity, set during wiring.
        telemetry: Option<EntityTelemetryGuard>,
        /// The lifecycle state (active with channels, or passive).
        lifecycle: ExtensionLifecycle<
            std::rc::Rc<dyn local_ext::Extension>,
            LocalReceiver<ExtensionControlMsg>,
        >,
        /// Mints instances of the extension's concrete type for
        /// capability consumers. The engine's generated capability
        /// registration glue downcasts back to `E` and wraps in the
        /// requested trait object.
        instance_factory: LocalInstanceFactory,
    },
    /// A shared (Send) extension variant.
    Shared {
        /// Extension identifier.
        name: ExtensionId,
        /// User-provided extension configuration.
        user_config: Arc<ExtensionUserConfig>,
        /// Engine runtime configuration for this extension.
        runtime_config: ExtensionConfig,
        /// Telemetry guard for this extension entity, set during wiring.
        telemetry: Option<EntityTelemetryGuard>,
        /// The lifecycle state (active with channels, or passive).
        lifecycle:
            ExtensionLifecycle<Box<dyn shared_ext::Extension>, SharedReceiver<ExtensionControlMsg>>,
        /// Mints instances of the extension's concrete type for
        /// capability consumers.
        instance_factory: SharedInstanceFactory,
    },
}

/// The set of extension wrappers produced by a single extension factory.
///
/// An extension can have at most one local variant and one shared variant.
/// The builder enforces that at least one is present and that dual
/// registrations use different concrete types.
pub struct ExtensionBundle {
    /// The local (!Send) extension variant, if any.
    local: Option<ExtensionWrapper>,
    /// The shared (Send) extension variant, if any.
    shared: Option<ExtensionWrapper>,
}

impl ExtensionWrapper {
    /// Returns the name of this extension.
    #[must_use]
    pub fn name(&self) -> ExtensionId {
        match self {
            ExtensionWrapper::Local { name, .. } | ExtensionWrapper::Shared { name, .. } => {
                name.clone()
            }
        }
    }

    /// Returns the user configuration for this extension.
    #[must_use]
    pub fn user_config(&self) -> Arc<ExtensionUserConfig> {
        match self {
            ExtensionWrapper::Local { user_config, .. }
            | ExtensionWrapper::Shared { user_config, .. } => user_config.clone(),
        }
    }

    /// Returns the variant kind of this extension wrapper.
    #[must_use]
    pub fn variant(&self) -> ExtensionVariant {
        match self {
            ExtensionWrapper::Local { .. } => ExtensionVariant::Local,
            ExtensionWrapper::Shared { .. } => ExtensionVariant::Shared,
        }
    }

    /// Returns `true` if this extension is passive (no active lifecycle).
    #[must_use]
    pub fn is_passive(&self) -> bool {
        match self {
            ExtensionWrapper::Local { lifecycle, .. } => {
                matches!(lifecycle, ExtensionLifecycle::Passive)
            }
            ExtensionWrapper::Shared { lifecycle, .. } => {
                matches!(lifecycle, ExtensionLifecycle::Passive)
            }
        }
    }

    /// Returns the shared instance factory, if this is a shared variant.
    #[must_use]
    pub fn shared_instance_factory(&self) -> Option<&SharedInstanceFactory> {
        match self {
            ExtensionWrapper::Shared {
                instance_factory, ..
            } => Some(instance_factory),
            _ => None,
        }
    }

    /// Returns the local instance factory, if this is a local variant.
    #[must_use]
    pub fn local_instance_factory(&self) -> Option<&LocalInstanceFactory> {
        match self {
            ExtensionWrapper::Local {
                instance_factory, ..
            } => Some(instance_factory),
            _ => None,
        }
    }

    pub(crate) fn with_entity_telemetry_guard(mut self, guard: EntityTelemetryGuard) -> Self {
        match &mut self {
            ExtensionWrapper::Local { telemetry, .. }
            | ExtensionWrapper::Shared { telemetry, .. } => *telemetry = Some(guard),
        }
        self
    }

    /// Extracts the wrapper's `EntityTelemetryGuard`, if one was wired.
    pub(crate) fn take_telemetry_guard(&mut self) -> Option<EntityTelemetryGuard> {
        match self {
            ExtensionWrapper::Local { telemetry, .. }
            | ExtensionWrapper::Shared { telemetry, .. } => telemetry.take(),
        }
    }

    pub(crate) fn with_control_channel_metrics(
        self,
        entity_handle: &EntityTelemetryHandle,
        ext_ctx: &ExtensionContext,
        channel_metrics: &mut ChannelMetricsRegistry,
        channel_metrics_enabled: bool,
    ) -> Self {
        match self {
            ExtensionWrapper::Local {
                name,
                user_config,
                runtime_config,
                telemetry,
                lifecycle,
                instance_factory,
            } => {
                let lifecycle = match lifecycle {
                    ExtensionLifecycle::Passive => ExtensionLifecycle::Passive,
                    ExtensionLifecycle::Active {
                        extension,
                        control_sender,
                        control_receiver,
                        shutdown_sender,
                        shutdown_receiver,
                        readiness_probe,
                        readiness_signaller,
                    } => {
                        let capacity = runtime_config.control_channel.capacity as u64;
                        let (local_sender, local_receiver) = match control_sender {
                            Sender::Local(s) => (s, control_receiver),
                            _ => unreachable!("Local variant always has local sender"),
                        };
                        let (s, r) = wrap_extension_control_channel_metrics::<
                            LocalMode,
                            ExtensionControlMsg,
                        >(
                            name.clone(),
                            ExtensionVariant::Local,
                            entity_handle,
                            ext_ctx,
                            channel_metrics,
                            channel_metrics_enabled,
                            capacity,
                            local_sender,
                            local_receiver,
                        );
                        ExtensionLifecycle::Active {
                            extension,
                            control_sender: Sender::Local(s),
                            control_receiver: r,
                            shutdown_sender,
                            shutdown_receiver,
                            readiness_probe,
                            readiness_signaller,
                        }
                    }
                };
                ExtensionWrapper::Local {
                    name,
                    user_config,
                    runtime_config,
                    telemetry,
                    lifecycle,
                    instance_factory,
                }
            }
            ExtensionWrapper::Shared {
                name,
                user_config,
                runtime_config,
                telemetry,
                lifecycle,
                instance_factory,
            } => {
                let lifecycle = match lifecycle {
                    ExtensionLifecycle::Passive => ExtensionLifecycle::Passive,
                    ExtensionLifecycle::Active {
                        extension,
                        control_sender,
                        control_receiver,
                        shutdown_sender,
                        shutdown_receiver,
                        readiness_probe,
                        readiness_signaller,
                    } => {
                        let capacity = runtime_config.control_channel.capacity as u64;
                        let shared_sender = match control_sender {
                            Sender::Shared(s) => s,
                            _ => unreachable!("Shared variant always has shared sender"),
                        };
                        let (s, r) = wrap_extension_control_channel_metrics::<
                            SharedMode,
                            ExtensionControlMsg,
                        >(
                            name.clone(),
                            ExtensionVariant::Shared,
                            entity_handle,
                            ext_ctx,
                            channel_metrics,
                            channel_metrics_enabled,
                            capacity,
                            shared_sender,
                            control_receiver,
                        );
                        ExtensionLifecycle::Active {
                            extension,
                            control_sender: Sender::Shared(s),
                            control_receiver: r,
                            shutdown_sender,
                            shutdown_receiver,
                            readiness_probe,
                            readiness_signaller,
                        }
                    }
                };
                ExtensionWrapper::Shared {
                    name,
                    user_config,
                    runtime_config,
                    telemetry,
                    lifecycle,
                    instance_factory,
                }
            }
        }
    }

    /// Returns the `ExtensionControlSender` for sending control messages,
    /// or `None` if this extension is passive.
    #[allow(dead_code)]
    pub(crate) fn extension_control_sender(
        &self,
    ) -> Option<crate::control::ExtensionControlSender> {
        let sender = match self {
            ExtensionWrapper::Local {
                lifecycle: ExtensionLifecycle::Active { control_sender, .. },
                ..
            }
            | ExtensionWrapper::Shared {
                lifecycle: ExtensionLifecycle::Active { control_sender, .. },
                ..
            } => control_sender,
            _ => return None,
        };
        Some(crate::control::ExtensionControlSender {
            sender: sender.clone(),
        })
    }

    /// Takes the owning sender half of the extension's dedicated
    /// shutdown oneshot. Returns `None` for passive extensions and on
    /// subsequent calls (single-use by construction).
    pub(crate) fn take_shutdown_sender(
        &mut self,
    ) -> Option<crate::control::ExtensionShutdownChannel> {
        let (name, sender_slot) = match self {
            ExtensionWrapper::Local {
                name,
                lifecycle:
                    ExtensionLifecycle::Active {
                        shutdown_sender, ..
                    },
                ..
            }
            | ExtensionWrapper::Shared {
                name,
                lifecycle:
                    ExtensionLifecycle::Active {
                        shutdown_sender, ..
                    },
                ..
            } => (name, shutdown_sender),
            _ => return None,
        };
        let sender = sender_slot.take()?;
        Some(crate::control::ExtensionShutdownChannel {
            name: name.clone(),
            sender,
        })
    }

    /// Takes the readiness probe placed on this variant when its
    /// extension opted in. Single-use; returns `None` for passive
    /// extensions and for variants that did not opt in.
    pub(crate) fn take_readiness_probe(&mut self) -> Option<super::readiness::ReadinessProbe> {
        let probe_slot = match self {
            ExtensionWrapper::Local {
                lifecycle:
                    ExtensionLifecycle::Active {
                        readiness_probe, ..
                    },
                ..
            }
            | ExtensionWrapper::Shared {
                lifecycle:
                    ExtensionLifecycle::Active {
                        readiness_probe, ..
                    },
                ..
            } => readiness_probe,
            _ => return None,
        };
        probe_slot.take()
    }

    /// Returns `true` if a readiness signaller is attached to this
    /// variant -- i.e. the extension opted in via the builder's
    /// `with_readiness_probe`. The signaller itself is private to the
    /// engine; it is consumed by [`Self::start`] and threaded into the
    /// extension's [`EffectHandler`]. Intended for tests that need to
    /// verify the builder wired both halves of the probe pair.
    #[cfg(test)]
    pub(crate) fn has_readiness_signaller(&self) -> bool {
        match self {
            ExtensionWrapper::Local {
                lifecycle:
                    ExtensionLifecycle::Active {
                        readiness_signaller,
                        ..
                    },
                ..
            }
            | ExtensionWrapper::Shared {
                lifecycle:
                    ExtensionLifecycle::Active {
                        readiness_signaller,
                        ..
                    },
                ..
            } => readiness_signaller.is_some(),
            _ => false,
        }
    }

    /// Starts the extension lifecycle.
    ///
    /// # Errors
    ///
    /// Returns an error if this is a passive extension (no active lifecycle).
    pub(crate) async fn start(
        self,
        metrics_reporter: MetricsReporter,
    ) -> Result<TerminalState, Error> {
        match self {
            ExtensionWrapper::Local {
                name,
                lifecycle:
                    ExtensionLifecycle::Active {
                        extension,
                        control_receiver,
                        shutdown_receiver,
                        readiness_signaller,
                        ..
                    },
                ..
            } => {
                otel_debug!("extension.start.local", name = name.as_ref());
                let effect_handler =
                    EffectHandler::new(name, metrics_reporter, readiness_signaller);
                let shutdown_rx = shutdown_receiver.expect(
                    "shutdown_receiver must be present when an active extension is started",
                );
                extension
                    .start(
                        local_ext::ControlChannel::new(control_receiver, shutdown_rx),
                        effect_handler,
                    )
                    .await
            }
            ExtensionWrapper::Shared {
                name,
                lifecycle:
                    ExtensionLifecycle::Active {
                        extension,
                        control_receiver,
                        shutdown_receiver,
                        readiness_signaller,
                        ..
                    },
                ..
            } => {
                otel_debug!("extension.start.shared", name = name.as_ref());
                let effect_handler =
                    EffectHandler::new(name, metrics_reporter, readiness_signaller);
                let shutdown_rx = shutdown_receiver.expect(
                    "shutdown_receiver must be present when an active extension is started",
                );
                extension
                    .start(
                        shared_ext::ControlChannel::new(control_receiver, shutdown_rx),
                        effect_handler,
                    )
                    .await
            }
            _ => Err(Error::InternalError {
                message: "start() called on passive extension".into(),
            }),
        }
    }
}

impl ExtensionBundle {
    /// Assemble an [`ExtensionBundle`] from already-constructed wrappers.
    /// Called by [`super::builder::ExtensionBundleBuilder::build`].
    pub(super) fn from_parts(
        local: Option<ExtensionWrapper>,
        shared: Option<ExtensionWrapper>,
    ) -> Self {
        ExtensionBundle { local, shared }
    }

    /// Returns a reference to the local extension variant, if any.
    #[must_use]
    pub fn local(&self) -> Option<&ExtensionWrapper> {
        self.local.as_ref()
    }

    /// Returns a reference to the shared extension variant, if any.
    #[must_use]
    pub fn shared(&self) -> Option<&ExtensionWrapper> {
        self.shared.as_ref()
    }

    /// Takes ownership of the local extension variant, if any.
    pub fn take_local(&mut self) -> Option<ExtensionWrapper> {
        self.local.take()
    }

    /// Takes ownership of the shared extension variant, if any.
    pub fn take_shared(&mut self) -> Option<ExtensionWrapper> {
        self.shared.take()
    }

    /// Returns an iterator over the extension wrappers in this set.
    pub fn iter(&self) -> impl Iterator<Item = &ExtensionWrapper> {
        self.local.iter().chain(self.shared.iter())
    }

    /// Register this bundle's capabilities into the given registry.
    ///
    /// No-op when `capabilities` is `None` (Background extensions).
    ///
    /// # Errors
    ///
    /// - [`Error::InternalError`](crate::capability::registry::Error::InternalError)
    ///   if `capabilities` advertises an execution model the bundle
    ///   does not contain (metadata-vs-bundle drift caught early).
    /// - Any error returned by the capability registry
    ///   (e.g. duplicate `(capability, extension)` registration).
    pub fn register_into(
        &self,
        capabilities: Option<&crate::capability::ExtensionCapabilities>,
        registry: &mut crate::capability::registry::CapabilityRegistry,
    ) -> Result<(), crate::capability::registry::Error> {
        // Background extensions: no capability registration.
        let Some(capabilities) = capabilities else {
            return Ok(());
        };

        // Fail fast on metadata-vs-bundle drift instead of letting it
        // surface later as a confusing config error.
        let bundle_name = self
            .local
            .as_ref()
            .or(self.shared.as_ref())
            .map(|w| w.name().to_string())
            .unwrap_or_else(|| "<empty bundle>".to_string());
        if !capabilities.shared.is_empty() && self.shared.is_none() {
            return Err(crate::capability::registry::Error::InternalError {
                message: format!(
                    "extension '{bundle_name}': extension_capabilities! advertises shared \
                     capabilities {caps:?} but the ExtensionBundle has no shared variant \
                     - either add `.shared(...)` to the builder chain or remove the \
                     capabilities from the macro list",
                    caps = capabilities.shared,
                ),
            });
        }
        if !capabilities.local.is_empty() && self.local.is_none() {
            return Err(crate::capability::registry::Error::InternalError {
                message: format!(
                    "extension '{bundle_name}': extension_capabilities! advertises local \
                     capabilities {caps:?} but the ExtensionBundle has no local variant \
                     - either add `.local(...)` to the builder chain or remove the \
                     capabilities from the macro list",
                    caps = capabilities.local,
                ),
            });
        }

        if let Some(shared) = self.shared.as_ref()
            && let Some(factory) = shared.shared_instance_factory()
        {
            (capabilities.register_shared)(shared.name(), factory.clone(), registry)?;
        }
        if let Some(local) = self.local.as_ref()
            && let Some(factory) = local.local_instance_factory()
        {
            (capabilities.register_local)(local.name(), factory.clone(), registry)?;
        }
        Ok(())
    }

    /// Wire entity telemetry and control-channel metrics onto each variant
    /// present in this bundle. Local and Shared variants of the same
    /// `ExtensionId` get **distinct** entity registrations so their
    /// telemetry stays separable downstream.
    pub(crate) fn wire_telemetry(
        &mut self,
        extension_id: ExtensionId,
        ext_ctx: &ExtensionContext,
        channel_metrics: &mut ChannelMetricsRegistry,
        channel_metrics_enabled: bool,
    ) -> ExtensionEntityKeys {
        let mut keys = ExtensionEntityKeys::default();
        if let Some(w) = self.local.take() {
            let entity_key =
                ext_ctx.register_extension_entity(extension_id.clone(), ExtensionVariant::Local);
            let handle = EntityTelemetryHandle::new(ext_ctx.metrics_registry(), entity_key);
            keys.local = Some(entity_key);
            let w = w.with_control_channel_metrics(
                &handle,
                ext_ctx,
                channel_metrics,
                channel_metrics_enabled,
            );
            self.local = Some(w.with_entity_telemetry_guard(EntityTelemetryGuard::new(handle)));
        }
        if let Some(w) = self.shared.take() {
            let entity_key =
                ext_ctx.register_extension_entity(extension_id, ExtensionVariant::Shared);
            let handle = EntityTelemetryHandle::new(ext_ctx.metrics_registry(), entity_key);
            keys.shared = Some(entity_key);
            let w = w.with_control_channel_metrics(
                &handle,
                ext_ctx,
                channel_metrics,
                channel_metrics_enabled,
            );
            self.shared = Some(w.with_entity_telemetry_guard(EntityTelemetryGuard::new(handle)));
        }
        keys
    }
}

/// Entity keys assigned to each variant in a wired [`ExtensionBundle`].
#[derive(Debug, Default, Clone, Copy)]
pub(crate) struct ExtensionEntityKeys {
    pub local: Option<otel_arrow_dfe_telemetry::registry::EntityKey>,
    pub shared: Option<otel_arrow_dfe_telemetry::registry::EntityKey>,
}