edgee_components_runtime/
context.rs

1use std::collections::HashMap;
2
3use wasmtime::{
4    component::{Component, Linker, ResourceTable},
5    Engine, Store,
6};
7use wasmtime_wasi::{WasiCtx, WasiView};
8
9use crate::config::{ComponentsConfiguration, DataCollectionComponents};
10use crate::consent_mapping::{ConsentMapping, ConsentMappingPre};
11use crate::data_collection::{DataCollection, DataCollectionPre};
12
13pub struct ComponentsContext {
14    pub engine: Engine,
15    pub components: Components,
16}
17
18pub struct Components {
19    pub data_collection: HashMap<String, DataCollectionPre<HostState>>,
20    pub consent_mapping: HashMap<String, ConsentMappingPre<HostState>>,
21}
22
23pub fn pre_instanciate_data_collection_component_internal(
24    engine: &Engine,
25    component_config: &DataCollectionComponents,
26) -> anyhow::Result<DataCollectionPre<HostState>> {
27    let mut linker = Linker::new(engine);
28    wasmtime_wasi::add_to_linker_async(&mut linker)?;
29
30    let span = tracing::info_span!("component-context", component = %component_config.id, category = "data-collection");
31    let _span = span.enter();
32
33    tracing::debug!("Loading new data collection component");
34
35    let component = Component::from_file(engine, &component_config.file)?;
36    let instance_pre = linker.instantiate_pre(&component)?;
37    let instance_pre = DataCollectionPre::new(instance_pre)?;
38
39    tracing::debug!("loaded new data collection component");
40
41    Ok(instance_pre)
42}
43
44impl ComponentsContext {
45    pub fn new(config: &ComponentsConfiguration) -> anyhow::Result<Self> {
46        let mut engine_config = wasmtime::Config::new();
47        engine_config
48            .wasm_backtrace_details(wasmtime::WasmBacktraceDetails::Enable)
49            .wasm_component_model(true)
50            .async_support(true);
51
52        if let Some(path) = config.cache.as_deref() {
53            engine_config.cache_config_load(path)?;
54        };
55
56        let engine = Engine::new(&engine_config)?;
57
58        let mut linker = Linker::new(&engine);
59        wasmtime_wasi::add_to_linker_async(&mut linker)?;
60
61        // Data collection components
62        let data_collection_components = config
63            .data_collection
64            .iter()
65            .map(|entry| {
66                let instance_pre =
67                    pre_instanciate_data_collection_component_internal(&engine, entry)?;
68                Ok((entry.id.clone(), instance_pre))
69            })
70            .collect::<anyhow::Result<_>>()?;
71
72        // Consent mapping components
73        let consent_mapping_components = config
74            .consent_mapping
75            .iter()
76            .map(|entry| {
77                let span = tracing::info_span!("component-context", component = %entry.name, category = "consent-mapping");
78                let _span = span.enter();
79
80                tracing::debug!("Start pre-instantiate consent mapping component");
81
82                let component = Component::from_file(&engine, &entry.component)?;
83                let instance_pre = linker.instantiate_pre(&component)?;
84                let instance_pre = ConsentMappingPre::new(instance_pre)?;
85
86                tracing::debug!("Finished pre-instantiate consent mapping component");
87
88                Ok((entry.name.clone(), instance_pre))
89            })
90            .collect::<anyhow::Result<_>>()?;
91
92        let components = Components {
93            data_collection: data_collection_components,
94            consent_mapping: consent_mapping_components,
95        };
96
97        Ok(Self { engine, components })
98    }
99
100    pub fn pre_instanciate_data_collection_component(
101        &self,
102        component_config: DataCollectionComponents,
103    ) -> anyhow::Result<DataCollectionPre<HostState>> {
104        let instance_pre =
105            pre_instanciate_data_collection_component_internal(&self.engine, &component_config)?;
106        Ok(instance_pre)
107    }
108
109    pub fn add_data_collection_component(
110        &mut self,
111        component_config: DataCollectionComponents,
112        instance_pre: DataCollectionPre<HostState>,
113    ) {
114        if !self
115            .components
116            .data_collection
117            .contains_key(&component_config.id)
118        {
119            self.components
120                .data_collection
121                .insert(component_config.id.clone(), instance_pre);
122        }
123    }
124
125    pub fn empty_store(&self) -> Store<HostState> {
126        Store::new(&self.engine, HostState::new())
127    }
128
129    pub fn empty_store_with_stdout(&self) -> Store<HostState> {
130        Store::new(&self.engine, HostState::new_with_stdout())
131    }
132
133    pub async fn get_data_collection_instance(
134        &self,
135        id: &str,
136        store: &mut Store<HostState>,
137    ) -> anyhow::Result<DataCollection> {
138        let instance_pre = self.components.data_collection.get(id);
139
140        if instance_pre.is_none() {
141            return Err(anyhow::anyhow!("component not found: {}", id));
142        }
143
144        instance_pre.unwrap().instantiate_async(store).await
145    }
146
147    pub async fn get_consent_mapping_instance(
148        &self,
149        id: &str,
150        store: &mut Store<HostState>,
151    ) -> anyhow::Result<ConsentMapping> {
152        let instance_pre = self.components.consent_mapping.get(id);
153
154        if instance_pre.is_none() {
155            return Err(anyhow::anyhow!("component not found: {}", id));
156        }
157
158        instance_pre.unwrap().instantiate_async(store).await
159    }
160}
161
162pub struct HostState {
163    ctx: WasiCtx,
164    table: ResourceTable,
165}
166
167impl HostState {
168    fn new() -> Self {
169        Self::new_with_ctx(WasiCtx::builder().build())
170    }
171
172    fn new_with_stdout() -> Self {
173        Self::new_with_ctx(WasiCtx::builder().inherit_stdout().build())
174    }
175
176    fn new_with_ctx(ctx: WasiCtx) -> Self {
177        let table = ResourceTable::new();
178        Self { ctx, table }
179    }
180}
181
182impl WasiView for HostState {
183    fn ctx(&mut self) -> &mut WasiCtx {
184        &mut self.ctx
185    }
186
187    fn table(&mut self) -> &mut ResourceTable {
188        &mut self.table
189    }
190}
edgee_components_runtime/context.rs

edgee_components_runtime/
context.rs
