vector/topology/

builder.rs

use std::{
    collections::HashMap,
    future::ready,
    num::NonZeroUsize,
    sync::{Arc, LazyLock, Mutex},
    time::Instant,
};

use futures::{FutureExt, StreamExt, TryStreamExt, stream::FuturesOrdered};
use futures_util::stream::FuturesUnordered;
use metrics::gauge;
use stream_cancel::{StreamExt as StreamCancelExt, Trigger, Tripwire};
use tokio::{
    select,
    sync::{mpsc::UnboundedSender, oneshot},
    time::timeout,
};
use tracing::Instrument;
use vector_lib::{
    EstimatedJsonEncodedSizeOf,
    buffers::{
        BufferType, WhenFull,
        topology::{
            builder::TopologyBuilder,
            channel::{BufferReceiver, BufferSender, ChannelMetricMetadata, LimitedReceiver},
        },
    },
    internal_event::{self, CountByteSize, EventsSent, InternalEventHandle as _, Registered},
    latency::LatencyRecorder,
    schema::Definition,
    source_sender::{CHUNK_SIZE, SourceSenderItem},
    transform::update_runtime_schema_definition,
};
use vector_vrl_metrics::MetricsStorage;

use super::{
    BuiltBuffer, ConfigDiff,
    fanout::{self, Fanout},
    schema,
    task::{Task, TaskOutput, TaskResult},
};
use crate::{
    SourceSender,
    config::{
        ComponentKey, Config, DataType, EnrichmentTableConfig, Input, Inputs, OutputId,
        ProxyConfig, SinkContext, SourceContext, TransformContext, TransformOuter, TransformOutput,
    },
    event::{EventArray, EventContainer},
    extra_context::ExtraContext,
    internal_events::EventsReceived,
    shutdown::SourceShutdownCoordinator,
    spawn_named,
    topology::task::TaskError,
    transforms::{SyncTransform, TaskTransform, Transform, TransformOutputs, TransformOutputsBuf},
    utilization::{
        OutputUtilization, Utilization, UtilizationComponentSender, UtilizationEmitter,
        UtilizationRegistry,
    },
};

static ENRICHMENT_TABLES: LazyLock<vector_lib::enrichment::TableRegistry> =
    LazyLock::new(vector_lib::enrichment::TableRegistry::default);
static METRICS_STORAGE: LazyLock<MetricsStorage> = LazyLock::new(MetricsStorage::default);

pub(crate) static SOURCE_SENDER_BUFFER_SIZE: LazyLock<usize> =
    LazyLock::new(|| *TRANSFORM_CONCURRENCY_LIMIT * CHUNK_SIZE);

const READY_ARRAY_CAPACITY: NonZeroUsize = NonZeroUsize::new(CHUNK_SIZE * 4).unwrap();
pub(crate) const TOPOLOGY_BUFFER_SIZE: NonZeroUsize = NonZeroUsize::new(100).unwrap();
const TRANSFORM_CHANNEL_METRIC_PREFIX: &str = "transform_buffer";

static TRANSFORM_CONCURRENCY_LIMIT: LazyLock<usize> = LazyLock::new(|| {
    crate::app::worker_threads()
        .map(std::num::NonZeroUsize::get)
        .unwrap_or_else(crate::num_threads)
});

const INTERNAL_SOURCES: [&str; 2] = ["internal_logs", "internal_metrics"];

struct Builder<'a> {
    config: &'a super::Config,
    diff: &'a ConfigDiff,
    shutdown_coordinator: SourceShutdownCoordinator,
    errors: Vec<String>,
    outputs: HashMap<OutputId, UnboundedSender<fanout::ControlMessage>>,
    tasks: HashMap<ComponentKey, Task>,
    buffers: HashMap<ComponentKey, BuiltBuffer>,
    inputs: HashMap<ComponentKey, (BufferSender<EventArray>, Inputs<OutputId>)>,
    healthchecks: HashMap<ComponentKey, Task>,
    detach_triggers: HashMap<ComponentKey, Trigger>,
    extra_context: ExtraContext,
    utilization_emitter: Option<UtilizationEmitter>,
    utilization_registry: UtilizationRegistry,
}

impl<'a> Builder<'a> {
    fn new(
        config: &'a super::Config,
        diff: &'a ConfigDiff,
        buffers: HashMap<ComponentKey, BuiltBuffer>,
        extra_context: ExtraContext,
        utilization_registry: Option<UtilizationRegistry>,
    ) -> Self {
        // If registry is not passed, we need to build a whole new utilization emitter + registry
        // Otherwise, we just store the registry and reuse it for this build
        let (emitter, registry) = if let Some(registry) = utilization_registry {
            (None, registry)
        } else {
            let (emitter, registry) = UtilizationEmitter::new();
            (Some(emitter), registry)
        };
        Self {
            config,
            diff,
            buffers,
            shutdown_coordinator: SourceShutdownCoordinator::default(),
            errors: vec![],
            outputs: HashMap::new(),
            tasks: HashMap::new(),
            inputs: HashMap::new(),
            healthchecks: HashMap::new(),
            detach_triggers: HashMap::new(),
            extra_context,
            utilization_emitter: emitter,
            utilization_registry: registry,
        }
    }

