saluki_core/accounting/

registry.rs

use std::marker::PhantomData;
use std::{
    collections::HashMap,
    sync::{Arc, Mutex, MutexGuard},
};

use saluki_common::resource_tracking::{ResourceGroupRegistry, ResourceGroupToken};

use super::{
    api::ResourceAPIHandler, BoundsVerifier, ComponentBounds, MemoryBounds, MemoryGrant, UsageExpr, VerifiedBounds,
    VerifierError,
};
use crate::support::SubsystemIdentifier;

pub(crate) struct ComponentMetadata {
    full_name: Option<String>,
    bounds: ComponentBounds,
    token: Option<ResourceGroupToken>,
    subcomponents: HashMap<String, Arc<Mutex<ComponentMetadata>>>,
}

impl ComponentMetadata {
    fn from_full_name(full_name: Option<String>) -> Self {
        Self {
            full_name,
            bounds: ComponentBounds::default(),
            token: None,
            subcomponents: HashMap::new(),
        }
    }

    pub fn get_or_create<S>(&mut self, name: S) -> Arc<Mutex<Self>>
    where
        S: AsRef<str>,
    {
        // Split the name into the current level name and the remaining name.
        //
        // This lets us handle names which refer to a target nested component instead of having to chain a ton of calls
        // together.
        let name = name.as_ref();
        let (current_level_name, remaining_name) = match name.split_once('.') {
            Some((current_level_name, remaining_name)) => (current_level_name, Some(remaining_name)),
            None => (name, None),
        };

        // Now we need to see if we have an existing component here or if we need to create a new one.
        match self.subcomponents.get(current_level_name) {
            Some(existing) => match remaining_name {
                Some(remaining_name) => {
                    // We found an intermediate subcomponent, so keep recursing.
                    existing.lock().unwrap().get_or_create(remaining_name)
                }
                None => {
                    // We've found the leaf subcomponent.
                    Arc::clone(existing)
                }
            },
            None => {
                // We couldn't find the component at this level, so we need to create it.
                //
                // We do all of our name calculation and so on, but we also leave the token empty for now. We do this to
                // avoid registering intermediate components that aren't actually used by the code, but are simply a
                // consequence of wanting to having a nicely nested structure.
                //
                // We'll register a token for the component the first time it's requested.
                let full_name = match self.full_name.as_ref() {
                    Some(parent_full_name) => format!("{}.{}", parent_full_name, current_level_name),
                    None => current_level_name.to_string(),
                };

                let inner = self
                    .subcomponents
                    .entry(current_level_name.to_string())
                    .or_insert_with(|| Arc::new(Mutex::new(Self::from_full_name(Some(full_name)))));

                // If we still need to recurse further, do so here.. otherwise, return the subcomponent we just created
                // as-is.
                match remaining_name {
                    Some(remaining_name) => inner.lock().unwrap().get_or_create(remaining_name),
                    None => Arc::clone(inner),
                }
            }
        }
    }

    fn token(&mut self) -> ResourceGroupToken {
        match self.token {
            Some(token) => token,
            None => match self.full_name.as_deref() {
                Some(full_name) => {
                    let allocator_component_registry = ResourceGroupRegistry::global();
                    let token = allocator_component_registry.register_resource_group(full_name);
                    self.token = Some(token);

                    token
                }
                None => ResourceGroupToken::root(),
            },
        }
    }

    fn reset(&mut self) {
        self.bounds = ComponentBounds::default();
        self.subcomponents.clear();
    }

    pub fn self_bounds(&self) -> &ComponentBounds {
        &self.bounds
    }

    pub fn as_bounds(&self) -> ComponentBounds {
        let mut bounds = ComponentBounds::default();
        bounds.self_firm_limit_bytes = self.bounds.self_firm_limit_bytes.clone();
        bounds.self_minimum_required_bytes = self.bounds.self_minimum_required_bytes.clone();

        for (name, subcomponent) in self.subcomponents.iter() {
            let subcomponent = subcomponent.lock().unwrap();
            let subcomponent_bounds = subcomponent.as_bounds();
            bounds.subcomponents.insert(name.clone(), subcomponent_bounds);
        }

        bounds
    }
}

