1use crate::worker::Worker;
2use anyhow::Result;
3use async_trait::async_trait;
4use deadpool::managed;
5use tokio::time::Instant;
6use tracing::{debug, debug_span};
7
8#[derive(Debug)]
9pub struct Manager {
10 path: String,
11}
12
13impl Manager {
14 pub fn new(path: &str) -> Self {
15 Self {
16 path: String::from(path),
17 }
18 }
19}
20
21#[async_trait]
22impl managed::Manager for Manager {
23 type Type = Worker;
24 type Error = anyhow::Error;
25
26 async fn create(&self) -> Result<Self::Type, Self::Error> {
27 let start_time = Instant::now();
28 let worker = Worker::new(&self.path).await?;
29 debug_span!("[Worker]", path = &self.path).in_scope(|| {
30 debug!(eplased = ?start_time.elapsed(), "create, ok");
31 });
32 Ok(worker)
33 }
34
35 async fn recycle(&self, _obj: &mut Self::Type) -> managed::RecycleResult<Self::Error> {
36 Ok(())
37 }
38}
39
40pub type WorkerPool = managed::Pool<Manager>;
41
42pub fn create_pool(path: &str) -> Result<WorkerPool> {
44 let mgr = Manager::new(path);
45 Ok(managed::Pool::builder(mgr).build().unwrap())
46}
47
48#[cfg(test)]
49mod tests {
50 use crate::{host_call::Request, worker::Context};
51 use hyper::Body;
52
53 #[tokio::test]
54 async fn run_worker_pool_test() {
55 let wasm_file = "../../tests/data/rust_impl.component.wasm";
56 let pool = super::create_pool(wasm_file).unwrap();
57
58 let status = pool.status();
59 assert_eq!(status.size, 0);
60 assert_eq!(status.available, 0);
61
62 {
63 let mut worker = pool.get().await.unwrap();
64 let worker = worker.as_mut();
65
66 let mut context = Context::default();
67 let body = Body::from("test request body");
68 let body_handle = context.set_body(body);
69
70 let headers: Vec<(String, String)> = vec![];
71 let req = Request {
72 method: "GET",
73 uri: "/abc",
74 headers: &headers,
75 body: Some(body_handle),
76 };
77
78 let (resp, _body) = worker.handle_request(req, context).await.unwrap();
79 assert_eq!(resp.status, 200);
80 assert_eq!(resp.body, Some(2));
81
82 let headers = resp.headers;
83 for (key, value) in headers {
84 if key == "X-Request-Method" {
85 assert_eq!(value, "GET");
86 }
87 if key == "X-Request-Url" {
88 assert_eq!(value, "/abc");
89 }
90 }
91 }
92
93 let status = pool.status();
94 assert_eq!(status.size, 1);
95 assert_eq!(status.available, 1);
96 }
97}