kinetic/topology/

builder.rs

//! Builds and wires the DAG of pipeline components.

use crate::topology::Fanout;
use crate::topology::running::{ComponentTask, RunningTopology};
use kinetic_common::config::{SinkContext, SourceContext, TransformContext};
use kinetic_config::model::Config;
use kinetic_core::{ComponentId, OutputId, ShutdownCoordinator};
use snafu::Snafu;
use std::collections::{HashMap, HashSet};
use tracing::error;

#[derive(Debug, Snafu)]
pub enum Error {
    #[snafu(display("Component '{}' references unknown input '{}'", component, input))]
    UnknownInput { component: String, input: String },

    #[snafu(display("Cycle detected in pipeline topology"))]
    CycleDetected,

    #[snafu(display("Internal error: {}", message))]
    InternalError { message: String },
}

pub type Result<T, E = Error> = std::result::Result<T, E>;

/// A parsed and validated reference to an upstream output.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct InputRef {
    component: String,
    output_name: String,
}

impl InputRef {
    /// Parses an input string like "source_name" or "source_name.logs" into an `InputRef`.
    /// If no output name is specified, defaults to "default".
    fn parse(input: &str) -> Self {
        match input.split_once('.') {
            Some((comp, out)) => Self {
                component: comp.to_string(),
                output_name: out.to_string(),
            },
            None => Self {
                component: input.to_string(),
                output_name: "default".to_string(),
            },
        }
    }
}

/// Context for running topology components
pub struct TopologyContext {
    pub transform_receivers: HashMap<String, kinetic_buffers::BufferReceiver>,
    pub sink_receivers: HashMap<String, kinetic_buffers::BufferReceiver>,
    pub shadow_archive_receiver: Option<kinetic_buffers::BufferReceiver>,
    pub fanouts: HashMap<OutputId, Fanout>,
}

/// The topology builder validates the configuration and wires the channels between components.
pub struct TopologyBuilder {
    config: Config,
    data_dir: std::path::PathBuf,
}

impl TopologyBuilder {
    pub fn new(config: Config, data_dir: std::path::PathBuf) -> Self {
        Self { config, data_dir }
    }