/// A registry for components for tracking memory bounds and runtime memory usage.
///
/// This registry provides a unified interface for declaring the memory bounds of a _component_, as well as registering
/// that component for runtime memory usage tracking when using the tracking allocator implementation in `saluki-common`.
///
/// ## Components
///
/// **Components** are any logical grouping of memory usage within a program, and they can be arbitrarily nested.
///
/// For example, a data plane will generally have a topology that defines the components used to accept, process, and
/// forward data. The topology itself could be considered a component, and each individual source, transform, and
/// destination within it could be subcomponents of the topology.
///
/// Components are generally meant to be tied to something that has its own memory bounds and is somewhat standalone,
/// but this isn't an absolute requirement and components can be nested more granularly for organizational/aesthetic
/// purposes. Again, for example, one might opt to create a component in their topology for each component type --
/// sources, transforms, and destinations -- and then add the actual instances of those components as subcomponents to
/// each grouping, leading to a nested structure such as `topology/sources/source1`, `topology/transforms/transform1`,
/// and so on.
///
/// ## Bounds
///
/// Every component is able to define memory bounds for itself and its subcomponents. A builder-style API is exposed to
/// allow for ergonomically defining these bounds -- both minimum and firm -- for components, as well as extending the
/// nestable aspect of the registry itself to the bounds builder, allowing for flexibility in where components are
/// defined from and how they're nested.
///
/// ## Allocation tracking
///
/// Every component is also able to be registered with its own resource group when using the tracking allocator
/// implementation. This is done on demand when the component's token is requested, which avoids polluting the tracking
/// allocator with components that are never actually used, such as those used for organizational/aesthetic purposes.
#[derive(Clone)]
pub struct ComponentRegistry {
    root: Arc<Mutex<ComponentMetadata>>,
}

impl ComponentRegistry {
    /// Creates a handle to this registry.
    ///
    /// The handle provides read-only access to root-level operations like creating API handlers and verifying bounds. It
    /// can be freely cloned and shared.
    pub fn root(&self) -> ComponentRegistryHandle {
        ComponentRegistryHandle {
            inner: Arc::clone(&self.root),
        }
    }

    /// Gets a bounds builder for the component identified by `id`, creating the component -- and any intermediate
    /// components -- if it doesn't already exist.
    ///
    /// `id` is a fully qualified, absolute identifier: the component is addressed directly from the root, so callers
    /// never need to track how deeply nested a registry is.
    pub fn bounds_builder(&self, id: &SubsystemIdentifier) -> MemoryBoundsBuilder<'_> {
        let node = self.root.lock().unwrap().get_or_create(id.to_string());
        MemoryBoundsBuilder { node, _lt: PhantomData }
    }

    /// Gets the resource group tracking token for the component identified by `id`, creating the component -- and any
    /// intermediate components -- if it doesn't already exist.
    ///
    /// The component's resource group is registered (using its full name) the first time its token is requested. `id`
    /// is a fully qualified, absolute identifier, addressed directly from the root.
    pub fn get_resource_group_token(&self, id: &SubsystemIdentifier) -> ResourceGroupToken {
        let node = self.root.lock().unwrap().get_or_create(id.to_string());
        let mut node = node.lock().unwrap();
        node.token()
    }

    /// Gets the memory bounds for all components in the registry.
    pub fn as_bounds(&self) -> ComponentBounds {
        self.root.lock().unwrap().as_bounds()
    }
}

/// A cloneable, read-only handle to a component registry.
///
/// This handle provides access to read-only operations such as creating an API handler or verifying bounds. Unlike
/// [`ComponentRegistry`], it can be freely cloned and shared across ownership boundaries.
///
/// Obtained via [`ComponentRegistry::root`].
#[derive(Clone)]
pub struct ComponentRegistryHandle {
    inner: Arc<Mutex<ComponentMetadata>>,
}

impl ComponentRegistryHandle {
    /// Validates that all components are able to respect the calculated effective limit.
    ///
    /// If validation succeeds, `VerifiedBounds` is returned, which provides information about the effective limit that
    /// can be used for allocating memory.
    ///
    /// ## Errors
    ///
    /// A number of invalid conditions are checked and will cause an error to be returned:
    ///
    /// - when a component has invalid bounds (for example, minimum required bytes higher than firm limit)
    /// - when the combined total of the firm limit for all components exceeds the effective limit
    pub fn verify_bounds(&self, initial_grant: MemoryGrant) -> Result<VerifiedBounds, VerifierError> {
        let bounds = self.inner.lock().unwrap().as_bounds();
        BoundsVerifier::new(initial_grant, bounds).verify()
    }

    /// Gets an API handler for reporting the memory bounds and resource usage of all component in the registry.
    ///
    /// This handler exposes routes for querying the memory bounds and usage of all registered components. See
    /// [`ResourceAPIHandler`] for more information about routes and responses.
    pub fn api_handler(&self) -> ResourceAPIHandler {
        ResourceAPIHandler::from_state(Arc::clone(&self.inner))
    }

