diff --git a/Cargo.lock b/Cargo.lock index 54c47ee..934a9a7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2760,6 +2760,7 @@ dependencies = [ "parking_lot 0.12.1", "prost", "rand", + "rayon", "reqwest", "serde", "sled", diff --git a/wacker/Cargo.toml b/wacker/Cargo.toml index 3732631..59e8630 100644 --- a/wacker/Cargo.toml +++ b/wacker/Cargo.toml @@ -43,6 +43,7 @@ env_logger = "0.11.3" chrono = "0.4.38" sled = "0.34.7" log = "0.4.21" +rayon = "1.10.0" [build-dependencies] anyhow.workspace = true diff --git a/wacker/src/server.rs b/wacker/src/server.rs index ccd49c4..62c2152 100644 --- a/wacker/src/server.rs +++ b/wacker/src/server.rs @@ -12,6 +12,7 @@ use async_stream::try_stream; use async_trait::async_trait; use log::{error, info, warn}; use parking_lot::Mutex; +use rayon::prelude::*; use sled::Db; use std::fmt::Display; use std::fs::{remove_file, OpenOptions}; @@ -135,6 +136,23 @@ impl Server { Err(err) => Err(Status::internal(err.to_string())), } } + + fn get_program_keys(&self) -> Vec { + let programs = self.programs.lock(); + programs.par_iter().map(|(key, _)| key.clone()).collect() + } +} + +fn search_id(keys: &Vec, id: &str) -> Result { + let positions: Vec<_> = keys.par_iter().positions(|s| s.starts_with(id)).collect(); + match positions.len() { + 0 => Err(anyhow!("program {} not found", id)), + 1 => Ok(keys[positions[0]].clone()), + _ => Err(anyhow!( + "ambiguous program id {}, more than one program starts with this id", + id + )), + } } fn to_status(err: E) -> Status { @@ -223,19 +241,18 @@ impl Wacker for Server { async fn stop(&self, request: Request) -> Result, Status> { let req = request.into_inner(); + let keys = self.get_program_keys(); let mut programs = self.programs.lock(); for id in req.ids { + let id = search_id(keys.as_ref(), id.as_str()).map_err(to_status)?; + info!("Stop the program: {}", id); - match programs.get_mut(id.as_str()) { - Some(program) => { - if !program.handler.is_finished() { - program.handler.abort(); - program.status = PROGRAM_STATUS_STOPPED; - } - } - None => return Err(Status::not_found(format!("program {} not exists", id))), + let program = programs.get_mut(id.as_str()).unwrap(); + if !program.handler.is_finished() { + program.handler.abort(); + program.status = PROGRAM_STATUS_STOPPED; } } Ok(Response::new(())) @@ -243,18 +260,16 @@ impl Wacker for Server { async fn restart(&self, request: Request) -> Result, Status> { let req = request.into_inner(); + let keys = self.get_program_keys(); for id in req.ids { + let id = search_id(keys.as_ref(), id.as_str()).map_err(to_status)?; + info!("Restart the program: {}", id); let meta = { let programs = self.programs.lock(); - let program = programs.get(id.as_str()); - if program.is_none() { - return Err(Status::not_found(format!("program {} not exists", id))); - } - - let program = program.unwrap(); + let program = programs.get(id.as_str()).unwrap(); if !program.handler.is_finished() { program.handler.abort(); } @@ -268,29 +283,31 @@ impl Wacker for Server { async fn delete(&self, request: Request) -> Result, Status> { let req = request.into_inner(); + let keys = self.get_program_keys(); let mut programs = self.programs.lock(); for id in req.ids { + let id = search_id(keys.as_ref(), id.as_str()).map_err(to_status)?; + info!("Delete the program: {}", id); - if let Some(program) = programs.get(id.as_str()) { - if !program.handler.is_finished() { - program.handler.abort(); - } + let program = programs.get(id.as_str()).unwrap(); + if !program.handler.is_finished() { + program.handler.abort(); + } - if let Err(err) = remove_file(self.logs_dir.join(id.as_str())) { - if err.kind() != ErrorKind::NotFound { - return Err(Status::internal(format!( - "failed to remove the log file for {}: {}", - id.as_str(), - err - ))); - } + if let Err(err) = remove_file(self.logs_dir.join(id.as_str())) { + if err.kind() != ErrorKind::NotFound { + return Err(Status::internal(format!( + "failed to remove the log file for {}: {}", + id.as_str(), + err + ))); } - - self.db.remove(id.as_str()).map_err(to_status)?; - programs.remove(id.as_str()); } + + self.db.remove(id.as_str()).map_err(to_status)?; + programs.remove(id.as_str()); } Ok(Response::new(())) } @@ -299,8 +316,10 @@ impl Wacker for Server { async fn logs(&self, request: Request) -> Result, Status> { let req = request.into_inner(); + let keys = self.get_program_keys(); + let id = search_id(keys.as_ref(), req.id.as_str()).map_err(to_status)?; - let mut file = File::open(self.logs_dir.join(req.id)).await?; + let mut file = File::open(self.logs_dir.join(id)).await?; let mut contents = String::new(); let last_position = file.read_to_string(&mut contents).await?; let lines: Vec<&str> = contents.split_inclusive('\n').collect(); diff --git a/wacker/tests/server.rs b/wacker/tests/server.rs index 8a029fa..59af63c 100644 --- a/wacker/tests/server.rs +++ b/wacker/tests/server.rs @@ -98,16 +98,19 @@ async fn stop() -> Result<()> { server.start().await; let mut client = server.client().await; - let response = client + client .run(RunRequest { path: "./tests/wasm/time.wasm".parse()?, args: vec![], }) - .await? - .into_inner(); + .await?; sleep(Duration::from_secs(1)).await; - client.stop(StopRequest { ids: vec![response.id] }).await?; + client + .stop(StopRequest { + ids: vec!["t".to_string()], + }) + .await?; sleep(Duration::from_secs(1)).await; let response = client.list(()).await?.into_inner(); @@ -139,11 +142,9 @@ async fn restart() -> Result<()> { .into_inner(); sleep(Duration::from_secs(1)).await; - let response = client.restart(RestartRequest { ids: vec![run_resp.id] }).await; - assert!(response.is_ok()); let response = client .restart(RestartRequest { - ids: vec![serve_resp.id], + ids: vec![run_resp.id, serve_resp.id], }) .await; assert!(response.is_ok()); @@ -158,16 +159,19 @@ async fn delete() -> Result<()> { server.start().await; let mut client = server.client().await; - let response = client + client .run(RunRequest { path: "./tests/wasm/hello.wasm".parse()?, args: vec![], }) - .await? - .into_inner(); + .await?; sleep(Duration::from_secs(1)).await; - client.delete(DeleteRequest { ids: vec![response.id] }).await?; + client + .delete(DeleteRequest { + ids: vec!["h".to_string()], + }) + .await?; let response = client.list(()).await?.into_inner(); assert_eq!(response.programs.len(), 0); @@ -177,23 +181,57 @@ async fn delete() -> Result<()> { } #[tokio::test(flavor = "multi_thread")] -async fn logs() -> Result<()> { +async fn delete_ambiguous_id() -> Result<()> { let mut server = TestServer::new(); server.start().await; let mut client = server.client().await; + client + .run(RunRequest { + path: "./tests/wasm/hello.wasm".parse()?, + args: vec![], + }) + .await?; + client + .run(RunRequest { + path: "./tests/wasm/hello.wasm".parse()?, + args: vec![], + }) + .await?; + sleep(Duration::from_secs(1)).await; + let response = client + .delete(DeleteRequest { + ids: vec!["hello".to_string()], + }) + .await; + assert!(response.is_err()); + assert_eq!( + response.err().unwrap().message(), + "ambiguous program id hello, more than one program starts with this id" + ); + + server.shutdown().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread")] +async fn logs() -> Result<()> { + let mut server = TestServer::new(); + server.start().await; + + let mut client = server.client().await; + client .run(RunRequest { path: "./tests/wasm/hello.wasm".parse()?, args: vec![], }) - .await? - .into_inner(); + .await?; sleep(Duration::from_secs(3)).await; let mut response = client .logs(LogRequest { - id: response.id, + id: "hello".to_string(), follow: false, tail: 1, })