edgee_components_runtime/
context.rs

1use std::collections::HashMap;
2
3use crate::config::ComponentsConfiguration;
4use crate::data_collection::versions::v1_0_0::data_collection::DataCollectionV100Pre;
5use crate::data_collection::versions::v1_0_0::pre_instanciate_data_collection_component_1_0_0;
6use crate::data_collection::versions::v1_0_1::data_collection::DataCollectionV101Pre;
7use crate::data_collection::versions::v1_0_1::pre_instanciate_data_collection_component_1_0_1;
8use crate::data_collection::versions::DataCollectionWitVersion;
9use http::HeaderValue;
10use wasmtime::{Cache, Engine, Store};
11use wasmtime_wasi::ResourceTable;
12use wasmtime_wasi::{WasiCtx, WasiCtxView, WasiView};
13use wasmtime_wasi_http::{WasiHttpCtx, WasiHttpView};
14
15use crate::edge_function::versions::v1_0_0::edge_function::EdgeFunctionV100Pre;
16use crate::edge_function::versions::v1_0_0::pre_instanciate_edge_function_component_1_0_0;
17use crate::edge_function::versions::EdgeFunctionWitVersion;
18
19#[derive(Clone)]
20pub struct ComponentsContext {
21    pub engine: Engine,
22    pub components: Components,
23}
24
25#[derive(Clone)]
26pub struct Components {
27    pub data_collection_1_0_0: HashMap<String, DataCollectionV100Pre<HostState>>,
28    pub data_collection_1_0_1: HashMap<String, DataCollectionV101Pre<HostState>>,
29    pub edge_function_1_0_0: HashMap<String, EdgeFunctionV100Pre<HostState>>,
30}
31
32impl ComponentsContext {
33    pub fn new(config: &ComponentsConfiguration) -> anyhow::Result<Self> {
34        let mut engine_config = wasmtime::Config::new();
35        engine_config
36            .wasm_backtrace_details(wasmtime::WasmBacktraceDetails::Enable)
37            .wasm_component_model(true)
38            .async_support(true);
39
40        match Cache::from_file(config.cache.as_deref())
41            .map(|cache| engine_config.cache(Some(cache)))
42        {
43            Ok(_) => {}
44            Err(e) => {
45                tracing::warn!("Failed to load cache: {e}");
46            }
47        }
48
49        let engine = Engine::new(&engine_config)?;
50
51        // Data collection components
52        let data_collection_1_0_0_components = config
53            .data_collection
54            .iter()
55            .filter(|entry| entry.wit_version == DataCollectionWitVersion::V1_0_0)
56            .map(|entry| {
57                let instance_pre = pre_instanciate_data_collection_component_1_0_0(&engine, entry)?;
58                Ok((entry.id.clone(), instance_pre))
59            })
60            .collect::<anyhow::Result<_>>()?;
61
62        let data_collection_1_0_1_components = config
63            .data_collection
64            .iter()
65            .filter(|entry| entry.wit_version == DataCollectionWitVersion::V1_0_1)
66            .map(|entry| {
67                let instance_pre = pre_instanciate_data_collection_component_1_0_1(&engine, entry)?;
68                Ok((entry.id.clone(), instance_pre))
69            })
70            .collect::<anyhow::Result<_>>()?;
71
72        let edge_function_1_0_0_components = config
73            .edge_function
74            .iter()
75            .filter(|entry| entry.wit_version == EdgeFunctionWitVersion::V1_0_0)
76            .map(|entry| {
77                let instance_pre = pre_instanciate_edge_function_component_1_0_0(&engine, entry)?;
78                Ok((entry.id.clone(), instance_pre))
79            })
80            .collect::<anyhow::Result<_>>()?;
81
82        let components = Components {
83            data_collection_1_0_0: data_collection_1_0_0_components,
84            data_collection_1_0_1: data_collection_1_0_1_components,
85            edge_function_1_0_0: edge_function_1_0_0_components,
86        };
87
88        Ok(Self { engine, components })
89    }
90
91    pub fn empty_store(&self) -> Store<HostState> {
92        Store::new(&self.engine, HostState::new())
93    }
94
95    pub fn empty_store_with_stdout(&self) -> Store<HostState> {
96        Store::new(&self.engine, HostState::new_with_stdout())
97    }
98
99    pub fn serialize_component(
100        &self,
101        component_path: &str,
102        component_type: &str,
103        component_wit_version: &str,
104    ) -> anyhow::Result<Vec<u8>> {
105        let component = match component_type {
106            "data-collection" => match component_wit_version {
107                "1.0.0" => self
108                    .components
109                    .data_collection_1_0_0
110                    .get(component_path)
111                    .map(|c| c.instance_pre().component())
112                    .ok_or_else(|| anyhow::anyhow!("Component {} not found", component_path))?,
113                "1.0.1" => self
114                    .components
115                    .data_collection_1_0_1
116                    .get(component_path)
117                    .map(|c| c.instance_pre().component())
118                    .ok_or_else(|| anyhow::anyhow!("Component {} not found", component_path))?,
119                _ => anyhow::bail!("Invalid WIT version: {}", component_wit_version),
120            },
121            "edge-function" => match component_wit_version {
122                "1.0.0" => self
123                    .components
124                    .edge_function_1_0_0
125                    .get(component_path)
126                    .map(|c| c.instance_pre().component())
127                    .ok_or_else(|| anyhow::anyhow!("Component {} not found", component_path))?,
128                _ => anyhow::bail!("Invalid WIT version: {}", component_wit_version),
129            },
130            _ => anyhow::bail!("Invalid component type: {}", component_type),
131        };
132        component.serialize()
133    }
134}
135
136pub struct HostState {
137    ctx: WasiCtx,
138    table: ResourceTable,
139    http: WasiHttpCtx,
140    component_id: Option<String>,
141}
142
143impl HostState {
144    fn new() -> Self {
145        Self::new_with_ctx(WasiCtx::builder().build())
146    }
147
148    fn new_with_stdout() -> Self {
149        Self::new_with_ctx(WasiCtx::builder().inherit_stdout().inherit_stderr().build())
150    }
151
152    fn new_with_ctx(ctx: WasiCtx) -> Self {
153        Self {
154            ctx,
155            table: ResourceTable::new(),
156            http: WasiHttpCtx::new(),
157            component_id: None,
158        }
159    }
160
161    pub fn set_component_id(&mut self, component_id: String) {
162        self.component_id = Some(component_id);
163    }
164}
165
166impl WasiHttpView for HostState {
167    fn ctx(&mut self) -> &mut WasiHttpCtx {
168        &mut self.http
169    }
170    fn table(&mut self) -> &mut ResourceTable {
171        &mut self.table
172    }
173    fn send_request(
174        &mut self,
175        mut request: hyper::Request<wasmtime_wasi_http::body::HyperOutgoingBody>,
176        config: wasmtime_wasi_http::types::OutgoingRequestConfig,
177    ) -> wasmtime_wasi_http::HttpResult<wasmtime_wasi_http::types::HostFutureIncomingResponse> {
178        tracing::info!("Sending outbound http request: {:?}", request);
179        if let Some(component_id) = &self.component_id {
180            request.headers_mut().insert(
181                "x-edgee-component-id",
182                HeaderValue::from_str(component_id).unwrap(),
183            );
184        }
185        Ok(wasmtime_wasi_http::types::default_send_request(
186            request, config,
187        ))
188    }
189}
190
191impl WasiView for HostState {
192    fn ctx(&mut self) -> WasiCtxView<'_> {
193        WasiCtxView {
194            ctx: &mut self.ctx,
195            table: &mut self.table,
196        }
197    }
198}