    /// Returns a JSON snapshot of the current memory bounds and live usage for all registered components.
    ///
    /// Each component appears as a key in the returned JSON object with `minimum_required_bytes`,
    /// `firm_limit_bytes`, and `actual_live_bytes` fields. This produces the same data as the `/memory/status`
    /// HTTP endpoint but is collected directly from the in-process registry for use outside of the HTTP handler
    /// path (for example, when building a diagnostic artifact).
    pub fn memory_snapshot_json(&self) -> String {
        use std::collections::BTreeMap;

        use saluki_common::resource_tracking::{ResourceGroupRegistry, ResourceStatsSnapshot};

        #[derive(serde::Serialize)]
        struct ComponentUsage {
            minimum_required_bytes: usize,
            firm_limit_bytes: usize,
            actual_live_bytes: usize,
        }

        let empty_snapshot = ResourceStatsSnapshot::empty();
        let mut component_usage: BTreeMap<String, ComponentUsage> = BTreeMap::new();
        let mut inner = self.inner.lock().unwrap();

        ResourceGroupRegistry::global().visit_resource_groups(|component_name, component_stats| {
            let component_meta = inner.get_or_create(component_name);
            let component_meta = component_meta.lock().unwrap();
            let bounds = component_meta.self_bounds();
            let stats_snapshot = component_stats.snapshot_delta(&empty_snapshot);

            component_usage.insert(
                component_name.to_string(),
                ComponentUsage {
                    minimum_required_bytes: bounds.total_minimum_required_bytes(),
                    firm_limit_bytes: bounds.total_firm_limit_bytes(),
                    actual_live_bytes: stats_snapshot.live_bytes(),
                },
            );
        });

        serde_json::to_string_pretty(&component_usage).unwrap_or_else(|e| format!("{{\"error\": \"{e}\"}}"))
    }

    /// Gets the total minimum required bytes for all components in the registry.
    ///
    /// See [`ComponentRegistry::as_bounds`] for more details.
    pub fn as_bounds(&self) -> ComponentBounds {
        self.inner.lock().unwrap().as_bounds()
    }
}

impl Default for ComponentRegistry {
    fn default() -> Self {
        Self {
            root: Arc::new(Mutex::new(ComponentMetadata::from_full_name(None))),
        }
    }
}

pub struct Minimum;
pub struct Firm;

pub(crate) mod private {
    pub trait Sealed {}

    impl Sealed for super::Minimum {}
    impl Sealed for super::Firm {}
}

// Simple trait-based builder state approach so we can use a single builder view to modify either the minimum required
// or firm limit amounts.
pub trait BoundsMutator: private::Sealed {
    fn add_usage(bounds: &mut ComponentBounds, expr: UsageExpr);
}

impl BoundsMutator for Minimum {
    fn add_usage(bounds: &mut ComponentBounds, expr: UsageExpr) {
        bounds.self_minimum_required_bytes.push(expr)
    }
}

impl BoundsMutator for Firm {
    fn add_usage(bounds: &mut ComponentBounds, expr: UsageExpr) {
        bounds.self_firm_limit_bytes.push(expr)
    }
}

/// Builder for defining the memory bounds of a component and its subcomponents.
///
/// This builder provides a simple interface for defining the minimum and firm bounds of a component, as well as
/// declaring subcomponents. For example, a topology can contain its own "self" memory bounds, and then define the
/// individual bounds for each component in the topology.
pub struct MemoryBoundsBuilder<'a> {
    node: Arc<Mutex<ComponentMetadata>>,
    _lt: PhantomData<&'a ()>,
}

impl MemoryBoundsBuilder<'static> {
    #[cfg(test)]
    pub(crate) fn for_test() -> Self {
        Self {
            node: Arc::new(Mutex::new(ComponentMetadata::from_full_name(None))),
            _lt: PhantomData,
        }
    }
}