    /// Builds the new pieces of the topology found in `self.diff`.
    async fn build(mut self) -> Result<TopologyPieces, Vec<String>> {
        let enrichment_tables = self.load_enrichment_tables().await;
        let source_tasks = self.build_sources(enrichment_tables).await;
        self.build_transforms(enrichment_tables).await;
        self.build_sinks(enrichment_tables).await;

        // We should have all the data for the enrichment tables loaded now, so switch them over to
        // readonly.
        enrichment_tables.finish_load();

        if self.errors.is_empty() {
            Ok(TopologyPieces {
                inputs: self.inputs,
                outputs: Self::finalize_outputs(self.outputs),
                tasks: self.tasks,
                source_tasks,
                healthchecks: self.healthchecks,
                shutdown_coordinator: self.shutdown_coordinator,
                detach_triggers: self.detach_triggers,
                metrics_storage: METRICS_STORAGE.clone(),
                utilization: self
                    .utilization_emitter
                    .map(|e| (e, self.utilization_registry)),
            })
        } else {
            Err(self.errors)
        }
    }

    fn finalize_outputs(
        outputs: HashMap<OutputId, UnboundedSender<fanout::ControlMessage>>,
    ) -> HashMap<ComponentKey, HashMap<Option<String>, UnboundedSender<fanout::ControlMessage>>>
    {
        let mut finalized_outputs = HashMap::new();
        for (id, output) in outputs {
            let entry = finalized_outputs
                .entry(id.component)
                .or_insert_with(HashMap::new);
            entry.insert(id.port, output);
        }

        finalized_outputs
    }

    /// Loads, or reloads the enrichment tables.
    /// The tables are stored in the `ENRICHMENT_TABLES` global variable.
    async fn load_enrichment_tables(&mut self) -> &'static vector_lib::enrichment::TableRegistry {
        let mut enrichment_tables = HashMap::new();

