[rust/axum] Axum performance improvements (#9484)

* perf: switch plaintext and json to one runtime per thread, improving performance.

* perf: remove need for additional vec

* perf: reduce length of query parameter

* perf: increase strength of inlining hint

* perf: reduce query length

* perf: shorten path and use references

* bug: increased length of route in line with requirements

* build: upgraded dependencies

* bug: increased mongodb pool size

* perf: replacement of moka  with quick_cache

* perf: simplification of mass update query

* perf: improvements to random number generation

* build: upgrade rust version

* fix: re-added ordering to remove deadlocks

* perf: use of mimalloc

frameworks/Rust/axum/Cargo.toml

@@ -39,48 +39,49 @@ simd-json = [
 ]
 
 [dependencies]
-axum = { version = "0.7.6", default-features = false, features = [
+axum = { version = "0.7.9", default-features = false, features = [
     "json",
     "query",
     "http1",
     "tokio",
 ] }
 deadpool = { version = "0.12.1", features = ["rt_tokio_1", "serde", "managed"] }
-deadpool-postgres = { version = "0.14.0", features = ["rt_tokio_1", "serde"] }
+deadpool-postgres = { version = "0.14.1", features = ["rt_tokio_1", "serde"] }
 dotenv = "0.15.0"
-futures = "0.3.30"
-futures-util = "0.3.30"
-mongodb = { version = "2.8.0", features = [
+futures = "0.3.31"
+futures-util = "0.3.31"
+mongodb = { version = "3.1.1", features = [
     "zstd-compression",
     "snappy-compression",
     "zlib-compression",
 ] }
 num_cpus = "1.16.0"
 rand = { version = "0.8.5", features = ["small_rng"] }
-serde = { version = "1.0.196", features = ["derive"] }
-serde_json = "1.0.127"
-sqlx = { version = "0.7.3", features = [
+serde = { version = "1.0.216", features = ["derive"] }
+serde_json = "1.0.134"
+sqlx = { version = "0.8.2", features = [
     "postgres",
     "macros",
     "runtime-tokio",
     "tls-rustls",
 ] }
-tokio = { version = "1.39.3", features = ["full"] }
+tokio = { version = "1.42.0", features = ["full"] }
 tokio-pg-mapper = { version = "0.2.0" }
 tokio-pg-mapper-derive = { version = "0.2.0" }
-tokio-postgres = { version = "0.7.11" }
-tower = { version = "0.5.0", features = ["util"] }
-tower-http = { version = "0.5.2", features = ["set-header"] }
+tokio-postgres = { version = "0.7.12" }
+tower = { version = "0.5.2", features = ["util"] }
+tower-http = { version = "0.6.2", features = ["set-header"] }
 yarte = "0.15.7"
-simd-json = { version = "0.13.8", optional = true }
-axum-core = { version = "0.4.3", optional = true }
+simd-json = { version = "0.14.3", optional = true }
+axum-core = { version = "0.4.5", optional = true }
 mime = { version = "0.3.17", optional = true }
-bytes = { version = "1.5.0", optional = true }
-serde_path_to_error = { version = "0.1.15", optional = true }
-moka = { version = "0.12.8", features = ["future"] }
-socket2 = "0.5.7"
-hyper = { version = "1.4", features = ["server", "http1"] }
+bytes = { version = "1.9.0", optional = true }
+serde_path_to_error = { version = "0.1.16", optional = true }
+socket2 = "0.5.8"
+hyper = { version = "1.5", features = ["server", "http1"] }
 hyper-util = { version = "0.1", features = ["tokio", "server-auto", "http1"] }
+quick_cache = "0.6.9"
+mimalloc = "0.1.43"
 
 
 [profile.release]

frameworks/Rust/axum/README.md

@@ -27,10 +27,11 @@ built with Tokio, Tower, and Hyper.
 ## Notable Points (both performance and build)
 
 - Use of `async`.
-- Use of most recent versions of Rust, `axum` and dependencies.
+- Use of the most recent versions of Rust, `axum` and dependencies.
 - (Disabled by default) Compile-time swap-in of `simd-json` instead of `serde_json` for faster JSON serialization.
 - Release binaries are stripped and compiled with CPU native.
-- Sockets configured with TCP_NODELAY and to support an increased number of pending connections.
+- Sockets configured with `TCP_NODELAY` and to support an increased number of pending connections.
+- For very simple benchmarks, use of a separate, single-threaded Tokio runtime for each thread.
 - Server configured to serve HTTP/1 only, with no need for websockets.
 - Separation of build and deployment containers using multi-stage builds.
 - Deployment into Google's minimal `distroless-cc` container.
@@ -39,8 +40,5 @@ built with Tokio, Tower, and Hyper.
 - Use of PostgreSQL prepared statements cache (where supported).
 - Use of PostgreSQL arrays to execute multi-row database updates with a single `UPDATE` query.
   - This is permitted by the [test requirements](https://github.com/TechEmpower/FrameworkBenchmarks/wiki/Project-Information-Framework-Tests-Overview#database-updates), step (ix).
-- In version 0.7.6 (as yet unreleased), a native API to set TCP_NODELAY will be included.
-  - https://github.com/tokio-rs/axum/pull/2653/
-  - https://github.com/tokio-rs/axum/issues/2521
 - More performance improvements are to be expected in version 0.8:
   - https://github.com/tokio-rs/axum/issues/1827

frameworks/Rust/axum/axum.dockerfile

@@ -1,4 +1,4 @@
-FROM docker.io/rust:1.80-slim-bookworm AS builder
+FROM docker.io/rust:1.83-slim-bookworm AS builder
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
     pkg-config libssl-dev \
@@ -18,7 +18,7 @@ ENV POSTGRES_MIN_POOL_SIZE=56
 ENV POSTGRES_MAX_POOL_SIZE=56
 ENV MONGODB_URL=mongodb://tfb-database:27017
 ENV MONGODB_MIN_POOL_SIZE=28
-ENV MONGODB_MAX_POOL_SIZE=14
+ENV MONGODB_MAX_POOL_SIZE=28
 COPY --from=builder /build/target/release/axum* /app/
 EXPOSE 8000
 CMD ["/app/axum"]

frameworks/Rust/axum/src/common/mod.rs

@@ -17,10 +17,9 @@ pub const SELECT_WORLD_BY_ID: &str =
 pub const SELECT_ALL_CACHED_WORLDS: &str =
     "SELECT id, randomnumber FROM world ORDER BY id";
 #[allow(dead_code)]
-pub const UPDATE_WORLDS: &str = "WITH vals AS (SELECT * FROM UNNEST($1::int[], $2::int[]) AS v(id, rnum))
-    UPDATE world SET randomnumber = new.rnum FROM
-  (SELECT w.id, v.rnum FROM world w INNER JOIN vals v ON v.id = w.id ORDER BY w.id FOR UPDATE) AS new
-  WHERE world.id = new.id";
+pub const UPDATE_WORLDS: &str = r#"UPDATE world SET randomnumber = new.rnum FROM
+    (SELECT * FROM UNNEST($1::int[], $2::int[]) AS v(id, rnum) ORDER BY 1) AS new
+WHERE world.id = new.id"#;
 
 /// Return the value of an environment variable.
 #[allow(dead_code)]
@@ -41,11 +40,10 @@ pub fn random_id(rng: &mut SmallRng) -> i32 {
     rng.gen_range(1..10_001)
 }
 
-/// Generate vector of integers in the range 1 to 10,000 (inclusive)
+/// Generate an iterator of integers in the range 1 to 10,000 (inclusive)
 #[allow(dead_code)]
 #[inline(always)]
-pub fn random_ids(rng: &mut SmallRng, count: usize) -> Vec<i32> {
+pub fn random_ids(rng: &mut SmallRng, count: usize) -> impl Iterator<Item = i32> + use<'_> {
     rng.sample_iter(Uniform::new(1, 10_001))
         .take(count)
-        .collect()
-}
+}

frameworks/Rust/axum/src/main.rs

@@ -4,6 +4,10 @@ mod server;
 use axum::{http::StatusCode, response::IntoResponse, routing::get, Router};
 use common::models::Message;
 use dotenv::dotenv;
+use mimalloc::MiMalloc;
+
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
 
 #[cfg(not(feature = "simd-json"))]
 use axum::Json;

frameworks/Rust/axum/src/main_mongo.rs

@@ -14,8 +14,7 @@ use axum::Json;
 #[cfg(feature = "simd-json")]
 use common::simd_json::Json;
 use common::{
-    models::{FortuneInfo, World},
-    random_ids,
+    models::{FortuneInfo, World}, random_id
 };
 use dotenv::dotenv;
 use mongodb::{
@@ -24,6 +23,10 @@ use mongodb::{
 };
 use rand::{rngs::SmallRng, thread_rng, Rng, SeedableRng};
 use yarte::Template;
+use mimalloc::MiMalloc;
+
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
 
 use common::{
     get_env,
@@ -58,9 +61,7 @@ async fn queries(
     let q = parse_params(params);
 
     let mut rng = SmallRng::from_rng(&mut thread_rng()).unwrap();
-    let ids = random_ids(&mut rng, q);
-
-    let worlds = find_worlds(db, ids).await;
+    let worlds = find_worlds(db, &mut rng, q).await;
     let results = worlds.expect("worlds could not be retrieved");
 
     (StatusCode::OK, Json(results))
@@ -73,17 +74,14 @@ async fn updates(
     let q = parse_params(params);
 
     let mut rng = SmallRng::from_rng(&mut thread_rng()).unwrap();
-    let ids = random_ids(&mut rng, q);
 
-    let worlds = find_worlds(db.clone(), ids)
+    let worlds = find_worlds(db.clone(), &mut  rng, q)
         .await
         .expect("worlds could not be retrieved");
     let mut updated_worlds: Vec<World> = Vec::with_capacity(q);
 
     for mut world in worlds {
-        let random_number = (rng.gen::<u32>() % 10_000 + 1) as i32;
-
-        world.random_number = random_number;
+        world.random_number = random_id(&mut rng);
         updated_worlds.push(world);
     }
 

frameworks/Rust/axum/src/main_mongo_raw.rs

@@ -2,7 +2,7 @@ mod common;
 mod mongo_raw;
 mod server;
 
-use common::{models::World, random_id, random_ids};
+use common::{models::World, random_id};
 use mongo_raw::database::{
     find_world_by_id, find_worlds, update_worlds, DatabaseConnection,
 };
@@ -17,6 +17,11 @@ use axum::{
     extract::Query, http::StatusCode, response::IntoResponse, routing::get, Router,
 };
 
+use mimalloc::MiMalloc;
+
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
+
 #[cfg(not(feature = "simd-json"))]
 use axum::Json;
 #[cfg(feature = "simd-json")]
@@ -27,7 +32,7 @@ use mongodb::{
     options::{ClientOptions, Compressor},
     Client,
 };
-use rand::{rngs::SmallRng, thread_rng, Rng, SeedableRng};
+use rand::{rngs::SmallRng, thread_rng, SeedableRng};
 
 async fn db(DatabaseConnection(db): DatabaseConnection) -> impl IntoResponse {
     let mut rng = SmallRng::from_rng(&mut thread_rng()).unwrap();
@@ -48,9 +53,7 @@ async fn queries(
     let q = parse_params(params);
 
     let mut rng = SmallRng::from_rng(&mut thread_rng()).unwrap();
-    let ids = random_ids(&mut rng, q);
-
-    let worlds = find_worlds(db, ids).await;
+    let worlds = find_worlds(db, &mut rng, q).await;
     let results = worlds.expect("worlds could not be retrieved");
 
     (StatusCode::OK, Json(results))
@@ -64,16 +67,13 @@ async fn updates(
 
     let mut rng = SmallRng::from_rng(&mut thread_rng()).unwrap();
 
-    let ids = random_ids(&mut rng, q);
-    let worlds = find_worlds(db.clone(), ids)
+    let worlds = find_worlds(db.clone(), &mut rng, q)
         .await
         .expect("worlds could not be retrieved");
     let mut updated_worlds: Vec<World> = Vec::with_capacity(q);
 
     for mut world in worlds {
-        let random_number = (rng.gen::<u32>() % 10_000 + 1) as i32;
-
-        world.random_number = random_number;
+        world.random_number = random_id(&mut rng);
         updated_worlds.push(world);
     }
 

frameworks/Rust/axum/src/main_pg.rs

@@ -7,6 +7,10 @@ use axum::{
 use dotenv::dotenv;
 use rand::{rngs::SmallRng, thread_rng, SeedableRng};
 use yarte::Template;
+use mimalloc::MiMalloc;
+
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
 
 #[cfg(not(feature = "simd-json"))]
 use axum::Json;

frameworks/Rust/axum/src/main_pg_pool.rs

@@ -15,6 +15,10 @@ use dotenv::dotenv;
 use futures_util::{stream::FuturesUnordered, TryStreamExt};
 use rand::{rngs::SmallRng, thread_rng, SeedableRng};
 use yarte::Template;
+use mimalloc::MiMalloc;
+
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
 
 mod server;
 

frameworks/Rust/axum/src/main_sqlx.rs

@@ -12,10 +12,14 @@ use axum::{
     Router,
 };
 use dotenv::dotenv;
-use moka::future::Cache;
+use quick_cache::sync::Cache;
 use rand::{rngs::SmallRng, thread_rng, SeedableRng};
 use sqlx::models::World;
 use yarte::Template;
+use mimalloc::MiMalloc;
+
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
 
 #[cfg(not(feature = "simd-json"))]
 use axum::Json;
@@ -55,10 +59,9 @@ async fn queries(
 ) -> impl IntoResponse {
     let mut rng = SmallRng::from_rng(&mut thread_rng()).unwrap();
     let count = parse_params(params);
-    let ids = random_ids(&mut rng, count);
     let mut worlds: Vec<World> = Vec::with_capacity(count);
 
-    for id in &ids {
+    for id in random_ids(&mut rng, count) {
         let world: World = ::sqlx::query_as(common::SELECT_WORLD_BY_ID)
             .bind(id)
             .fetch_one(&mut *db.acquire().await.unwrap())
@@ -98,10 +101,10 @@ async fn cache(
 ) -> impl IntoResponse {
     let count = parse_params(params);
     let mut rng = SmallRng::from_rng(&mut thread_rng()).unwrap();
-    let mut worlds: Vec<Option<Arc<World>>> = Vec::with_capacity(count);
-
+    let mut worlds: Vec<Option<World>> = Vec::with_capacity(count);
+    
     for id in random_ids(&mut rng, count) {
-        worlds.push(cache.get(&id).await);
+        worlds.push(cache.get(&id));
     }
 
     (StatusCode::OK, Json(worlds))
@@ -115,7 +118,7 @@ async fn preload_cache(AppState { db, cache }: &AppState) {
         .expect("error loading worlds");
 
     for world in worlds {
-        cache.insert(world.id, Arc::new(world)).await;
+        cache.insert(world.id, world);
     }
 }
 
@@ -123,7 +126,7 @@ async fn preload_cache(AppState { db, cache }: &AppState) {
 #[derive(Clone)]
 struct AppState {
     db: PgPool,
-    cache: Cache<i32, Arc<World>>,
+    cache: Arc<Cache<i32, World>>,
 }
 
 #[tokio::main]
@@ -136,10 +139,7 @@ async fn main() {
 
     let state = AppState {
         db: create_pool(database_url, max_pool_size, min_pool_size).await,
-        cache: Cache::builder()
-        .initial_capacity(10000)
-        .max_capacity(10000)
-        .build()
+        cache: Arc::new(Cache::new(10_000))
     };
 
     // Prime the cache with CachedWorld objects

frameworks/Rust/axum/src/mongo/database.rs

@@ -3,8 +3,9 @@ use std::{convert::Infallible, io};
 use axum::{async_trait, extract::FromRequestParts, http::request::Parts};
 use futures_util::{stream::FuturesUnordered, StreamExt, TryStreamExt};
 use mongodb::{bson::doc, Database};
+use rand::rngs::SmallRng;
 
-use crate::common::models::{Fortune, World};
+use crate::common::{models::{Fortune, World}, random_ids};
 
 pub struct DatabaseConnection(pub Database);
 
@@ -45,17 +46,17 @@ pub async fn find_world_by_id(db: Database, id: i32) -> Result<World, MongoError
     let filter = doc! { "_id": id as f32 };
 
     let world: World = world_collection
-        .find_one(Some(filter), None)
+        .find_one(filter)
         .await
         .unwrap()
         .expect("expected world, found none");
     Ok(world)
 }
 
-pub async fn find_worlds(db: Database, ids: Vec<i32>) -> Result<Vec<World>, MongoError> {
+pub async fn find_worlds(db: Database, rng: &mut SmallRng, count: usize) -> Result<Vec<World>, MongoError> {
     let future_worlds = FuturesUnordered::new();
 
-    for id in ids {
+    for id in random_ids(rng, count) {
         future_worlds.push(find_world_by_id(db.clone(), id));
     }
 
@@ -67,7 +68,7 @@ pub async fn fetch_fortunes(db: Database) -> Result<Vec<Fortune>, MongoError> {
     let fortune_collection = db.collection::<Fortune>("fortune");
 
     let mut fortune_cursor = fortune_collection
-        .find(None, None)
+        .find(doc! {})
         .await
         .expect("fortunes could not be loaded");
 
@@ -99,8 +100,7 @@ pub async fn update_worlds(
     }
 
     db.run_command(
-        doc! {"update": "world", "updates": updates, "ordered": false},
-        None,
+        doc! {"update": "world", "updates": updates, "ordered": false}
     )
     .await
     .expect("could not update worlds");

frameworks/Rust/axum/src/mongo_raw/database.rs

@@ -6,8 +6,9 @@ use mongodb::{
     bson::{doc, RawDocumentBuf},
     Database,
 };
+use rand::rngs::SmallRng;
 
-use crate::common::models::World;
+use crate::common::{models::World, random_ids};
 
 pub struct DatabaseConnection(pub Database);
 
@@ -48,7 +49,7 @@ pub async fn find_world_by_id(db: Database, id: i32) -> Result<World, MongoError
     let filter = doc! { "_id": id as f32 };
 
     let raw: RawDocumentBuf = world_collection
-        .find_one(Some(filter), None)
+        .find_one(filter)
         .await
         .unwrap()
         .expect("expected world, found none");
@@ -69,10 +70,10 @@ pub async fn find_world_by_id(db: Database, id: i32) -> Result<World, MongoError
     })
 }
 
-pub async fn find_worlds(db: Database, ids: Vec<i32>) -> Result<Vec<World>, MongoError> {
+pub async fn find_worlds(db: Database, rng: &mut SmallRng, count: usize) -> Result<Vec<World>, MongoError> {
     let future_worlds = FuturesUnordered::new();
 
-    for id in ids {
+    for id in random_ids(rng, count) {
         future_worlds.push(find_world_by_id(db.clone(), id));
     }
 
@@ -93,8 +94,7 @@ pub async fn update_worlds(
     }
 
     db.run_command(
-        doc! {"update": "world", "updates": updates, "ordered": false},
-        None,
+        doc! {"update": "world", "updates": updates, "ordered": false}
     )
     .await
     .expect("could not update worlds");