land_runtime/
worker.rs

1use crate::host_call::{HttpContext, HttpHandler, HttpService, Request, Response};
2use anyhow::Result;
3use hyper::body::Body;
4use std::fmt::Debug;
5use wasmtime::component::{Component, InstancePre, Linker};
6use wasmtime::{Config, Engine, InstanceAllocationStrategy, PoolingAllocationConfig, Store};
7use wasmtime_wasi::preview2::{Table, WasiCtx, WasiCtxBuilder, WasiView};
8
9pub struct Context {
10    wasi_ctx: WasiCtx,
11    table: Table,
12    http_ctx: HttpContext,
13}
14
15impl Default for Context {
16    fn default() -> Self {
17        Self::new(uuid::Uuid::new_v4().to_string())
18    }
19}
20
21impl WasiView for Context {
22    fn table(&self) -> &Table {
23        &self.table
24    }
25    fn table_mut(&mut self) -> &mut Table {
26        &mut self.table
27    }
28    fn ctx(&self) -> &WasiCtx {
29        &self.wasi_ctx
30    }
31    fn ctx_mut(&mut self) -> &mut WasiCtx {
32        &mut self.wasi_ctx
33    }
34}
35
36impl Context {
37    pub fn new(req_id: String) -> Self {
38        let mut table = Table::new();
39        Context {
40            wasi_ctx: WasiCtxBuilder::new()
41                .inherit_stdio()
42                .build(&mut table)
43                .unwrap(),
44            http_ctx: HttpContext::new(req_id),
45            table,
46        }
47    }
48
49    /// get http_ctx
50    pub fn http_ctx(&mut self) -> &mut HttpContext {
51        &mut self.http_ctx
52    }
53
54    /// set body
55    pub fn set_body(&mut self, body: Body) -> u32 {
56        self.http_ctx.set_body(body)
57    }
58
59    /// take body
60    pub fn take_body(&mut self, handle: u32) -> Option<Body> {
61        self.http_ctx.take_body(handle)
62    }
63
64    /// get request id
65    pub fn req_id(&self) -> String {
66        self.http_ctx.req_id.clone()
67    }
68}
69
70fn create_wasmtime_config() -> Config {
71    let mut config = Config::new();
72    config.wasm_component_model(true);
73    config.async_support(true);
74
75    const MB: usize = 1 << 20;
76    let mut pooling_allocation_config = PoolingAllocationConfig::default();
77    pooling_allocation_config.instance_size(MB);
78    pooling_allocation_config.instance_memory_pages(128 * (MB as u64) / (64 * 1024));
79    config.allocation_strategy(InstanceAllocationStrategy::Pooling(
80        pooling_allocation_config,
81    ));
82
83    config
84}
85
86/// Worker is used to run wasm component
87pub struct Worker {
88    path: String,
89    engine: Engine,
90    // component: Component,
91    instance_pre: InstancePre<Context>,
92}
93
94impl Debug for Worker {
95    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
96        f.debug_struct("Worker").field("path", &self.path).finish()
97    }
98}
99
100impl Worker {
101    /// new a worker
102    pub async fn new(path: &str) -> Result<Self> {
103        let binary = std::fs::read(path)?;
104        Self::from_binary(&binary).await
105    }
106
107    // from_binary is used to create worker from bytes
108    pub async fn from_binary(bytes: &[u8]) -> Result<Self> {
109        // create component
110        let config = create_wasmtime_config();
111        let engine = Engine::new(&config)?;
112        let component = Component::from_binary(&engine, bytes)?;
113
114        // create linker
115        let mut linker: Linker<Context> = Linker::new(&engine);
116        // init wasi context
117        wasmtime_wasi::preview2::wasi::command::add_to_linker(&mut linker)
118            .expect("add wasmtime_wasi::preview2 failed");
119        HttpService::add_to_linker(&mut linker, Context::http_ctx)?;
120
121        Ok(Self {
122            path: "bytes".to_string(),
123            engine,
124            instance_pre: linker.instantiate_pre(&component)?,
125        })
126    }
127
128    /// handle_request is used to handle http request
129    pub async fn handle_request(
130        &mut self,
131        req: Request<'_>,
132        context: Context,
133    ) -> Result<(Response, Body)> {
134        // create store
135        let mut store = Store::new(&self.engine, context);
136
137        // get exports and call handle_request
138        let (exports, _instance) =
139            HttpHandler::instantiate_pre(&mut store, &self.instance_pre).await?;
140        let resp = exports
141            .land_http_http_incoming()
142            .call_handle_request(&mut store, req)
143            .await?;
144        let body = store.data_mut().take_body(resp.body.unwrap()).unwrap();
145        Ok((resp, body))
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use crate::{
152        host_call::Request,
153        worker::{Context, Worker},
154    };
155    use hyper::Body;
156
157    #[tokio::test]
158    async fn run_wasm() {
159        let wasm_file = "../../tests/data/rust_impl.component.wasm";
160        let mut worker = Worker::new(wasm_file).await.unwrap();
161
162        for _ in 1..10 {
163            let headers: Vec<(String, String)> = vec![];
164
165            let mut context = Context::default();
166            let body = Body::from("test request body");
167            let body_handle = context.set_body(body);
168
169            let req = Request {
170                method: "GET",
171                uri: "/abc",
172                headers: &headers,
173                body: Some(body_handle),
174            };
175
176            let (resp, _body) = worker.handle_request(req, context).await.unwrap();
177            assert_eq!(resp.status, 200);
178            // this wasm return request's body
179            // so the body handler u32 is 2, same as request's body
180            assert_eq!(resp.body, Some(2));
181
182            let headers = resp.headers;
183            for (key, value) in headers {
184                if key == "X-Request-Method" {
185                    assert_eq!(value, "GET");
186                }
187                if key == "X-Request-Url" {
188                    assert_eq!(value, "/abc");
189                }
190            }
191        }
192    }
193}
land_runtime/worker.rs

land_runtime/
worker.rs