        // Build enrichment tables
        'tables: for (name, table_outer) in self.config.enrichment_tables.iter() {
            let table_name = name.to_string();
            if ENRICHMENT_TABLES.needs_reload(&table_name) {
                let indexes = if !self.diff.enrichment_tables.is_added(name) {
                    // If this is an existing enrichment table, we need to store the indexes to reapply
                    // them again post load.
                    Some(ENRICHMENT_TABLES.index_fields(&table_name))
                } else {
                    None
                };

                let mut table = match table_outer.inner.build(&self.config.global).await {
                    Ok(table) => table,
                    Err(error) => {
                        self.errors
                            .push(format!("Enrichment Table \"{name}\": {error}"));
                        continue;
                    }
                };

                if let Some(indexes) = indexes {
                    for (case, index) in indexes {
                        match table
                            .add_index(case, &index.iter().map(|s| s.as_ref()).collect::<Vec<_>>())
                        {
                            Ok(_) => (),
                            Err(error) => {
                                // If there is an error adding an index we do not want to use the reloaded
                                // data, the previously loaded data will still need to be used.
                                // Just report the error and continue.
                                error!(message = "Unable to add index to reloaded enrichment table.",
                                    table = ?name.to_string(),
                                    %error);
                                continue 'tables;
                            }
                        }
                    }
                }

                enrichment_tables.insert(table_name, table);
            }
        }

        ENRICHMENT_TABLES.load(enrichment_tables);

        &ENRICHMENT_TABLES
    }

    async fn build_sources(
        &mut self,
        enrichment_tables: &vector_lib::enrichment::TableRegistry,
    ) -> HashMap<ComponentKey, Task> {
        let mut source_tasks = HashMap::new();

        let table_sources = self
            .config
            .enrichment_tables
            .iter()
            .filter_map(|(key, table)| table.as_source(key))
            .collect::<Vec<_>>();
        for (key, source) in self
            .config
            .sources()
            .filter(|(key, _)| self.diff.sources.contains_new(key))
            .chain(
                table_sources
                    .iter()
                    .map(|(key, source)| (key, source))
                    .filter(|(key, _)| self.diff.enrichment_tables.contains_new(key)),
            )
        {
            debug!(component_id = %key, "Building new source.");

            let typetag = source.inner.get_component_name();
            let source_outputs = source.inner.outputs(self.config.schema.log_namespace());

            let span = error_span!(
                "source",
                component_kind = "source",
                component_id = %key.id(),
                component_type = %source.inner.get_component_name(),
            );
            let _entered_span = span.enter();

            let task_name = format!(
                ">> {} ({}, pump) >>",
                source.inner.get_component_name(),
                key.id()
            );

            let mut builder = SourceSender::builder()
                .with_buffer(*SOURCE_SENDER_BUFFER_SIZE)
                .with_timeout(source.inner.send_timeout())
                .with_ewma_half_life_seconds(
                    self.config.global.buffer_utilization_ewma_half_life_seconds,
                );
            let mut pumps = Vec::new();
            let mut controls = HashMap::new();
            let mut schema_definitions = HashMap::with_capacity(source_outputs.len());

            for output in source_outputs.into_iter() {
                let rx = builder.add_source_output(output.clone(), key.clone());

                let (fanout, control) = Fanout::new();
                let source_type = source.inner.get_component_name();
                let source = Arc::new(key.clone());

                let pump = run_source_output_pump(rx, fanout, source, source_type);

                pumps.push(pump.instrument(span.clone()));
                controls.insert(
                    OutputId {
                        component: key.clone(),
                        port: output.port.clone(),
                    },
                    control,
                );

                let port = output.port.clone();
                if let Some(definition) = output.schema_definition(self.config.schema.enabled) {
                    schema_definitions.insert(port, definition);
                }
            }

            let (pump_error_tx, mut pump_error_rx) = oneshot::channel();
            let pump = async move {
                debug!("Source pump supervisor starting.");

                // Spawn all of the per-output pumps and then await their completion.
                //
                // If any of the pumps complete with an error, or panic/are cancelled, we return
                // immediately.
                let mut handles = FuturesUnordered::new();
                for pump in pumps {
                    handles.push(spawn_named(pump, task_name.as_ref()));
                }

                let mut had_pump_error = false;
                while let Some(output) = handles.try_next().await? {
                    if let Err(e) = output {
                        // Immediately send the error to the source's wrapper future, but ignore any
                        // errors during the send, since nested errors wouldn't make any sense here.
                        _ = pump_error_tx.send(e);
                        had_pump_error = true;
                        break;
                    }
                }

                if had_pump_error {
                    debug!("Source pump supervisor task finished with an error.");
                } else {
                    debug!("Source pump supervisor task finished normally.");
                }
                Ok(TaskOutput::Source)
            };
            let pump = Task::new(key.clone(), typetag, pump);

            let (shutdown_signal, force_shutdown_tripwire) = self
                .shutdown_coordinator
                .register_source(key, INTERNAL_SOURCES.contains(&typetag));

            let context = SourceContext {
                key: key.clone(),
                globals: self.config.global.clone(),
                enrichment_tables: enrichment_tables.clone(),
                metrics_storage: METRICS_STORAGE.clone(),
                shutdown: shutdown_signal,
                out: builder.build(),
                proxy: ProxyConfig::merge_with_env(&self.config.global.proxy, &source.proxy),
                acknowledgements: source.sink_acknowledgements,
                schema_definitions,
                schema: self.config.schema,
                extra_context: self.extra_context.clone(),
            };
            let server = match source.inner.build(context).await {
                Err(error) => {
                    self.errors.push(format!("Source \"{key}\": {error}"));
                    continue;
                }
                Ok(server) => server,
            };

            // Build a wrapper future that drives the actual source future, but returns early if we've
            // been signalled to forcefully shutdown, or if the source pump encounters an error.
            //
            // The forceful shutdown will only resolve if the source itself doesn't shutdown gracefully
            // within the allotted time window. This can occur normally for certain sources, like stdin,
            // where the I/O is blocking (in a separate thread) and won't wake up to check if it's time
            // to shutdown unless some input is given.
            let server = async move {
                debug!("Source starting.");

                let mut result = select! {
                    biased;

                    // We've been told that we must forcefully shut down.
                    _ = force_shutdown_tripwire => Ok(()),

                    // The source pump encountered an error, which we're now bubbling up here to stop
                    // the source as well, since the source running makes no sense without the pump.
                    //
                    // We only match receiving a message, not the error of the sender being dropped,
                    // just to keep things simpler.
                    Ok(e) = &mut pump_error_rx => Err(e),

                    // The source finished normally.
                    result = server => result.map_err(|_| TaskError::Opaque),
                };

                // Even though we already tried to receive any pump task error above, we may have exited
                // on the source itself returning an error due to task scheduling, where the pump task
                // encountered an error, sent it over the oneshot, but we were polling the source
                // already and hit an error trying to send to the now-shutdown pump task.
                //
                // Since the error from the source is opaque at the moment (i.e. `()`), we try a final
                // time to see if the pump task encountered an error, using _that_ instead if so, to
                // propagate the true error that caused the source to have to stop.
                if let Ok(e) = pump_error_rx.try_recv() {
                    result = Err(e);
                }

                match result {
                    Ok(()) => {
                        debug!("Source finished normally.");
                        Ok(TaskOutput::Source)
                    }
                    Err(e) => {
                        debug!("Source finished with an error.");
                        Err(e)
                    }
                }
            };
            let server = Task::new(key.clone(), typetag, server);

            self.outputs.extend(controls);
            self.tasks.insert(key.clone(), pump);
            source_tasks.insert(key.clone(), server);
        }

        source_tasks
    }

    async fn build_transforms(
        &mut self,
        enrichment_tables: &vector_lib::enrichment::TableRegistry,
    ) {
        let mut definition_cache = HashMap::default();

        for (key, transform) in self
            .config
            .transforms()
            .filter(|(key, _)| self.diff.transforms.contains_new(key))
        {
            debug!(component_id = %key, "Building new transform.");

            let input_definitions = match schema::input_definitions(
                &transform.inputs,
                self.config,
                enrichment_tables.clone(),
                &mut definition_cache,
            ) {
                Ok(definitions) => definitions,
                Err(_) => {
                    // We have received an error whilst retrieving the definitions,
                    // there is no point in continuing.

                    return;
                }
            };

            let merged_definition: Definition = input_definitions
                .iter()
                .map(|(_output_id, definition)| definition.clone())
                .reduce(Definition::merge)
                // We may not have any definitions if all the inputs are from metrics sources.
                .unwrap_or_else(Definition::any);

            let span = error_span!(
                "transform",
                component_kind = "transform",
                component_id = %key.id(),
                component_type = %transform.inner.get_component_name(),
            );
            let _span = span.enter();

            // Create a map of the outputs to the list of possible definitions from those outputs.
            let schema_definitions = transform
                .inner
                .outputs(
                    &TransformContext {
                        enrichment_tables: enrichment_tables.clone(),
                        metrics_storage: METRICS_STORAGE.clone(),
                        schema: self.config.schema,
                        ..Default::default()
                    },
                    &input_definitions,
                )
                .into_iter()
                .map(|output| {
                    let definitions = output.schema_definitions(self.config.schema.enabled);
                    (output.port, definitions)
                })
                .collect::<HashMap<_, _>>();

            let context = TransformContext {
                key: Some(key.clone()),
                globals: self.config.global.clone(),
                enrichment_tables: enrichment_tables.clone(),
                metrics_storage: METRICS_STORAGE.clone(),
                schema_definitions,
                merged_schema_definition: merged_definition.clone(),
                schema: self.config.schema,
                extra_context: self.extra_context.clone(),
            };

            let node =
                TransformNode::from_parts(key.clone(), &context, transform, &input_definitions);

            let transform = match transform
                .inner
                .build(&context)
                .instrument(span.clone())
                .await
            {
                Err(error) => {
                    self.errors.push(format!("Transform \"{key}\": {error}"));
                    continue;
                }
                Ok(transform) => transform,
            };

            let metrics = ChannelMetricMetadata::new(TRANSFORM_CHANNEL_METRIC_PREFIX, None);
            let (input_tx, input_rx) = TopologyBuilder::standalone_memory(
                TOPOLOGY_BUFFER_SIZE,
                WhenFull::Block,
                &span,
                Some(metrics),
                self.config.global.buffer_utilization_ewma_half_life_seconds,
            );

            self.inputs
                .insert(key.clone(), (input_tx, node.inputs.clone()));

            let (transform_task, transform_outputs) =
                self.build_transform(transform, node, input_rx);

            self.outputs.extend(transform_outputs);
            self.tasks.insert(key.clone(), transform_task);
        }
    }

    async fn build_sinks(&mut self, enrichment_tables: &vector_lib::enrichment::TableRegistry) {
        let table_sinks = self
            .config
            .enrichment_tables
            .iter()
            .filter_map(|(key, table)| table.as_sink(key))
            .collect::<Vec<_>>();
        for (key, sink) in self
            .config
            .sinks()
            .filter(|(key, _)| self.diff.sinks.contains_new(key))
            .chain(
                table_sinks
                    .iter()
                    .map(|(key, sink)| (key, sink))
                    .filter(|(key, _)| self.diff.enrichment_tables.contains_new(key)),
            )
        {
            debug!(component_id = %key, "Building new sink.");

            let sink_inputs = &sink.inputs;
            let healthcheck = sink.healthcheck();
            let enable_healthcheck = healthcheck.enabled && self.config.healthchecks.enabled;
            let healthcheck_timeout = healthcheck.timeout;

            let typetag = sink.inner.get_component_name();
            let input_type = sink.inner.input().data_type();

            let span = error_span!(
                "sink",
                component_kind = "sink",
                component_id = %key.id(),
                component_type = %sink.inner.get_component_name(),
            );
            let _entered_span = span.enter();

            // At this point, we've validated that all transforms are valid, including any
            // transform that mutates the schema provided by their sources. We can now validate the
            // schema expectations of each individual sink.
            if let Err(mut err) = schema::validate_sink_expectations(
                key,
                sink,
                self.config,
                enrichment_tables.clone(),
            ) {
                self.errors.append(&mut err);
            };

            let (tx, rx) = match self.buffers.remove(key) {
                Some(buffer) => buffer,
                _ => {
                    let buffer_type =
                        match sink.buffer.stages().first().expect("cant ever be empty") {
                            BufferType::Memory { .. } => "memory",
                            BufferType::DiskV2 { .. } => "disk",
                        };
                    let buffer_span = error_span!("sink", buffer_type);
                    let buffer = sink
                        .buffer
                        .build(
                            self.config.global.data_dir.clone(),
                            key.to_string(),
                            buffer_span,
                        )
                        .await;
                    match buffer {
                        Err(error) => {
                            self.errors.push(format!("Sink \"{key}\": {error}"));
                            continue;
                        }
                        Ok((tx, rx)) => (tx, Arc::new(Mutex::new(Some(rx.into_stream())))),
                    }
                }
            };

            let cx = SinkContext {
                healthcheck,
                globals: self.config.global.clone(),
                enrichment_tables: enrichment_tables.clone(),
                metrics_storage: METRICS_STORAGE.clone(),
                proxy: ProxyConfig::merge_with_env(&self.config.global.proxy, sink.proxy()),
                schema: self.config.schema,
                app_name: crate::get_app_name().to_string(),
                app_name_slug: crate::get_slugified_app_name(),
                extra_context: self.extra_context.clone(),
            };

            let (sink, healthcheck) = match sink.inner.build(cx).await {
                Err(error) => {
                    self.errors.push(format!("Sink \"{key}\": {error}"));
                    continue;
                }
                Ok(built) => built,
            };

            let (trigger, tripwire) = Tripwire::new();

            let utilization_sender = self
                .utilization_registry
                .add_component(key.clone(), gauge!("utilization"));
            let component_key = key.clone();
            let sink = async move {
                debug!("Sink starting.");

                // Why is this Arc<Mutex<Option<_>>> needed you ask.
                // In case when this function build_pieces errors
                // this future won't be run so this rx won't be taken
                // which will enable us to reuse rx to rebuild
                // old configuration by passing this Arc<Mutex<Option<_>>>
                // yet again.
                let rx = rx
                    .lock()
                    .unwrap()
                    .take()
                    .expect("Task started but input has been taken.");

                let mut rx = Utilization::new(utilization_sender, component_key.clone(), rx);

                let events_received = register!(EventsReceived);
                sink.run(
                    rx.by_ref()
                        .filter(|events: &EventArray| ready(filter_events_type(events, input_type)))
                        .inspect(|events| {
                            events_received.emit(CountByteSize(
                                events.len(),
                                events.estimated_json_encoded_size_of(),
                            ))
                        })
                        .take_until_if(tripwire),
                )
                .await
                .map(|_| {
                    debug!("Sink finished normally.");
                    TaskOutput::Sink(rx)
                })
                .map_err(|_| {
                    debug!("Sink finished with an error.");
                    TaskError::Opaque
                })
            };

            let task = Task::new(key.clone(), typetag, sink);

            let component_key = key.clone();
            let healthcheck_task = async move {
                if enable_healthcheck {
                    timeout(healthcheck_timeout, healthcheck)
                        .map(|result| match result {
                            Ok(Ok(_)) => {
                                info!("Healthcheck passed.");
                                Ok(TaskOutput::Healthcheck)
                            }
                            Ok(Err(error)) => {
                                error!(
                                    msg = "Healthcheck failed.",
                                    %error,
                                    component_kind = "sink",
                                    component_type = typetag,
                                    component_id = %component_key.id(),
                                );
                                Err(TaskError::wrapped(error))
                            }
                            Err(e) => {
                                error!(
                                    msg = "Healthcheck timed out.",
                                    component_kind = "sink",
                                    component_type = typetag,
                                    component_id = %component_key.id(),
                                );
                                Err(TaskError::wrapped(Box::new(e)))
                            }
                        })
                        .await
                } else {
                    info!("Healthcheck disabled.");
                    Ok(TaskOutput::Healthcheck)
                }
            };

            let healthcheck_task = Task::new(key.clone(), typetag, healthcheck_task);

            self.inputs.insert(key.clone(), (tx, sink_inputs.clone()));
            self.healthchecks.insert(key.clone(), healthcheck_task);
            self.tasks.insert(key.clone(), task);
            self.detach_triggers.insert(key.clone(), trigger);
        }
    }

    fn build_transform(
        &self,
        transform: Transform,
        node: TransformNode,
        input_rx: BufferReceiver<EventArray>,
    ) -> (Task, HashMap<OutputId, fanout::ControlChannel>) {
        match transform {
            // TODO: avoid the double boxing for function transforms here
            Transform::Function(t) => self.build_sync_transform(Box::new(t), node, input_rx),
            Transform::Synchronous(t) => self.build_sync_transform(t, node, input_rx),
            Transform::Task(t) => self.build_task_transform(
                t,
                input_rx,
                node.input_details.data_type(),
                node.typetag,
                &node.key,
                &node.outputs,
            ),
        }
    }

    fn build_sync_transform(
        &self,
        t: Box<dyn SyncTransform>,
        node: TransformNode,
        input_rx: BufferReceiver<EventArray>,
    ) -> (Task, HashMap<OutputId, fanout::ControlChannel>) {
        let (outputs, controls) = TransformOutputs::new(node.outputs, &node.key);

        let sender = self
            .utilization_registry
            .add_component(node.key.clone(), gauge!("utilization"));
        let runner = Runner::new(
            t,
            input_rx,
            sender,
            node.input_details.data_type(),
            outputs,
            LatencyRecorder::new(self.config.global.latency_ewma_alpha),
        );
        let transform = if node.enable_concurrency {
            runner.run_concurrently().boxed()
        } else {
            runner.run_inline().boxed()
        };

        let transform = async move {
            debug!("Synchronous transform starting.");

            match transform.await {
                Ok(v) => {
                    debug!("Synchronous transform finished normally.");
                    Ok(v)
                }
                Err(e) => {
                    debug!("Synchronous transform finished with an error.");
                    Err(e)
                }
            }
        };

        let mut output_controls = HashMap::new();
        for (name, control) in controls {
            let id = name
                .map(|name| OutputId::from((&node.key, name)))
                .unwrap_or_else(|| OutputId::from(&node.key));
            output_controls.insert(id, control);
        }

        let task = Task::new(node.key.clone(), node.typetag, transform);

        (task, output_controls)
    }

    fn build_task_transform(
        &self,
        t: Box<dyn TaskTransform<EventArray>>,
        input_rx: BufferReceiver<EventArray>,
        input_type: DataType,
        typetag: &str,
        key: &ComponentKey,
        outputs: &[TransformOutput],
    ) -> (Task, HashMap<OutputId, fanout::ControlChannel>) {
        let (mut fanout, control) = Fanout::new();

        let sender = self
            .utilization_registry
            .add_component(key.clone(), gauge!("utilization"));
        let output_sender = sender.clone();
        let input_rx = Utilization::new(sender, key.clone(), input_rx.into_stream());

        let events_received = register!(EventsReceived);
        let filtered = input_rx
            .filter(move |events| ready(filter_events_type(events, input_type)))
            .inspect(move |events| {
                events_received.emit(CountByteSize(
                    events.len(),
                    events.estimated_json_encoded_size_of(),
                ))
            });
        let events_sent = register!(EventsSent::from(internal_event::Output(None)));
        let output_id = Arc::new(OutputId {
            component: key.clone(),
            port: None,
        });
        let latency_recorder = LatencyRecorder::new(self.config.global.latency_ewma_alpha);

        // Task transforms can only write to the default output, so only a single schema def map is needed
        let schema_definition_map = outputs
            .iter()
            .find(|x| x.port.is_none())
            .expect("output for default port required for task transforms")
            .log_schema_definitions
            .clone()
            .into_iter()
            .map(|(key, value)| (key, Arc::new(value)))
            .collect();

        let stream = t
            .transform(Box::pin(filtered))
            .map(move |mut events| {
                for event in events.iter_events_mut() {
                    update_runtime_schema_definition(event, &output_id, &schema_definition_map);
                }
                let now = Instant::now();
                latency_recorder.on_send(&mut events, now);
                (events, now)
            })
            .inspect(move |(events, _): &(EventArray, Instant)| {
                events_sent.emit(CountByteSize(
                    events.len(),
                    events.estimated_json_encoded_size_of(),
                ));
            });
        let stream = OutputUtilization::new(output_sender, stream);
        let transform = async move {
            debug!("Task transform starting.");

            match fanout.send_stream(stream).await {
                Ok(()) => {
                    debug!("Task transform finished normally.");
                    Ok(TaskOutput::Transform)
                }
                Err(e) => {
                    debug!("Task transform finished with an error.");
                    Err(TaskError::wrapped(e))
                }
            }
        }
        .boxed();

        let mut outputs = HashMap::new();
        outputs.insert(OutputId::from(key), control);

        let task = Task::new(key.clone(), typetag, transform);

        (task, outputs)
    }
}

