zen_engine/nodes/
extensions.rs

1use crate::loader::{DynamicLoader, NoopLoader};
2use crate::model::CompilationKey;
3use crate::nodes::custom::{DynamicCustomNode, NoopCustomNode};
4use crate::nodes::function::http_handler::DynamicHttpHandler;
5use crate::nodes::function::v2::function::{Function, FunctionConfig};
6use crate::nodes::function::v2::module::console::ConsoleListener;
7use crate::nodes::function::v2::module::http::listener::HttpListener;
8use crate::nodes::function::v2::module::zen::ZenListener;
9use crate::nodes::validator_cache::ValidatorCache;
10use ahash::HashMap;
11use anyhow::Context;
12use std::cell::OnceCell;
13use std::sync::Arc;
14use zen_expression::compiler::Opcode;
15
16/// This is created on every graph evaluation
17#[derive(Debug, Clone)]
18pub struct NodeHandlerExtensions {
19    pub(crate) function_runtime: Arc<tokio::sync::OnceCell<Function>>,
20    pub(crate) validator_cache: Arc<OnceCell<ValidatorCache>>,
21    pub(crate) loader: DynamicLoader,
22    pub(crate) custom_node: DynamicCustomNode,
23    pub(crate) http_handler: DynamicHttpHandler,
24    pub(crate) compiled_cache: Option<Arc<HashMap<CompilationKey, Vec<Opcode>>>>,
25}
26
27impl Default for NodeHandlerExtensions {
28    fn default() -> Self {
29        Self {
30            function_runtime: Default::default(),
31            validator_cache: Default::default(),
32
33            loader: Arc::new(NoopLoader::default()),
34            custom_node: Arc::new(NoopCustomNode::default()),
35            compiled_cache: 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}
zen_engine/nodes/extensions.rs

zen_engine/nodes/
extensions.rs
