fv-compute 0.2.0

//! One object to run transforms with.
//!
//! A [`Runtime`] bundles the three things every host otherwise wires by hand: a [`Registry`] over
//! the transform directories you point it at, the [`Backends`] you register, and a cache of loaded
//! transforms so `load` happens once per selector and `run` is the hot path. Build one at startup
//! and share it (`Arc<Runtime>`) across threads or partitions.
//!
//! ```no_run
//! use fv_compute::Runtime;
//! # use fv_compute::{Compute, ComputeError, ImplKind, TransformBackend, TransformManifest};
//! # struct Noop;
//! # impl TransformBackend for Noop {
//! #     fn kind(&self) -> ImplKind { ImplKind::Wasm }
//! #     fn load(&self, _: &TransformManifest, _: &std::path::Path) -> Result<Box<dyn Compute>, ComputeError> { unimplemented!() }
//! # }
//! # fn backend() -> Noop { Noop }
//! # fn batch() -> arrow::array::RecordBatch { unimplemented!() }
//!
//! let runtime = Runtime::builder()
//!     .root("./transforms")      // every `<dir>/transform.toml` under it
//!     .backend(backend())        // e.g. fv_compute_wasm::WasmBackend::new()?
//!     .build()?;
//!
//! let out = runtime.run("spikeTotal", &[batch()])?;   // "id" or "id@version"
//! # Ok::<(), Box<dyn std::error::Error>>(())
//! ```

use crate::contract::{Compute, ComputeError, TransformBackend};
use crate::registry::{Backends, DirRoot, Registry, RegistryBuilder, RegistryError, Root};
use arrow::array::RecordBatch;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};

/// A registry, its backends, and a cache of loaded transforms. See the [module docs](self).
pub struct Runtime {
    registry: Registry,
    backends: Backends,
    loaded: Mutex<HashMap<String, Arc<dyn Compute>>>,
}

impl Runtime {
    /// Start building a runtime: add roots and backends, then [`RuntimeBuilder::build`].
    pub fn builder() -> RuntimeBuilder {
        RuntimeBuilder::default()
    }

    /// Wrap an already-built registry and backend set.
    pub fn new(registry: Registry, backends: Backends) -> Self {
        Self {
            registry,
            backends,
            loaded: Mutex::new(HashMap::new()),
        }
    }

    /// The transforms this runtime can see.
    pub fn registry(&self) -> &Registry {
        &self.registry
    }

    /// The backends this runtime dispatches to.
    pub fn backends(&self) -> &Backends {
        &self.backends
    }

    /// Resolve `selector` (`id` or `id@version`) and load it through its backend. The loaded
    /// transform is cached by selector, so repeated calls are a map lookup; the expensive work
    /// (compiling a component, resolving an image) happens once.
    pub fn load(&self, selector: &str) -> Result<Arc<dyn Compute>, ComputeError> {
        if let Some(c) = self.loaded.lock().unwrap().get(selector) {
            return Ok(c.clone());
        }
        let compute: Arc<dyn Compute> = Arc::from(self.registry.load(selector, &self.backends)?);
        self.loaded
            .lock()
            .unwrap()
            .insert(selector.to_string(), compute.clone());
        Ok(compute)
    }

    /// Load (cached) and run `selector` over `inputs`, one batch per declared input.
    pub fn run(&self, selector: &str, inputs: &[RecordBatch]) -> Result<RecordBatch, ComputeError> {
        self.load(selector)?.run(inputs)
    }

    /// Load `selector` without running it. Call this when a job is configured, so a missing or
    /// broken transform fails the job up front instead of failing every batch later.
    pub fn ensure_loadable(&self, selector: &str) -> Result<(), ComputeError> {
        self.load(selector).map(drop)
    }

    /// Drop every cached transform. The next `load` of each selector goes through its backend
    /// again — use after the transform directories change on disk.
    pub fn clear_cache(&self) {
        self.loaded.lock().unwrap().clear();
    }
}

impl std::fmt::Debug for Runtime {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Runtime")
            .field("transforms", &self.registry.len())
            .field("backends", &self.backends.kinds().collect::<Vec<_>>())
            .field("loaded", &self.loaded.lock().unwrap().len())
            .finish()
    }
}

/// Builds a [`Runtime`]: roots in precedence order (later roots override earlier ones for the
/// same `id@version`) and the backends to dispatch to.
#[derive(Default)]
pub struct RuntimeBuilder {
    registry: RegistryBuilder,
    backends: Backends,
    roots: usize,
}