async fn run_source_output_pump(
    mut rx: LimitedReceiver<SourceSenderItem>,
    mut fanout: Fanout,
    source: Arc<ComponentKey>,
    source_type: &'static str,
) -> TaskResult {
    debug!("Source pump starting.");

    while let Some(SourceSenderItem {
        events: mut array,
        send_reference,
    }) = rx.next().await
    {
        // Even though we have a `send_reference` timestamp above, that reference time is when
        // the events were enqueued in the `SourceSender`, not when they were pulled out of the
        // `rx` stream on this end. Since those times can be quite different (due to blocking
        // inherent to the fanout send operation), we set the `last_transform_timestamp` to the
        // current time instead to get an accurate reference for when the events started waiting
        // for the first transform.
        let now = Instant::now();
        array.for_each_metadata_mut(|metadata| {
            metadata.set_source_id(Arc::clone(&source));
            metadata.set_source_type(source_type);
            metadata.set_last_transform_timestamp(now);
        });
        fanout
            .send(array, Some(send_reference))
            .await
            .map_err(|e| {
                debug!("Source pump finished with an error.");
                TaskError::wrapped(e)
            })?;
    }

    debug!("Source pump finished normally.");
    Ok(TaskOutput::Source)
}

pub async fn reload_enrichment_tables(config: &Config) {
    let mut enrichment_tables = HashMap::new();
    // Build enrichment tables
    'tables: for (name, table_outer) in config.enrichment_tables.iter() {
        let table_name = name.to_string();
        if ENRICHMENT_TABLES.needs_reload(&table_name) {
            let indexes = Some(ENRICHMENT_TABLES.index_fields(&table_name));

            let mut table = match table_outer.inner.build(&config.global).await {
                Ok(table) => table,
                Err(error) => {
                    error!("Enrichment table \"{name}\" reload failed: {error}");
                    continue;
                }
            };

            if let Some(indexes) = indexes {
                for (case, index) in indexes {
                    match table
                        .add_index(case, &index.iter().map(|s| s.as_ref()).collect::<Vec<_>>())
                    {
                        Ok(_) => (),
                        Err(error) => {
                            // If there is an error adding an index we do not want to use the reloaded
                            // data, the previously loaded data will still need to be used.
                            // Just report the error and continue.
                            error!(
                                message = "Unable to add index to reloaded enrichment table.",
                                table = ?name.to_string(),
                                %error
                            );
                            continue 'tables;
                        }
                    }
                }
            }

            enrichment_tables.insert(table_name, table);
        }
    }

    ENRICHMENT_TABLES.load(enrichment_tables);
    ENRICHMENT_TABLES.finish_load();
}

