kinetic/topology/

running.rs

//! Management of a running topology, including graceful shutdown and hot-reload.

use crate::topology::{Fanout, TopologyBuilder};
use kinetic_config::model::Config;
use kinetic_core::{OutputId, ShutdownCoordinator, ShutdownReason};
use std::collections::HashMap;
use tokio::task::JoinHandle;
use tracing::info;

/// A handle to a running component task.
pub struct ComponentTask {
    pub coordinator: ShutdownCoordinator,
    pub handle: JoinHandle<()>,
}

/// A handle to a running topology.
pub struct RunningTopology {
    config: Config,
    data_dir: std::path::PathBuf,
    pub sources: HashMap<String, ComponentTask>,
    pub transforms: HashMap<String, ComponentTask>,
    pub sinks: HashMap<String, ComponentTask>,
    pub fanouts: HashMap<OutputId, Fanout>,
}

impl RunningTopology {
    /// Create a new `RunningTopology` from a set of spawned tasks.
    pub fn new(
        config: Config,
        data_dir: std::path::PathBuf,
        sources: HashMap<String, ComponentTask>,
        transforms: HashMap<String, ComponentTask>,
        sinks: HashMap<String, ComponentTask>,
        fanouts: HashMap<OutputId, Fanout>,
    ) -> Self {
        Self {
            config,
            data_dir,
            sources,
            transforms,
            sinks,
            fanouts,
        }
    }

    /// Triggers a graceful shutdown and waits for all components to drain.
    pub async fn stop(mut self) {
        info!("Initiating graceful shutdown of all components...");

        // 1. Broadcast shutdown to all components first.
        // This ensures everyone knows to stop accepting new work and start draining.
        for task in self.sources.values() {
            task.coordinator.broadcast(ShutdownReason::Graceful);
        }
        for task in self.transforms.values() {
            task.coordinator.broadcast(ShutdownReason::Graceful);
        }
        for task in self.sinks.values() {
            task.coordinator.broadcast(ShutdownReason::Graceful);
        }

        // 2. Await them in order (Sources -> Transforms -> Sinks) to ensure data flows through.
        for (name, task) in self.sources.drain() {
            let _ = task.handle.await;
            info!("Source '{}' stopped", name);
        }

        for (name, task) in self.transforms.drain() {
            let _ = task.handle.await;
            info!("Transform '{}' stopped", name);
        }

        for (name, task) in self.sinks.drain() {
            let _ = task.handle.await;
            info!("Sink '{}' stopped", name);
        }

        info!("Shutdown complete.");
    }

    /// Reloads the topology with a new configuration using granular diffing.
    pub async fn reload(&mut self, new_config: Config) -> anyhow::Result<()> {
        let diff = self.config.diff(&new_config);
        if diff.is_empty() {
            info!("No changes detected in configuration reload.");
            return Ok(());
        }

        info!(
            "Reloading topology with granular changes: added={}, removed={}, changed={}",
            diff.sources_added.len() + diff.transforms_added.len() + diff.sinks_added.len(),
            diff.sources_removed.len() + diff.transforms_removed.len() + diff.sinks_removed.len(),
            diff.sources_changed.len() + diff.transforms_changed.len() + diff.sinks_changed.len()
        );

        let builder = TopologyBuilder::new(new_config.clone(), self.data_dir.clone());
        let new_context = builder.build().map_err(|e| anyhow::anyhow!(e))?;

        // 1. Update existing Fanouts with new destinations.
        // This ensures that existing components (which are still running and have a handle to their Fanout)
        // can see the new downstream destinations that were wired in the new context.
        for (output_id, new_fanout) in &new_context.fanouts {
            if let Some(existing_fanout) = self.fanouts.get(output_id) {
                existing_fanout.replace_senders(new_fanout.get_senders());
            } else {
                // This is a new output from an existing component (shouldn't happen for sources, maybe for transforms)
                self.fanouts.insert(output_id.clone(), new_fanout.clone());
            }
        }

        // 2. Stop removed and changed components
        for name in diff.sources_removed.iter().chain(&diff.sources_changed) {
            if let Some(task) = self.sources.remove(name) {
                task.coordinator.broadcast(ShutdownReason::Reload);
                let _ = task.handle.await;
            }
        }
        for name in diff
            .transforms_removed
            .iter()
            .chain(&diff.transforms_changed)
        {
            if let Some(task) = self.transforms.remove(name) {
                task.coordinator.broadcast(ShutdownReason::Reload);
                let _ = task.handle.await;
            }
        }
        for name in diff.sinks_removed.iter().chain(&diff.sinks_changed) {
            if let Some(task) = self.sinks.remove(name) {
                task.coordinator.broadcast(ShutdownReason::Reload);
                let _ = task.handle.await;
            }
        }

        // 3. Spawn added and changed components
        let new_running = builder.spawn_with_context(new_context, Some(&diff)).await?;

        // Merge the new tasks into the existing ones
        for (name, task) in new_running.sources {
            if diff.sources_added.contains(&name) || diff.sources_changed.contains(&name) {
                self.sources.insert(name, task);
            }
        }
        for (name, task) in new_running.transforms {
            if diff.transforms_added.contains(&name) || diff.transforms_changed.contains(&name) {
                self.transforms.insert(name, task);
            }
        }
        for (name, task) in new_running.sinks {
            if diff.sinks_added.contains(&name) || diff.sinks_changed.contains(&name) {
                self.sinks.insert(name, task);
            }
        }

        // Update the set of fanouts we track
        self.fanouts = new_running.fanouts;

        self.config = new_config;
        info!("Granular topology reload completed.");
        Ok(())
    }
}