From 56348fea0467854fda0dfd24d0e814f76bcc574f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=8B=8F=E5=90=91=E5=A4=9C?= Date: Sun, 28 Apr 2024 19:09:42 +0800 Subject: [PATCH] perf: use rayon and fastrand to imporve exectute speed --- .github/workflows/benchmark.yml | 26 ++++++++++++++++++++++++++ Cargo.lock | 8 ++++++++ Cargo.toml | 6 ++++++ benches/rand.rs | 9 +++++++++ src/models/router.rs | 28 ++++++++++++++++++---------- src/models/server.rs | 3 +-- src/utils/generator.rs | 6 +++--- 7 files changed, 71 insertions(+), 15 deletions(-) create mode 100644 .github/workflows/benchmark.yml create mode 100644 benches/rand.rs diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml new file mode 100644 index 0000000..d347a62 --- /dev/null +++ b/.github/workflows/benchmark.yml @@ -0,0 +1,26 @@ +name: Benchmark + +on: + push: + branches: ["main"] + pull_request: + branches: ["main"] + +env: + CARGO_TERM_COLOR: always + +jobs: + build: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + - name: Build + run: cargo build --verbose + - name: Run benchmarks + run: cargo bench --verbose + - name: Upload benchmark results + uses: actions/upload-artifact@v2 + with: + path: target/criterion/report/ + name: benchmark-results diff --git a/Cargo.lock b/Cargo.lock index 2f325fc..4147804 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -406,6 +406,12 @@ dependencies = [ "zeroize", ] +[[package]] +name = "fastrand" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fc0510504f03c51ada170672ac806f1f105a88aa97a5281117e1ddc3368e51a" + [[package]] name = "ff" version = "0.13.0" @@ -771,6 +777,7 @@ dependencies = [ "colored", "criterion", "elliptic-curve", + "fastrand", "futures", "hkdf", "oblivion-codegen", @@ -778,6 +785,7 @@ dependencies = [ "proc-macro2", "pyo3", "rand", + "rayon", "regex", "ring", "scrypt", diff --git a/Cargo.toml b/Cargo.toml index db5fe0f..76d32d6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,6 +16,7 @@ ring = "0.17" sha2 = "0.10" scrypt = "0.11" hkdf = "0.12.4" +fastrand = "2.1.0" # Utils oblivion-codegen = { path = "oblivion-codegen" } @@ -27,6 +28,7 @@ thiserror = "1" anyhow = "1.0" colored = "2.1" chrono = "0.4" +rayon = "1.10.0" # Optional pyo3 = { version = "0.20", optional = true } @@ -62,3 +64,7 @@ harness = false [[bench]] name = "socket" harness = false + +[[bench]] +name = "rand" +harness = false diff --git a/benches/rand.rs b/benches/rand.rs new file mode 100644 index 0000000..4671bfd --- /dev/null +++ b/benches/rand.rs @@ -0,0 +1,9 @@ +use criterion::{criterion_group, criterion_main, Criterion}; +use oblivion::utils::generator::generate_random_salt; + +fn criterion_benchmark_salt(c: &mut Criterion) { + c.bench_function("salt", |b| b.iter(|| generate_random_salt())); +} + +criterion_group!(benches, criterion_benchmark_salt); +criterion_main!(benches); diff --git a/src/models/router.rs b/src/models/router.rs index 83bec31..d22c499 100644 --- a/src/models/router.rs +++ b/src/models/router.rs @@ -3,6 +3,7 @@ use super::handler::not_found; use super::session::Session; use crate::types::server; use anyhow::Result; +use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; use regex::Regex; use std::collections::HashMap; @@ -18,7 +19,7 @@ impl Route { Self { handler } } - pub fn get_handler(&mut self) -> Handler { + pub fn get_handler(&self) -> Handler { self.handler.clone() } } @@ -44,7 +45,7 @@ impl RoutePath { } } - pub fn check(&mut self, olps: &str) -> Result { + pub fn check(&self, olps: &str) -> Result { if self.route_type == RouteType::RegexPath { let regex = Regex::new(&self.route)?; Ok(regex.is_match(olps)) @@ -69,7 +70,7 @@ impl Router { } pub fn route(&mut self, path: RoutePath, handler: Handler) -> &mut Self { - self.routes.insert(path.clone(), Route { handler: handler }); + self.routes.insert(path.clone(), Route { handler }); self } @@ -78,13 +79,20 @@ impl Router { self.routes.insert(path.clone(), route); } - pub fn get_handler(&self, path: &str) -> Result { - for (route_path, route) in &self.routes { - let mut route_path = route_path.clone(); - if route_path.check(path)? { - return Ok(route.clone()); - }; + pub fn get_handler(&self, path: &str) -> Result { + let handler = self + .routes + .par_iter() + .find_any(|values| { + let (route_path, _) = values; + route_path.check(path).unwrap_or(false) + }) + .map(|route| route.1.get_handler()); + + if let Some(route) = handler { + Ok(route) + } else { + Ok(not_found) } - Ok(Route::new(not_found)) } } diff --git a/src/models/server.rs b/src/models/server.rs index fabdebf..e24d530 100644 --- a/src/models/server.rs +++ b/src/models/server.rs @@ -69,8 +69,7 @@ async fn _handle(router: &Router, stream: TcpStream, peer: SocketAddr) -> Result let socket = Arc::clone(&session.socket); - let mut route = router.get_handler(&session.request.as_ref().unwrap().olps)?; - let callback = route.get_handler()(session).await?; + let callback = router.get_handler(&session.request.as_ref().unwrap().olps)?(session).await?; #[cfg(not(any(feature = "perf", feature = "bench")))] let status_code = callback.get_status_code()?; diff --git a/src/utils/generator.rs b/src/utils/generator.rs index 9a1e43d..dc6b948 100644 --- a/src/utils/generator.rs +++ b/src/utils/generator.rs @@ -13,7 +13,7 @@ use p256::{ecdh::EphemeralSecret, PublicKey}; use ring::agreement::{agree_ephemeral, EphemeralPrivateKey, PublicKey, UnparsedPublicKey, X25519}; use ring::aead::AES_128_GCM; -use ring::rand::{SecureRandom, SystemRandom}; +use ring::rand::SystemRandom; use scrypt::{scrypt, Params}; use sha2::Sha256; @@ -113,8 +113,8 @@ impl SharedKey { /// let salt = generate_random_salt(); /// ``` pub fn generate_random_salt() -> Vec { - let rand = SystemRandom::new(); + let mut rng = fastrand::Rng::new(); let mut key_bytes = vec![0; AES_128_GCM.key_len()]; - rand.fill(&mut key_bytes).unwrap(); + rng.fill(&mut key_bytes); key_bytes }