Skip to main content

zen_engine/nodes/
extensions.rs

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