pub struct TopologyPieces {
    pub(super) inputs: HashMap<ComponentKey, (BufferSender<EventArray>, Inputs<OutputId>)>,
    pub(crate) outputs: HashMap<ComponentKey, HashMap<Option<String>, fanout::ControlChannel>>,
    pub(super) tasks: HashMap<ComponentKey, Task>,
    pub(crate) source_tasks: HashMap<ComponentKey, Task>,
    pub(super) healthchecks: HashMap<ComponentKey, Task>,
    pub(crate) shutdown_coordinator: SourceShutdownCoordinator,
    pub(crate) detach_triggers: HashMap<ComponentKey, Trigger>,
    pub(crate) metrics_storage: MetricsStorage,
    pub(crate) utilization: Option<(UtilizationEmitter, UtilizationRegistry)>,
}

/// Builder for constructing TopologyPieces with a fluent API.
///
/// # Examples
///
/// ```ignore
/// let pieces = TopologyPiecesBuilder::new(&config, &diff)
///     .with_buffers(buffers)
///     .with_extra_context(extra_context)
///     .build()
///     .await?;
/// ```
pub struct TopologyPiecesBuilder<'a> {
    config: &'a Config,
    diff: &'a ConfigDiff,
    buffers: HashMap<ComponentKey, BuiltBuffer>,
    extra_context: ExtraContext,
    utilization_registry: Option<UtilizationRegistry>,
}

