Skip to main content

fv_compute/
runtime.rs

1//! One object to run transforms with.
2//!
3//! A [`Runtime`] bundles the three things every host otherwise wires by hand: a [`Registry`] over
4//! the transform directories you point it at, the [`Backends`] you register, and a cache of loaded
5//! transforms so `load` happens once per selector and `run` is the hot path. Build one at startup
6//! and share it (`Arc<Runtime>`) across threads or partitions.
7//!
8//! ```no_run
9//! use fv_compute::Runtime;
10//! # use fv_compute::{Compute, ComputeError, ImplKind, TransformBackend, TransformManifest};
11//! # struct Noop;
12//! # impl TransformBackend for Noop {
13//! #     fn kind(&self) -> ImplKind { ImplKind::Wasm }
14//! #     fn load(&self, _: &TransformManifest, _: &std::path::Path) -> Result<Box<dyn Compute>, ComputeError> { unimplemented!() }
15//! # }
16//! # fn backend() -> Noop { Noop }
17//! # fn batch() -> arrow::array::RecordBatch { unimplemented!() }
18//!
19//! let runtime = Runtime::builder()
20//!     .root("./transforms")      // every `<dir>/transform.toml` under it
21//!     .backend(backend())        // e.g. fv_compute_wasm::WasmBackend::new()?
22//!     .build()?;
23//!
24//! let out = runtime.run("spikeTotal", &[batch()])?;   // "id" or "id@version"
25//! # Ok::<(), Box<dyn std::error::Error>>(())
26//! ```
27
28use crate::contract::{Compute, ComputeError, TransformBackend};
29use crate::registry::{Backends, DirRoot, Registry, RegistryBuilder, RegistryError, Root};
30use arrow::array::RecordBatch;
31use std::collections::HashMap;
32use std::path::PathBuf;
33use std::sync::{Arc, Mutex};
34
35/// A registry, its backends, and a cache of loaded transforms. See the [module docs](self).
36pub struct Runtime {
37    registry: Registry,
38    backends: Backends,
39    loaded: Mutex<HashMap<String, Arc<dyn Compute>>>,
40}
41
42impl Runtime {
43    /// Start building a runtime: add roots and backends, then [`RuntimeBuilder::build`].
44    pub fn builder() -> RuntimeBuilder {
45        RuntimeBuilder::default()
46    }
47
48    /// Wrap an already-built registry and backend set.
49    pub fn new(registry: Registry, backends: Backends) -> Self {
50        Self {
51            registry,
52            backends,
53            loaded: Mutex::new(HashMap::new()),
54        }
55    }
56
57    /// The transforms this runtime can see.
58    pub fn registry(&self) -> &Registry {
59        &self.registry
60    }
61
62    /// The backends this runtime dispatches to.
63    pub fn backends(&self) -> &Backends {
64        &self.backends
65    }
66
67    /// Resolve `selector` (`id` or `id@version`) and load it through its backend. The loaded
68    /// transform is cached by selector, so repeated calls are a map lookup; the expensive work
69    /// (compiling a component, resolving an image) happens once.
70    pub fn load(&self, selector: &str) -> Result<Arc<dyn Compute>, ComputeError> {
71        if let Some(c) = self.loaded.lock().unwrap().get(selector) {
72            return Ok(c.clone());
73        }
74        let compute: Arc<dyn Compute> = Arc::from(self.registry.load(selector, &self.backends)?);
75        self.loaded
76            .lock()
77            .unwrap()
78            .insert(selector.to_string(), compute.clone());
79        Ok(compute)
80    }
81
82    /// Load (cached) and run `selector` over `inputs`, one batch per declared input.
83    pub fn run(&self, selector: &str, inputs: &[RecordBatch]) -> Result<RecordBatch, ComputeError> {
84        self.load(selector)?.run(inputs)
85    }
86
87    /// Load `selector` without running it. Call this when a job is configured, so a missing or
88    /// broken transform fails the job up front instead of failing every batch later.
89    pub fn ensure_loadable(&self, selector: &str) -> Result<(), ComputeError> {
90        self.load(selector).map(drop)
91    }
92
93    /// Drop every cached transform. The next `load` of each selector goes through its backend
94    /// again — use after the transform directories change on disk.
95    pub fn clear_cache(&self) {
96        self.loaded.lock().unwrap().clear();
97    }
98}
99
100impl std::fmt::Debug for Runtime {
101    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
102        f.debug_struct("Runtime")
103            .field("transforms", &self.registry.len())
104            .field("backends", &self.backends.kinds().collect::<Vec<_>>())
105            .field("loaded", &self.loaded.lock().unwrap().len())
106            .finish()
107    }
108}
109
110/// Builds a [`Runtime`]: roots in precedence order (later roots override earlier ones for the
111/// same `id@version`) and the backends to dispatch to.
112#[derive(Default)]
113pub struct RuntimeBuilder {
114    registry: RegistryBuilder,
115    backends: Backends,
116    roots: usize,
117}
118
119impl RuntimeBuilder {
120    /// A directory of transforms (`<path>/<name>/transform.toml`), named after its last path
121    /// segment. A path that does not exist contributes nothing.
122    pub fn root(mut self, path: impl Into<PathBuf>) -> Self {
123        let path = path.into();
124        let name = path
125            .file_name()
126            .map(|s| s.to_string_lossy().into_owned())
127            .unwrap_or_else(|| format!("root{}", self.roots));
128        self.roots += 1;
129        self.registry = self.registry.root(DirRoot::new(name, path));
130        self
131    }
132
133    /// A directory of transforms with an explicit name (shown in diagnostics and the catalog).
134    pub fn named_root(mut self, name: impl Into<String>, path: impl Into<PathBuf>) -> Self {
135        self.roots += 1;
136        self.registry = self.registry.root(DirRoot::new(name, path));
137        self
138    }
139
140    /// Any other source of transforms — your own [`Root`] implementation.
141    pub fn source(mut self, root: impl Root + 'static) -> Self {
142        self.registry = self.registry.root(root);
143        self
144    }
145
146    /// Register a backend. One per [`crate::ImplKind`]; a later registration for the same kind
147    /// replaces the earlier one.
148    pub fn backend(mut self, backend: impl TransformBackend + 'static) -> Self {
149        self.backends = self.backends.register(backend);
150        self
151    }
152
153    /// Discover and validate every transform under the roots. A malformed manifest anywhere is an
154    /// error here, not at run time.
155    pub fn build(self) -> Result<Runtime, RegistryError> {
156        Ok(Runtime::new(self.registry.build()?, self.backends))
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163    use crate::manifest::{ImplKind, TransformManifest};
164    use crate::registry::MemoryRoot;
165    use std::path::Path;
166    use std::sync::atomic::{AtomicUsize, Ordering};
167
168    const MANIFEST: &str = r#"
169        id = "double"
170        version = "1.0.0"
171        impl = "wasm"
172        entry = "main.wasm"
173        [[inputs]]
174        columns = [{ name = "x", type = "int64" }]
175        [output]
176        columns = [{ name = "x", type = "int64" }]
177    "#;
178
179    /// A backend that counts loads and echoes its input.
180    struct Counting(Arc<AtomicUsize>);
181
182    struct Echo(TransformManifest);
183
184    impl Compute for Echo {
185        fn manifest(&self) -> &TransformManifest {
186            &self.0
187        }
188        fn run(&self, inputs: &[RecordBatch]) -> Result<RecordBatch, ComputeError> {
189            Ok(inputs[0].clone())
190        }
191    }
192
193    impl TransformBackend for Counting {
194        fn kind(&self) -> ImplKind {
195            ImplKind::Wasm
196        }
197        fn load(&self, manifest: &TransformManifest, _root: &Path) -> Result<Box<dyn Compute>, ComputeError> {
198            self.0.fetch_add(1, Ordering::SeqCst);
199            Ok(Box::new(Echo(manifest.clone())))
200        }
201    }
202
203    fn batch() -> RecordBatch {
204        use arrow::array::Int64Array;
205        use arrow::datatypes::{DataType, Field, Schema};
206        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, true)]));
207        RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![1, 2, 3]))]).unwrap()
208    }
209
210    #[test]
211    fn loads_once_per_selector_and_runs() {
212        let loads = Arc::new(AtomicUsize::new(0));
213        let runtime = Runtime::builder()
214            .source(MemoryRoot::new("mem").unit("double", MANIFEST))
215            .backend(Counting(loads.clone()))
216            .build()
217            .unwrap();
218
219        runtime.ensure_loadable("double").unwrap();
220        let out = runtime.run("double", &[batch()]).unwrap();
221        let again = runtime.run("double@1.0.0", &[batch()]).unwrap();
222        assert_eq!(out.num_rows(), 3);
223        assert_eq!(again.num_rows(), 3);
224        // "double" and "double@1.0.0" are distinct selectors: two loads, never a third.
225        assert_eq!(loads.load(Ordering::SeqCst), 2);
226
227        runtime.clear_cache();
228        runtime.run("double", &[batch()]).unwrap();
229        assert_eq!(loads.load(Ordering::SeqCst), 3);
230    }
231
232    #[test]
233    fn missing_transform_and_missing_backend_are_distinct_errors() {
234        let runtime = Runtime::builder()
235            .source(MemoryRoot::new("mem").unit("double", MANIFEST))
236            .build()
237            .unwrap();
238        assert!(matches!(runtime.load("nope"), Err(ComputeError::Load { .. })));
239        assert!(matches!(
240            runtime.load("double"),
241            Err(ComputeError::NoBackend { kind: ImplKind::Wasm })
242        ));
243    }
244
245    #[test]
246    fn a_missing_directory_root_is_empty_not_an_error() {
247        let runtime = Runtime::builder().root("/definitely/not/here").build().unwrap();
248        assert!(runtime.registry().is_empty());
249        assert_eq!(
250            format!("{runtime:?}"),
251            "Runtime { transforms: 0, backends: [], loaded: 0 }"
252        );
253    }
254}