impl MemoryBoundsBuilder<'_> {
    /// Resets the bounds of the current component to a default state.
    ///
    /// This can be used in scenarios where the bounds of a component need to be redefined after they have been
    /// specified, as not all components are able to be defined in a single pass.
    pub fn reset(&mut self) {
        let mut node = self.node.lock().unwrap();
        node.reset();
    }

    /// Gets a builder object for defining the minimum bounds of the current component.
    pub fn minimum(&mut self) -> BoundsBuilder<'_, Minimum> {
        let bounds = self.node.lock().unwrap();
        BoundsBuilder::<'_, Minimum>::new(bounds)
    }

    /// Gets a builder object for defining the firm bounds of the current component.
    ///
    /// The firm limit is additive with the minimum required memory, so entries that are added via `minimum` don't need
    /// to be added again here.
    pub fn firm(&mut self) -> BoundsBuilder<'_, Firm> {
        let bounds = self.node.lock().unwrap();
        BoundsBuilder::<'_, Firm>::new(bounds)
    }

    /// Creates a nested subcomponent and gets a builder object for it.
    ///
    /// This allows for defining the bounds of various subcomponents within a larger component, which are then rolled up
    /// into the calculated bounds for the parent component.
    ///
    /// `name` is a relative label (single segment, or a dotted path for further nesting) that is sanitized and appended
    /// beneath the current component.
    pub fn subcomponent<S>(&mut self, name: S) -> MemoryBoundsBuilder<'_>
    where
        S: AsRef<str>,
    {
        let node = self
            .node
            .lock()
            .unwrap()
            .get_or_create(SubsystemIdentifier::from_dotted(name.as_ref()).to_string());
        MemoryBoundsBuilder { node, _lt: PhantomData }
    }

    /// Creates a nested subcomponent based on the given component.
    ///
    /// This allows for defining a subcomponent whose bounds come from an object that implements `MemoryBounds` directly.
    pub fn with_subcomponent<S, C>(&mut self, name: S, component: &C) -> &mut Self
    where
        S: AsRef<str>,
        C: MemoryBounds,
    {
        let mut builder = self.subcomponent(name);
        component.specify_bounds(&mut builder);

        self
    }

    #[cfg(test)]
    pub(crate) fn as_bounds(&self) -> ComponentBounds {
        self.node.lock().unwrap().as_bounds()
    }
}

/// Bounds builder.
///
/// Helper type for defining the bounds of a component in a field-driven manner.
pub struct BoundsBuilder<'a, S> {
    inner: MutexGuard<'a, ComponentMetadata>,
    _state: PhantomData<S>,
}

impl<'a, S: BoundsMutator> BoundsBuilder<'a, S> {
    fn new(inner: MutexGuard<'a, ComponentMetadata>) -> Self {
        Self {
            inner,
            _state: PhantomData,
        }
    }

    /// Accounts for the in-memory size of a single value.
    ///
    /// This is useful for tracking the expected memory usage of a single instance of a type if that type is heap
    /// allocated. For example, components that are spawned by a topology generally end up being boxed, which means a
    /// heap allocation exists that's the size of the component type.
    pub fn with_single_value<T>(&mut self, name: impl Into<String>) -> &mut Self {
        S::add_usage(&mut self.inner.bounds, UsageExpr::struct_size::<T>(name));
        self
    }

    /// Accounts for a fixed amount of memory usage.
    ///
    /// This is a catch-all for directly accounting for a specific number of bytes.
    pub fn with_fixed_amount(&mut self, name: impl Into<String>, chunk_size: usize) -> &mut Self {
        S::add_usage(&mut self.inner.bounds, UsageExpr::constant(name, chunk_size));
        self
    }

    /// Accounts for an item container of the given length.
    ///
    /// This can be used to track the expected memory usage of generalized containers like `Vec<T>`, where items are
    /// homogeneous and allocated contiguously.
    pub fn with_array<T>(&mut self, name: impl Into<String>, len: usize) -> &mut Self {
        S::add_usage(
            &mut self.inner.bounds,
            UsageExpr::product(
                "array",
                UsageExpr::struct_size::<T>(name),
                UsageExpr::constant("len", len),
            ),
        );
        self
    }

    /// Accounts for a map container of the given length.
    ///
    /// This can be used to track the expected memory usage of generalized maps like `HashMap<K, V>`, where keys and
    /// values are
    pub fn with_map<K, V>(&mut self, name: impl Into<String>, len: usize) -> &mut Self {
        S::add_usage(
            &mut self.inner.bounds,
            UsageExpr::product(
                "map",
                UsageExpr::sum(
                    name,
                    UsageExpr::struct_size::<K>("key"),
                    UsageExpr::struct_size::<V>("value"),
                ),
                UsageExpr::constant("len", len),
            ),
        );
        self
    }