impl<'a> TopologyPiecesBuilder<'a> {
    /// Creates a new builder with required parameters.
    pub fn new(config: &'a Config, diff: &'a ConfigDiff) -> Self {
        Self {
            config,
            diff,
            buffers: HashMap::new(),
            extra_context: ExtraContext::default(),
            utilization_registry: None,
        }
    }

    /// Sets the buffers for the topology.
    pub fn with_buffers(mut self, buffers: HashMap<ComponentKey, BuiltBuffer>) -> Self {
        self.buffers = buffers;
        self
    }

    /// Sets the extra context for the topology.
    pub fn with_extra_context(mut self, extra_context: ExtraContext) -> Self {
        self.extra_context = extra_context;
        self
    }

    /// Sets the utilization registry for the topology.
    pub fn with_utilization_registry(mut self, registry: Option<UtilizationRegistry>) -> Self {
        self.utilization_registry = registry;
        self
    }

    /// Builds the topology pieces, returning errors if any occur.
    ///
    /// Use this method when you need to handle errors explicitly,
    /// such as in tests or validation code.
    pub async fn build(self) -> Result<TopologyPieces, Vec<String>> {
        Builder::new(
            self.config,
            self.diff,
            self.buffers,
            self.extra_context,
            self.utilization_registry,
        )
        .build()
        .await
    }