    pub fn build(&self) -> Result<TopologyContext> {
        self.validate_references()?;
        self.detect_cycles()?;

        let mut fanouts: HashMap<OutputId, Fanout> = HashMap::new();
        let mut transform_receivers = HashMap::new();
        let mut sink_receivers = HashMap::new();
        let mut component_senders = HashMap::new();

        // Initialize shadow archive sink if configured.
        // This is an invisible sink that captures raw data from all sources.
        let mut shadow_archive_receiver = None;
        if let Some(_archive_config) = &self.config.archive {
            let id = "shadow_archive_sink".to_string();
            // Larger capacity for archive to minimize backpressure impact
            let (tx, rx) =
                kinetic_buffers::channel(1000, kinetic_buffers::WhenFull::Block, id.clone());
            shadow_archive_receiver = Some(rx);
            component_senders.insert(id, tx);
        }

        // Initialize channels for all transforms and sinks.
        for id in self.config.transforms.keys() {
            let (tx, rx) =
                kinetic_buffers::channel(100, kinetic_buffers::WhenFull::Block, id.clone());
            transform_receivers.insert(id.clone(), rx);
            component_senders.insert(id.clone(), tx);
        }
        for (id, sink) in &self.config.sinks {
            let (tx, rx) = match sink.buffer() {
                kinetic_config::model::BufferConfig::Memory { max_batches } => {
                    kinetic_buffers::channel(
                        *max_batches,
                        kinetic_buffers::WhenFull::Block,
                        id.clone(),
                    )
                }
                kinetic_config::model::BufferConfig::Disk {
                    path,
                    max_size_bytes,
                } => kinetic_buffers::disk::channel(path.clone(), id.clone(), *max_size_bytes)
                    .map_err(|e| Error::InternalError {
                        message: format!("Failed to create disk buffer for sink '{}': {}", id, e),
                    })?,
            };
            sink_receivers.insert(id.clone(), rx);
            component_senders.insert(id.clone(), tx);
        }

        for id in self.config.sources.keys() {
            let fanout = Fanout::new();

            if let Some(shadow_tx) = component_senders.get("shadow_archive_sink") {
                fanout.add_sender(shadow_tx.clone());
            }

            fanouts.insert(
                OutputId {
                    component: ComponentId(id.clone()),
                    output_name: "default".to_string(),
                },
                fanout,
            );
            fanouts.insert(
                OutputId {
                    component: ComponentId(id.clone()),
                    output_name: "error".to_string(),
                },
                Fanout::new(),
            );
        }
        for (id, transform) in &self.config.transforms {
            fanouts.insert(
                OutputId {
                    component: ComponentId(id.clone()),
                    output_name: "default".to_string(),
                },
                Fanout::new(),
            );
            fanouts.insert(
                OutputId {
                    component: ComponentId(id.clone()),
                    output_name: "error".to_string(),
                },
                Fanout::new(),
            );
            for output_name in transform.outputs().keys() {
                fanouts.insert(
                    OutputId {
                        component: ComponentId(id.clone()),
                        output_name: output_name.clone(),
                    },
                    Fanout::new(),
                );
            }
        }

        let inputs = self
            .config
            .transforms
            .iter()
            .map(|(id, t)| (id, t.inputs()))
            .chain(self.config.sinks.iter().map(|(id, s)| (id, s.inputs())));

        for (component_id, inputs) in inputs {
            for input_str in inputs {
                let parsed = InputRef::parse(input_str);
                let output_id = OutputId {
                    component: ComponentId(parsed.component),
                    output_name: parsed.output_name,
                };
                let fanout = fanouts
                    .get_mut(&output_id)
                    .ok_or_else(|| Error::UnknownInput {
                        component: component_id.clone(),
                        input: input_str.clone(),
                    })?;
                let sender = component_senders
                    .get(component_id)
                    .ok_or_else(|| Error::InternalError {
                        message: format!("Sender missing for component: {}", component_id),
                    })?
                    .clone();
                fanout.add_sender(sender);
            }
        }

        Ok(TopologyContext {
            transform_receivers,
            sink_receivers,
            shadow_archive_receiver,
            fanouts,
        })
    }

    pub async fn spawn(
        self,
        diff: Option<&kinetic_config::model::ConfigDiff>,
    ) -> anyhow::Result<RunningTopology> {
        let context = self.build()?;
        self.spawn_with_context(context, diff).await
    }

    pub async fn spawn_with_context(
        self,
        mut context: TopologyContext,
        diff: Option<&kinetic_config::model::ConfigDiff>,
    ) -> anyhow::Result<RunningTopology> {
        // 1. Run healthchecks for NEW or CHANGED sinks and sources
        let mut healthchecks = Vec::new();
        for (id, sink_config) in &self.config.sinks {
            if diff.is_none_or(|d| d.sinks_added.contains(id) || d.sinks_changed.contains(id))
                && let Some(hc) =
                    crate::topology::build_impl::build_sink_healthcheck(sink_config, id.clone())
                        .await?
            {
                healthchecks.push((id.clone(), hc));
            }
        }

        for (id, source_config) in &self.config.sources {
            if diff.is_none_or(|d| d.sources_added.contains(id) || d.sources_changed.contains(id))
                && let Some(hc) = crate::topology::build_impl::build_source_healthcheck(
                    source_config,
                    id.clone(),
                    self.config.pipeline_id.clone(),
                )
                .await?
            {
                healthchecks.push((id.clone(), hc));
            }
        }

        if !healthchecks.is_empty() {
            crate::topology::healthcheck::run_all(healthchecks)
                .await
                .map_err(|e| Error::InternalError {
                    message: format!("Healthchecks failed: {}", e),
                })?;
        }

        let mut sources = HashMap::new();
        let mut transforms = HashMap::new();
        let mut sinks = HashMap::new();

        // 2. Spawn components with error cleanup
        let result = self
            .spawn_inner(
                &mut context,
                diff,
                &mut sources,
                &mut transforms,
                &mut sinks,
            )
            .await;

        if let Err(e) = result {
            error!("Failed to spawn topology, cleaning up partial tasks: {}", e);
            let partial = RunningTopology::new(
                self.config.clone(),
                self.data_dir.clone(),
                sources,
                transforms,
                sinks,
                context.fanouts,
            );
            partial.stop().await;
            return Err(e);
        }

        Ok(RunningTopology::new(
            self.config,
            self.data_dir,
            sources,
            transforms,
            sinks,
            context.fanouts,
        ))
    }

