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    pub(crate) stripped_functions: Option<Arc<ahash::HashMap<Arc<str>, Arc<str>>>>,
24}
25
26impl Default for NodeHandlerExtensions {
27    fn default() -> Self {
28        Self {
29            function_runtime: Default::default(),
30            validator_cache: Default::default(),
31
32            loader: Arc::new(NoopLoader::default()),
33            custom_node: Arc::new(NoopCustomNode::default()),
34            compiled_cache: None,
35            stripped_functions: None,
36            http_handler: None,
37        }
38    }
39}
40
41impl NodeHandlerExtensions {
42    pub async fn function_runtime(&self) -> anyhow::Result<&Function> {
43        self.function_runtime
44            .get_or_try_init(|| {
45                Function::create(FunctionConfig {
46                    listeners: Some(vec![
47                        Box::new(ConsoleListener),
48                        Box::new(HttpListener {
49                            http_handler: self.http_handler.clone(),
50                        }),
51                        Box::new(ZenListener {
52                            loader: self.loader.clone(),
53                            custom_node: self.custom_node.clone(),
54                            http_handler: self.http_handler.clone(),
55                        }),
56                    ]),
57                })
58            })
59            .await
60            .context("Failed to create function")
61    }
62
63    pub fn validator_cache(&self) -> &ValidatorCache {
64        self.validator_cache
65            .get_or_init(|| ValidatorCache::default())
66    }
67
68    pub fn custom_node(&self) -> &DynamicCustomNode {
69        &self.custom_node
70    }
71
72    pub fn loader(&self) -> &DynamicLoader {
73        &self.loader
74    }
75
76    pub fn http_handler(&self) -> &DynamicHttpHandler {
77        &self.http_handler
78    }
79}