    /// Builds the topology pieces, logging any errors that occur.
    ///
    /// Use this method for runtime configuration loading where
    /// errors should be logged and execution should continue.
    pub async fn build_or_log_errors(self) -> Option<TopologyPieces> {
        match self.build().await {
            Err(errors) => {
                for error in errors {
                    error!(message = "Configuration error.", %error, internal_log_rate_limit = false);
                }
                None
            }
            Ok(new_pieces) => Some(new_pieces),
        }
    }
}

impl TopologyPieces {
    pub async fn build_or_log_errors(
        config: &Config,
        diff: &ConfigDiff,
        buffers: HashMap<ComponentKey, BuiltBuffer>,
        extra_context: ExtraContext,
        utilization_registry: Option<UtilizationRegistry>,
    ) -> Option<Self> {
        TopologyPiecesBuilder::new(config, diff)
            .with_buffers(buffers)
            .with_extra_context(extra_context)
            .with_utilization_registry(utilization_registry)
            .build_or_log_errors()
            .await
    }

    /// Builds only the new pieces, and doesn't check their topology.
    pub async fn build(
        config: &super::Config,
        diff: &ConfigDiff,
        buffers: HashMap<ComponentKey, BuiltBuffer>,
        extra_context: ExtraContext,
        utilization_registry: Option<UtilizationRegistry>,
    ) -> Result<Self, Vec<String>> {
        TopologyPiecesBuilder::new(config, diff)
            .with_buffers(buffers)
            .with_extra_context(extra_context)
            .with_utilization_registry(utilization_registry)
            .build()
            .await
    }
}

