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

zen_engine/nodes/
extensions.rs