    async fn spawn_inner(
        &self,
        context: &mut TopologyContext,
        diff: Option<&kinetic_config::model::ConfigDiff>,
        sources: &mut HashMap<String, ComponentTask>,
        transforms: &mut HashMap<String, ComponentTask>,
        sinks: &mut HashMap<String, ComponentTask>,
    ) -> anyhow::Result<()> {
        // Spawn shadow archive sink if configured
        if let Some(archive_config) = self.config.archive.clone()
            && let Some(receiver) = context.shadow_archive_receiver.take()
        {
            let coordinator = ShutdownCoordinator::new();
            let sink = crate::sinks::shadow_archive::ShadowArchiveSink::new(
                archive_config,
                receiver,
                coordinator.register(),
                self.config.pipeline_id.clone(),
            );
            let handle = sink.run();
            sinks.insert(
                "shadow_archive_sink".to_string(),
                ComponentTask {
                    coordinator,
                    handle,
                },
            );
        }

        for (id, source_config) in &self.config.sources {
            if diff.is_none_or(|d| d.sources_added.contains(id) || d.sources_changed.contains(id)) {
                let output_id = OutputId {
                    component: ComponentId(id.clone()),
                    output_name: "default".to_string(),
                };
                let fanout = context
                    .fanouts
                    .get(&output_id)
                    .ok_or_else(|| anyhow::anyhow!("Default fanout missing for source {}", id))?;
                let error_output_id = OutputId {
                    component: ComponentId(id.clone()),
                    output_name: "error".to_string(),
                };
                let error_fanout = context
                    .fanouts
                    .get(&error_output_id)
                    .ok_or_else(|| anyhow::anyhow!("Error fanout missing for source {}", id))?;

                let coordinator = ShutdownCoordinator::new();
                let cx = SourceContext {
                    id: ComponentId(id.clone()),
                    pipeline_id: self.config.pipeline_id.clone(),
                    out: fanout.clone_to_sender(id.clone()),
                    error_out: error_fanout.clone_to_sender(format!("{}_error", id)),
                    shutdown: coordinator.register(),
                    acknowledgements: Default::default(),
                    data_dir: Some(self.data_dir.join("sources").join(id)),
                };
                let handle = crate::topology::build_impl::build_source(source_config, cx)?;
                sources.insert(
                    id.clone(),
                    ComponentTask {
                        coordinator,
                        handle,
                    },
                );
            }
        }

        for (id, transform_config) in &self.config.transforms {
            if diff.is_none_or(|d| {
                d.transforms_added.contains(id) || d.transforms_changed.contains(id)
            }) {
                let receiver = context
                    .transform_receivers
                    .remove(id)
                    .ok_or_else(|| anyhow::anyhow!("Receiver missing"))?;

                let mut outs = HashMap::new();
                for output_name in transform_config.outputs().keys() {
                    let output_id = OutputId {
                        component: ComponentId(id.clone()),
                        output_name: output_name.clone(),
                    };
                    let fanout = context
                        .fanouts
                        .get(&output_id)
                        .ok_or_else(|| anyhow::anyhow!("Fanout missing"))?;
                    outs.insert(
                        output_name.clone(),
                        fanout.clone_to_sender(format!("{}_{}", id, output_name)),
                    );
                }
                if !outs.contains_key("default") {
                    let output_id = OutputId {
                        component: ComponentId(id.clone()),
                        output_name: "default".to_string(),
                    };
                    if let Some(fanout) = context.fanouts.get(&output_id) {
                        outs.insert(
                            "default".to_string(),
                            fanout.clone_to_sender(format!("{}_default", id)),
                        );
                    }
                }

                let coordinator = ShutdownCoordinator::new();
                let cx = TransformContext {
                    id: ComponentId(id.clone()),
                    pipeline_id: self.config.pipeline_id.clone(),
                    in_rx: receiver,
                    outs,
                    shutdown: coordinator.register(),
                    data_dir: Some(self.data_dir.join("transforms").join(id)),
                };
                let handle = crate::topology::build_impl::build_transform(transform_config, cx)?;
                transforms.insert(
                    id.clone(),
                    ComponentTask {
                        coordinator,
                        handle,
                    },
                );
            }
        }

        for (id, sink_config) in &self.config.sinks {
            if diff.is_none_or(|d| d.sinks_added.contains(id) || d.sinks_changed.contains(id)) {
                let receiver = context
                    .sink_receivers
                    .remove(id)
                    .ok_or_else(|| anyhow::anyhow!("Receiver missing"))?;

                let coordinator = ShutdownCoordinator::new();
                let cx = SinkContext {
                    id: ComponentId(id.clone()),
                    in_rx: receiver,
                    shutdown: coordinator.register(),
                    acknowledgements: Default::default(),
                    data_dir: Some(self.data_dir.join("sinks").join(id)),
                };
                let handle = crate::topology::build_impl::build_sink(sink_config, cx)?;
                sinks.insert(
                    id.clone(),
                    ComponentTask {
                        coordinator,
                        handle,
                    },
                );
            }
        }

        Ok(())
    }