    pub fn with_expr(&mut self, expr: UsageExpr) -> &mut Self {
        S::add_usage(&mut self.inner.bounds, expr);
        self
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn id(dotted: &str) -> SubsystemIdentifier {
        SubsystemIdentifier::from_dotted(dotted)
    }

    #[test]
    fn bounds_builder_records_bounds_under_absolute_id() {
        let registry = ComponentRegistry::default();
        registry
            .bounds_builder(&id("topology.primary.sources.dsd_in"))
            .firm()
            .with_fixed_amount("buffer", 128);

        assert_eq!(registry.as_bounds().total_firm_limit_bytes(), 128);
    }

    #[test]
    fn repeated_bounds_builder_for_same_id_targets_the_same_node() {
        // Two `bounds_builder` calls with the same identifier resolve to the same component node, so their bounds
        // accumulate rather than replace.
        let registry = ComponentRegistry::default();
        let component = id("topology.primary.transforms.enrich");

        registry
            .bounds_builder(&component)
            .firm()
            .with_fixed_amount("first", 100);
        registry
            .bounds_builder(&component)
            .firm()
            .with_fixed_amount("second", 200);

        assert_eq!(registry.as_bounds().total_firm_limit_bytes(), 300);
    }

    #[test]
    fn distinct_absolute_ids_roll_up_to_the_root() {
        // Distinct absolute identifiers that share a prefix create a nested structure whose bounds roll up.
        let registry = ComponentRegistry::default();
        registry
            .bounds_builder(&id("topology.primary.sources.in"))
            .firm()
            .with_fixed_amount("a", 100);
        registry
            .bounds_builder(&id("topology.primary.destinations.out"))
            .firm()
            .with_fixed_amount("b", 200);

        assert_eq!(registry.as_bounds().total_firm_limit_bytes(), 300);
    }

    #[test]
    fn relative_subcomponents_roll_up_under_their_component() {
        // The bounds builder's relative `subcomponent` nests beneath the component addressed by its absolute id.
        let registry = ComponentRegistry::default();
        {
            let mut builder = registry.bounds_builder(&id("env_provider.workload.remote_agent"));
            builder.firm().with_fixed_amount("self", 10);
            builder
                .subcomponent("string_interner")
                .firm()
                .with_fixed_amount("interner", 90);
        }

        assert_eq!(registry.as_bounds().total_firm_limit_bytes(), 100);
    }

    #[test]
    fn resource_group_token_and_bounds_builder_share_the_same_node() {
        // Taking a resource group token creates the component node; declaring bounds for the same identifier targets
        // that same node, and the read-only handle observes the result.
        let registry = ComponentRegistry::default();
        let component = id("topology.primary.sources.dsd_in");

        let _token = registry.get_resource_group_token(&component);
        registry
            .bounds_builder(&component)
            .minimum()
            .with_fixed_amount("buffer", 64);

        assert_eq!(registry.root().as_bounds().total_minimum_required_bytes(), 64);
    }

    #[test]
    fn with_array_accounts_for_element_size_times_length() {
        // `with_array` builds a `product(struct_size::<T>, len)` expression, so a 10-element array of `u64` should
        // account for exactly `size_of::<u64>() * 10` bytes.
        let mut builder = MemoryBoundsBuilder::for_test();
        builder.firm().with_array::<u64>("array", 10);

        let bounds = builder.as_bounds();
        assert_eq!(bounds.total_firm_limit_bytes(), std::mem::size_of::<u64>() * 10);
        assert_eq!(bounds.total_minimum_required_bytes(), 0);
    }

    #[test]
    fn with_map_accounts_for_entry_size_times_length() {
        // `with_map` builds a `product(sum(struct_size::<K>, struct_size::<V>), len)` expression, so a 5-entry map
        // of `u32 -> u64` should account for exactly `(size_of::<u32>() + size_of::<u64>()) * 5` bytes.
        let mut builder = MemoryBoundsBuilder::for_test();
        builder.firm().with_map::<u32, u64>("map", 5);

        let bounds = builder.as_bounds();
        let expected = (std::mem::size_of::<u32>() + std::mem::size_of::<u64>()) * 5;
        assert_eq!(bounds.total_firm_limit_bytes(), expected);
        assert_eq!(bounds.total_minimum_required_bytes(), 0);
    }

    #[test]
    fn minimum_bounds_are_included_in_the_firm_total() {
        // The firm limit is documented as additive with the minimum required bytes, so an array added to `minimum`
        // and a map added to `firm` should both roll up into `total_firm_limit_bytes`.
        let mut builder = MemoryBoundsBuilder::for_test();
        builder.minimum().with_array::<u64>("array", 10);
        builder.firm().with_map::<u32, u64>("map", 5);

        let bounds = builder.as_bounds();
        let array_bytes = std::mem::size_of::<u64>() * 10;
        let map_bytes = (std::mem::size_of::<u32>() + std::mem::size_of::<u64>()) * 5;

        assert_eq!(bounds.total_minimum_required_bytes(), array_bytes);
        assert_eq!(bounds.total_firm_limit_bytes(), array_bytes + map_bytes);
    }
}