1use 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
35pub struct Runtime {
37 registry: Registry,
38 backends: Backends,
39 loaded: Mutex<HashMap<String, Arc<dyn Compute>>>,
40}
41
42impl Runtime {
43 pub fn builder() -> RuntimeBuilder {
45 RuntimeBuilder::default()
46 }
47
48 pub fn new(registry: Registry, backends: Backends) -> Self {
50 Self {
51 registry,
52 backends,
53 loaded: Mutex::new(HashMap::new()),
54 }
55 }
56
57 pub fn registry(&self) -> &Registry {
59 &self.registry
60 }
61
62 pub fn backends(&self) -> &Backends {
64 &self.backends
65 }
66
67 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 pub fn run(&self, selector: &str, inputs: &[RecordBatch]) -> Result<RecordBatch, ComputeError> {
84 self.load(selector)?.run(inputs)
85 }
86
87 pub fn ensure_loadable(&self, selector: &str) -> Result<(), ComputeError> {
90 self.load(selector).map(drop)
91 }
92
93 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#[derive(Default)]
113pub struct RuntimeBuilder {
114 registry: RegistryBuilder,
115 backends: Backends,
116 roots: usize,
117}
118
119impl RuntimeBuilder {
120 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 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 pub fn source(mut self, root: impl Root + 'static) -> Self {
142 self.registry = self.registry.root(root);
143 self
144 }
145
146 pub fn backend(mut self, backend: impl TransformBackend + 'static) -> Self {
149 self.backends = self.backends.register(backend);
150 self
151 }
152
153 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 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 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}