    fn validate_references(&self) -> Result<()> {
        let valid_components: HashSet<&String> = self
            .config
            .sources
            .keys()
            .chain(self.config.transforms.keys())
            .collect();
        for (id, transform) in &self.config.transforms {
            self.check_inputs(id, transform.inputs(), &valid_components)?;
        }
        for (id, sink) in &self.config.sinks {
            self.check_inputs(id, sink.inputs(), &valid_components)?;
        }
        Ok(())
    }

    fn check_inputs(
        &self,
        component_id: &str,
        inputs: &[String],
        valid_components: &HashSet<&String>,
    ) -> Result<()> {
        for input in inputs {
            let parsed = InputRef::parse(input);
            if !valid_components.contains(&parsed.component) {
                return Err(Error::UnknownInput {
                    component: component_id.to_owned(),
                    input: input.clone(),
                });
            }
        }
        Ok(())
    }

    fn detect_cycles(&self) -> Result<()> {
        let mut visited = HashSet::new();
        let mut visiting = HashSet::new();
        let mut graph: HashMap<String, Vec<String>> = HashMap::new();
        for (id, transform) in &self.config.transforms {
            for input in transform.inputs() {
                graph
                    .entry(InputRef::parse(input).component)
                    .or_default()
                    .push(id.clone());
            }
        }
        for (id, sink) in &self.config.sinks {
            for input in sink.inputs() {
                graph
                    .entry(InputRef::parse(input).component)
                    .or_default()
                    .push(id.clone());
            }
        }
        for node in graph.keys() {
            if !visited.contains(node) && self.dfs_cycle(node, &graph, &mut visited, &mut visiting)
            {
                return Err(Error::CycleDetected);
            }
        }
        Ok(())
    }

    fn dfs_cycle(
        &self,
        node: &str,
        graph: &HashMap<String, Vec<String>>,
        visited: &mut HashSet<String>,
        visiting: &mut HashSet<String>,
    ) -> bool {
        visiting.insert(node.to_string());
        if let Some(neighbors) = graph.get(node) {
            for neighbor in neighbors {
                if visiting.contains(neighbor) {
                    return true;
                }
                if !visited.contains(neighbor) && self.dfs_cycle(neighbor, graph, visited, visiting)
                {
                    return true;
                }
            }
        }
        visiting.remove(node);
        visited.insert(node.to_string());
        false
    }
}