Skip to main content

zen_engine/nodes/
extensions.rs

1use crate::loader::{DynamicLoader, NoopLoader};
2use crate::nodes::custom::{DynamicCustomNode, NoopCustomNode};
3use crate::nodes::function::http_handler::DynamicHttpHandler;
4use crate::nodes::function::v2::function::{Function, FunctionConfig};
5use crate::nodes::function::v2::module::console::ConsoleListener;
6use crate::nodes::function::v2::module::http::listener::HttpListener;
7use crate::nodes::function::v2::module::zen::ZenListener;
8use crate::nodes::validator_cache::ValidatorCache;
9use anyhow::Context;
10use std::cell::OnceCell;
11use std::sync::Arc;
12use zen_expression::OpcodeCache;
13
14/// This is created on every graph evaluation
15#[derive(Debug, Clone)]
16pub struct NodeHandlerExtensions {
17    pub(crate) function_runtime: Arc<tokio::sync::OnceCell<Function>>,
18    pub(crate) validator_cache: Arc<OnceCell<ValidatorCache>>,
19    pub(crate) loader: DynamicLoader,
20    pub(crate) custom_node: DynamicCustomNode,
21    pub(crate) http_handler: DynamicHttpHandler,
22    pub(crate) compiled_cache: Option<Arc<OpcodeCache>>,
23}
24
25impl Default for NodeHandlerExtensions {
26    fn default() -> Self {
27        Self {
28            function_runtime: Default::default(),
29            validator_cache: Default::default(),
30
31            loader: Arc::new(NoopLoader::default()),
32            custom_node: Arc::new(NoopCustomNode::default()),
33            compiled_cache: None,
34            http_handler: None,
35        }
36    }
37}
38
39impl NodeHandlerExtensions {
40    pub async fn function_runtime(&self) -> anyhow::Result<&Function> {
41        self.function_runtime
42            .get_or_try_init(|| {
43                Function::create(FunctionConfig {
44                    listeners: Some(vec![
45                        Box::new(ConsoleListener),
46                        Box::new(HttpListener {
47                            http_handler: self.http_handler.clone(),
48                        }),
49                        Box::new(ZenListener {
50                            loader: self.loader.clone(),
51                            custom_node: self.custom_node.clone(),
52                            http_handler: self.http_handler.clone(),
53                        }),
54                    ]),
55                })
56            })
57            .await
58            .context("Failed to create function")
59    }
60
61    pub fn validator_cache(&self) -> &ValidatorCache {
62        self.validator_cache
63            .get_or_init(|| ValidatorCache::default())
64    }
65
66    pub fn custom_node(&self) -> &DynamicCustomNode {
67        &self.custom_node
68    }
69
70    pub fn loader(&self) -> &DynamicLoader {
71        &self.loader
72    }
73
74    pub fn http_handler(&self) -> &DynamicHttpHandler {
75        &self.http_handler
76    }
77}