const fn filter_events_type(events: &EventArray, data_type: DataType) -> bool {
    match events {
        EventArray::Logs(_) => data_type.contains(DataType::Log),
        EventArray::Metrics(_) => data_type.contains(DataType::Metric),
        EventArray::Traces(_) => data_type.contains(DataType::Trace),
    }
}

#[derive(Debug, Clone)]
struct TransformNode {
    key: ComponentKey,
    typetag: &'static str,
    inputs: Inputs<OutputId>,
    input_details: Input,
    outputs: Vec<TransformOutput>,
    enable_concurrency: bool,
}

impl TransformNode {
    pub fn from_parts(
        key: ComponentKey,
        context: &TransformContext,
        transform: &TransformOuter<OutputId>,
        schema_definition: &[(OutputId, Definition)],
    ) -> Self {
        Self {
            key,
            typetag: transform.inner.get_component_name(),
            inputs: transform.inputs.clone(),
            input_details: transform.inner.input(),
            outputs: transform.inner.outputs(context, schema_definition),
            enable_concurrency: transform.inner.enable_concurrency(),
        }
    }
}

struct Runner {
    transform: Box<dyn SyncTransform>,
    input_rx: Option<BufferReceiver<EventArray>>,
    input_type: DataType,
    outputs: TransformOutputs,
    timer_tx: UtilizationComponentSender,
    latency_recorder: LatencyRecorder,
    events_received: Registered<EventsReceived>,
}

impl Runner {
    fn new(
        transform: Box<dyn SyncTransform>,
        input_rx: BufferReceiver<EventArray>,
        timer_tx: UtilizationComponentSender,
        input_type: DataType,
        outputs: TransformOutputs,
        latency_recorder: LatencyRecorder,
    ) -> Self {
        Self {
            transform,
            input_rx: Some(input_rx),
            input_type,
            outputs,
            timer_tx,
            latency_recorder,
            events_received: register!(EventsReceived),
        }
    }

    fn on_events_received(&mut self, events: &EventArray) {
        self.timer_tx.try_send_stop_wait();

        self.events_received.emit(CountByteSize(
            events.len(),
            events.estimated_json_encoded_size_of(),
        ));
    }

    async fn send_outputs(&mut self, outputs_buf: &mut TransformOutputsBuf) -> crate::Result<()> {
        self.timer_tx.try_send_start_wait();
        let now = Instant::now();
        outputs_buf.for_each_array_mut(|array| self.latency_recorder.on_send(array, now));
        self.outputs.send(outputs_buf).await
    }

    async fn run_inline(mut self) -> TaskResult {
        // 128 is an arbitrary, smallish constant
        const INLINE_BATCH_SIZE: usize = 128;

        let mut outputs_buf = self.outputs.new_buf_with_capacity(INLINE_BATCH_SIZE);

        let mut input_rx = self
            .input_rx
            .take()
            .expect("can't run runner twice")
            .into_stream()
            .filter(move |events| ready(filter_events_type(events, self.input_type)));

        self.timer_tx.try_send_start_wait();
        while let Some(events) = input_rx.next().await {
            self.on_events_received(&events);
            self.transform.transform_all(events, &mut outputs_buf);
            self.send_outputs(&mut outputs_buf)
                .await
                .map_err(TaskError::wrapped)?;
        }

        Ok(TaskOutput::Transform)
    }

    async fn run_concurrently(mut self) -> TaskResult {
        let input_rx = self
            .input_rx
            .take()
            .expect("can't run runner twice")
            .into_stream()
            .filter(move |events| ready(filter_events_type(events, self.input_type)));

        let mut input_rx =
            super::ready_arrays::ReadyArrays::with_capacity(input_rx, READY_ARRAY_CAPACITY);

        let mut in_flight = FuturesOrdered::new();
        let mut shutting_down = false;

        self.timer_tx.try_send_start_wait();
        loop {
            tokio::select! {
                biased;

                result = in_flight.next(), if !in_flight.is_empty() => {
                    match result {
                        Some(Ok(mut outputs_buf)) => {
                            self.send_outputs(&mut outputs_buf).await
                                .map_err(TaskError::wrapped)?;
                        }
                        _ => unreachable!("join error or bad poll"),
                    }
                }

                input_arrays = input_rx.next(), if in_flight.len() < *TRANSFORM_CONCURRENCY_LIMIT && !shutting_down => {
                    match input_arrays {
                        Some(input_arrays) => {
                            let mut len = 0;
                            for events in &input_arrays {
                                self.on_events_received(events);
                                len += events.len();
                            }

                            let mut t = self.transform.clone();
                            let mut outputs_buf = self.outputs.new_buf_with_capacity(len);
                            let task = tokio::spawn(async move {
                                for events in input_arrays {
                                    t.transform_all(events, &mut outputs_buf);
                                }
                                outputs_buf
                            }.in_current_span());
                            in_flight.push_back(task);
                        }
                        None => {
                            shutting_down = true;
                            continue
                        }
                    }
                }

                else => {
                    if shutting_down {
                        break
                    }
                }
            }
        }

        Ok(TaskOutput::Transform)
    }
}