impl RuntimeBuilder {
    /// A directory of transforms (`<path>/<name>/transform.toml`), named after its last path
    /// segment. A path that does not exist contributes nothing.
    pub fn root(mut self, path: impl Into<PathBuf>) -> Self {
        let path = path.into();
        let name = path
            .file_name()
            .map(|s| s.to_string_lossy().into_owned())
            .unwrap_or_else(|| format!("root{}", self.roots));
        self.roots += 1;
        self.registry = self.registry.root(DirRoot::new(name, path));
        self
    }

    /// A directory of transforms with an explicit name (shown in diagnostics and the catalog).
    pub fn named_root(mut self, name: impl Into<String>, path: impl Into<PathBuf>) -> Self {
        self.roots += 1;
        self.registry = self.registry.root(DirRoot::new(name, path));
        self
    }

    /// Any other source of transforms — your own [`Root`] implementation.
    pub fn source(mut self, root: impl Root + 'static) -> Self {
        self.registry = self.registry.root(root);
        self
    }

    /// Register a backend. One per [`crate::ImplKind`]; a later registration for the same kind
    /// replaces the earlier one.
    pub fn backend(mut self, backend: impl TransformBackend + 'static) -> Self {
        self.backends = self.backends.register(backend);
        self
    }

    /// Discover and validate every transform under the roots. A malformed manifest anywhere is an
    /// error here, not at run time.
    pub fn build(self) -> Result<Runtime, RegistryError> {
        Ok(Runtime::new(self.registry.build()?, self.backends))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::manifest::{ImplKind, TransformManifest};
    use crate::registry::MemoryRoot;
    use std::path::Path;
    use std::sync::atomic::{AtomicUsize, Ordering};

    const MANIFEST: &str = r#"
        id = "double"
        version = "1.0.0"
        impl = "wasm"
        entry = "main.wasm"
        [[inputs]]
        columns = [{ name = "x", type = "int64" }]
        [output]
        columns = [{ name = "x", type = "int64" }]
    "#;

    /// A backend that counts loads and echoes its input.
    struct Counting(Arc<AtomicUsize>);

    struct Echo(TransformManifest);

    impl Compute for Echo {
        fn manifest(&self) -> &TransformManifest {
            &self.0
        }
        fn run(&self, inputs: &[RecordBatch]) -> Result<RecordBatch, ComputeError> {
            Ok(inputs[0].clone())
        }
    }

    impl TransformBackend for Counting {
        fn kind(&self) -> ImplKind {
            ImplKind::Wasm
        }
        fn load(&self, manifest: &TransformManifest, _root: &Path) -> Result<Box<dyn Compute>, ComputeError> {
            self.0.fetch_add(1, Ordering::SeqCst);
            Ok(Box::new(Echo(manifest.clone())))
        }
    }

    fn batch() -> RecordBatch {
        use arrow::array::Int64Array;
        use arrow::datatypes::{DataType, Field, Schema};
        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, true)]));
        RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![1, 2, 3]))]).unwrap()
    }

    #[test]
    fn loads_once_per_selector_and_runs() {
        let loads = Arc::new(AtomicUsize::new(0));
        let runtime = Runtime::builder()
            .source(MemoryRoot::new("mem").unit("double", MANIFEST))
            .backend(Counting(loads.clone()))
            .build()
            .unwrap();

        runtime.ensure_loadable("double").unwrap();
        let out = runtime.run("double", &[batch()]).unwrap();
        let again = runtime.run("double@1.0.0", &[batch()]).unwrap();
        assert_eq!(out.num_rows(), 3);
        assert_eq!(again.num_rows(), 3);
        // "double" and "double@1.0.0" are distinct selectors: two loads, never a third.
        assert_eq!(loads.load(Ordering::SeqCst), 2);

        runtime.clear_cache();
        runtime.run("double", &[batch()]).unwrap();
        assert_eq!(loads.load(Ordering::SeqCst), 3);
    }

    #[test]
    fn missing_transform_and_missing_backend_are_distinct_errors() {
        let runtime = Runtime::builder()
            .source(MemoryRoot::new("mem").unit("double", MANIFEST))
            .build()
            .unwrap();
        assert!(matches!(runtime.load("nope"), Err(ComputeError::Load { .. })));
        assert!(matches!(
            runtime.load("double"),
            Err(ComputeError::NoBackend { kind: ImplKind::Wasm })
        ));
    }

    #[test]
    fn a_missing_directory_root_is_empty_not_an_error() {
        let runtime = Runtime::builder().root("/definitely/not/here").build().unwrap();
        assert!(runtime.registry().is_empty());
        assert_eq!(
            format!("{runtime:?}"),
            "Runtime { transforms: 0, backends: [], loaded: 0 }"
        );
    }
}