This commit is contained in:
2026-05-26 23:22:32 +02:00
parent 842fca2b72
commit 730de334e0
17 changed files with 5380 additions and 0 deletions
+585
View File
@@ -0,0 +1,585 @@
//! axum-server — the hyper/axum side of the urus / axum / cowboy benchmark.
//!
//! Exposes the same routes as urus-server with materially the same
//! middleware stack:
//! logger (ring-buffer, no stdout) -> request_id -> auth -> router
//!
//! Backing store is switchable:
//! --store=memory Mutex<HashMap<u64, User>>
//! --store=sqlite sqlx::SqlitePool in WAL mode
//!
//! Routes:
//! GET /ping naked (bypasses auth)
//! GET /api/v1/users list (cap 100)
//! POST /api/v1/users create
//! GET /api/v1/users/:id get
//! PUT /api/v1/users/:id update
//! DELETE /api/v1/users/:id delete
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Instant;
use axum::{
body::Bytes,
extract::{Path, State},
http::{header, HeaderValue, Request, StatusCode},
middleware::{self, Next as AxumNext},
response::Response,
routing::get,
Router,
};
use serde::{Deserialize, Serialize};
use sqlx::{sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}, SqlitePool};
use subtle::ConstantTimeEq;
use tokio::sync::Mutex;
// ---------------------------------------------------------------------------
// CLI (hand-rolled, same shape as urus-server)
// ---------------------------------------------------------------------------
#[derive(Clone, Debug)]
struct Cli {
addr: String,
store: StoreKind,
token: String,
db_path: String,
no_auth: bool,
}
#[derive(Clone, Copy, Debug, PartialEq)]
enum StoreKind { Memory, Sqlite }
fn parse_cli() -> Result<Cli, String> {
let mut cli = Cli {
addr: "127.0.0.1:8080".into(),
store: StoreKind::Memory,
token: "test-token-aaaaaaaaaaaaaaaaaaaa".into(),
db_path: "/tmp/axum-bench.sqlite".into(),
no_auth: false,
};
for arg in std::env::args().skip(1) {
let (k, v) = match arg.split_once('=') {
Some((k, v)) => (k.to_string(), Some(v.to_string())),
None => (arg.clone(), None),
};
match (k.as_str(), v) {
("--addr", Some(v)) => cli.addr = v,
("--token", Some(v)) => cli.token = v,
("--db-path", Some(v)) => cli.db_path = v,
("--store", Some(v)) => cli.store = match v.as_str() {
"memory" => StoreKind::Memory,
"sqlite" => StoreKind::Sqlite,
other => return Err(format!("unknown store: {other}")),
},
("--no-auth", None) => cli.no_auth = true,
("--help" | "-h", _) => { print_usage(); std::process::exit(0); }
(k, _) => return Err(format!("unknown flag: {k}")),
}
}
Ok(cli)
}
fn print_usage() {
eprintln!("axum-server [--addr=HOST:PORT] [--store=memory|sqlite] \
[--token=TOKEN] [--db-path=PATH] [--no-auth]");
}
// ---------------------------------------------------------------------------
// Model
// ---------------------------------------------------------------------------
#[derive(Clone, Debug, Serialize, Deserialize, sqlx::FromRow)]
struct User {
id: i64,
name: String,
email: String,
}
#[derive(Deserialize)]
struct NewUser {
name: String,
email: String,
}
// ---------------------------------------------------------------------------
// Store trait — memory and sqlite implementations behind one type.
// ---------------------------------------------------------------------------
//
// We use an enum rather than a trait object so handler code stays
// monomorphic. Both variants are async (sqlx is genuinely async; the
// memory variant uses a tokio Mutex for parity — std::sync::Mutex inside
// an async handler is a known footgun, even though for our short
// critical sections it'd be measurably faster).
#[derive(Clone)]
enum Store {
Memory(Arc<MemoryStore>),
Sqlite(SqlitePool),
}
struct MemoryStore {
map: Mutex<HashMap<u64, User>>,
next_id: AtomicU64,
}
impl MemoryStore {
fn new() -> Self {
Self {
map: Mutex::new(HashMap::with_capacity(1024)),
next_id: AtomicU64::new(1),
}
}
}
async fn open_sqlite(path: &str) -> Result<SqlitePool, sqlx::Error> {
let opts = SqliteConnectOptions::new()
.filename(path)
.create_if_missing(true)
.journal_mode(SqliteJournalMode::Wal)
.busy_timeout(std::time::Duration::from_secs(5));
// Pool sized to taste; 8 connections matches urus-side's 1 writer + 4
// readers plus headroom for axum's tokio worker count.
let pool = SqlitePoolOptions::new()
.max_connections(8)
.connect_with(opts)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS users (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT NOT NULL,
email TEXT NOT NULL
);"
).execute(&pool).await?;
Ok(pool)
}
// ---------------------------------------------------------------------------
// AppState
// ---------------------------------------------------------------------------
#[derive(Clone)]
struct AppState {
store: Store,
log_ring: LogRing,
}
// ---------------------------------------------------------------------------
// Bounded log ring — analogue of urus-server's LogRing.
// ---------------------------------------------------------------------------
//
// We can't easily use the same crossbeam ArrayQueue type because axum's
// middleware closures aren't `&` (`tokio::sync::Mutex` is the wrong
// shape here — we want lock-free push). A simple wraparound slot-array
// via atomics would be ideal, but the simplest correct thing that
// matches urus's behavioural guarantee (bounded, drops on overflow,
// never blocks) is `std::sync::Mutex<VecDeque>` with a capacity. The
// critical section is microseconds long.
#[derive(Clone)]
struct LogRing {
inner: Arc<std::sync::Mutex<std::collections::VecDeque<LogEntry>>>,
cap: usize,
}
#[derive(Clone, Debug)]
#[allow(dead_code)]
struct LogEntry {
method: String,
path: String,
status: u16,
elapsed_us: u64,
}
impl LogRing {
fn with_capacity(n: usize) -> Self {
Self {
inner: Arc::new(std::sync::Mutex::new(
std::collections::VecDeque::with_capacity(n))),
cap: n,
}
}
fn push(&self, e: LogEntry) {
if let Ok(mut q) = self.inner.lock() {
if q.len() >= self.cap { q.pop_front(); }
q.push_back(e);
}
}
}
// ---------------------------------------------------------------------------
// Middleware
// ---------------------------------------------------------------------------
//
// We hand-write three middleware functions rather than reaching for
// tower-http for everything. Reason: tower-http's auth and request-id
// layers are great, but for the bench we want behavioural parity with
// urus-server's plugs — same observable behaviour, same trivial work
// shape. Mixing tower-http for some and hand-rolled for others would
// confuse the comparison.
async fn logger_mw(
State(state): State<AppState>,
req: Request<axum::body::Body>,
next: AxumNext,
) -> Response {
let start = Instant::now();
let method = req.method().to_string();
let path = req.uri().path().to_string();
let response = next.run(req).await;
let elapsed_us = start.elapsed().as_micros() as u64;
state.log_ring.push(LogEntry {
method,
path,
status: response.status().as_u16(),
elapsed_us,
});
response
}
async fn request_id_mw(req: Request<axum::body::Body>, next: AxumNext) -> Response {
let mut raw = [0u8; 12];
fill_random_bytes(&mut raw);
let id = b32_encode(&raw);
let mut resp = next.run(req).await;
if let Ok(v) = HeaderValue::from_str(&id) {
resp.headers_mut().insert("x-request-id", v);
}
resp
}
fn b32_encode(input: &[u8; 12]) -> String {
static ALPHABET: &[u8; 32] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";
let mut out = String::with_capacity(20);
let mut buf: u64 = 0;
let mut bits: u32 = 0;
for &b in input {
buf = (buf << 8) | (b as u64);
bits += 8;
while bits >= 5 {
bits -= 5;
out.push(ALPHABET[((buf >> bits) & 0x1F) as usize] as char);
}
}
if bits > 0 {
out.push(ALPHABET[((buf << (5 - bits)) & 0x1F) as usize] as char);
}
out
}
// SplitMix64-based RNG, same as urus-server's. Kept identical so the
// request_id middleware costs are comparable.
thread_local! {
static RNG_STATE: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
}
fn splitmix64(s: u64) -> (u64, u64) {
let new_state = s.wrapping_add(0x9E3779B97F4A7C15);
let mut z = new_state;
z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
(new_state, z ^ (z >> 31))
}
fn fill_random_bytes(out: &mut [u8; 12]) {
RNG_STATE.with(|cell| {
let mut s = cell.get();
if s == 0 {
let p = cell as *const _ as usize as u64;
let t = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_nanos() as u64).unwrap_or(0);
s = p ^ t ^ 0xA5A5_A5A5_A5A5_A5A5;
if s == 0 { s = 1; }
}
let (s1, r1) = splitmix64(s);
let (s2, r2) = splitmix64(s1);
cell.set(s2);
out[0..8].copy_from_slice(&r1.to_le_bytes());
out[8..12].copy_from_slice(&r2.to_le_bytes()[..4]);
});
}
async fn auth_mw(
State(token): State<Arc<Vec<u8>>>,
req: Request<axum::body::Body>,
next: AxumNext,
) -> Response {
// /ping bypass, same exception as urus-server (see middleware.rs in
// that crate for the rationale).
if req.uri().path() == "/ping" {
return next.run(req).await;
}
let presented = req.headers()
.get(header::AUTHORIZATION)
.and_then(|h| h.to_str().ok())
.and_then(|s| s.strip_prefix("Bearer "))
.map(str::as_bytes);
let ok = match presented {
Some(p) if p.len() == token.len() => bool::from(p.ct_eq(&token)),
_ => false,
};
if ok {
next.run(req).await
} else {
let body = b"{\"error\":\"unauthorized\"}";
Response::builder()
.status(StatusCode::UNAUTHORIZED)
.header(header::CONTENT_TYPE, "application/json")
.body(axum::body::Body::from(&body[..]))
.unwrap()
}
}
// ---------------------------------------------------------------------------
// Handlers
// ---------------------------------------------------------------------------
async fn h_ping() -> (StatusCode, [(header::HeaderName, &'static str); 1], &'static str) {
(StatusCode::OK, [(header::CONTENT_TYPE, "text/plain; charset=utf-8")], "pong")
}
async fn h_list(State(state): State<AppState>) -> Response {
let users: Vec<User> = match &state.store {
Store::Memory(m) => {
let g = m.map.lock().await;
g.values().take(100).cloned().collect()
}
Store::Sqlite(p) => {
match sqlx::query_as::<_, User>("SELECT id, name, email FROM users LIMIT 100")
.fetch_all(p).await {
Ok(v) => v,
Err(_) => return db_error(),
}
}
};
json_response(StatusCode::OK, &users)
}
async fn h_get(
State(state): State<AppState>,
Path(id): Path<u64>,
) -> Response {
match &state.store {
Store::Memory(m) => {
let g = m.map.lock().await;
match g.get(&id) {
Some(u) => json_response(StatusCode::OK, u),
None => not_found(),
}
}
Store::Sqlite(p) => {
match sqlx::query_as::<_, User>("SELECT id, name, email FROM users WHERE id=?1")
.bind(id as i64).fetch_optional(p).await {
Ok(Some(u)) => json_response(StatusCode::OK, &u),
Ok(None) => not_found(),
Err(_) => db_error(),
}
}
}
}
async fn h_create(
State(state): State<AppState>,
body: Bytes,
) -> Response {
let nu: NewUser = match serde_json::from_slice(&body) {
Ok(v) => v,
Err(_) => return bad_request(),
};
match &state.store {
Store::Memory(m) => {
let id = m.next_id.fetch_add(1, Ordering::Relaxed);
let u = User { id: id as i64, name: nu.name, email: nu.email };
let mut g = m.map.lock().await;
// bounded growth same as urus-side
if g.len() > 100_000 {
let to_remove: Vec<u64> = g.keys().take(50_000).copied().collect();
for k in to_remove { g.remove(&k); }
}
g.insert(id, u.clone());
json_response(StatusCode::CREATED, &u)
}
Store::Sqlite(p) => {
let res = sqlx::query("INSERT INTO users (name, email) VALUES (?1, ?2)")
.bind(&nu.name).bind(&nu.email)
.execute(p).await;
match res {
Ok(r) => {
let u = User { id: r.last_insert_rowid(), name: nu.name, email: nu.email };
json_response(StatusCode::CREATED, &u)
}
Err(_) => db_error(),
}
}
}
}
async fn h_update(
State(state): State<AppState>,
Path(id): Path<u64>,
body: Bytes,
) -> Response {
let nu: NewUser = match serde_json::from_slice(&body) {
Ok(v) => v,
Err(_) => return bad_request(),
};
match &state.store {
Store::Memory(m) => {
let mut g = m.map.lock().await;
match g.get_mut(&id) {
Some(u) => {
u.name = nu.name;
u.email = nu.email;
let snap = u.clone();
json_response(StatusCode::OK, &snap)
}
None => not_found(),
}
}
Store::Sqlite(p) => {
let res = sqlx::query("UPDATE users SET name=?1, email=?2 WHERE id=?3")
.bind(&nu.name).bind(&nu.email).bind(id as i64)
.execute(p).await;
match res {
Ok(r) if r.rows_affected() > 0 => {
let u = User { id: id as i64, name: nu.name, email: nu.email };
json_response(StatusCode::OK, &u)
}
Ok(_) => not_found(),
Err(_) => db_error(),
}
}
}
}
async fn h_delete(
State(state): State<AppState>,
Path(id): Path<u64>,
) -> Response {
match &state.store {
Store::Memory(m) => {
let mut g = m.map.lock().await;
match g.remove(&id) {
Some(_) => Response::builder()
.status(StatusCode::NO_CONTENT)
.body(axum::body::Body::empty()).unwrap(),
None => not_found(),
}
}
Store::Sqlite(p) => {
let res = sqlx::query("DELETE FROM users WHERE id=?1")
.bind(id as i64).execute(p).await;
match res {
Ok(r) if r.rows_affected() > 0 => Response::builder()
.status(StatusCode::NO_CONTENT)
.body(axum::body::Body::empty()).unwrap(),
Ok(_) => not_found(),
Err(_) => db_error(),
}
}
}
}
// ---------------------------------------------------------------------------
// Response helpers
// ---------------------------------------------------------------------------
fn json_response<T: Serialize>(status: StatusCode, body: &T) -> Response {
match serde_json::to_vec(body) {
Ok(bytes) => Response::builder()
.status(status)
.header(header::CONTENT_TYPE, "application/json")
.body(axum::body::Body::from(bytes))
.unwrap(),
Err(_) => db_error(),
}
}
fn not_found() -> Response {
Response::builder()
.status(StatusCode::NOT_FOUND)
.header(header::CONTENT_TYPE, "application/json")
.body(axum::body::Body::from(&b"{\"error\":\"not found\"}"[..]))
.unwrap()
}
fn bad_request() -> Response {
Response::builder()
.status(StatusCode::BAD_REQUEST)
.header(header::CONTENT_TYPE, "application/json")
.body(axum::body::Body::from(&b"{\"error\":\"invalid body\"}"[..]))
.unwrap()
}
fn db_error() -> Response {
Response::builder()
.status(StatusCode::INTERNAL_SERVER_ERROR)
.header(header::CONTENT_TYPE, "application/json")
.body(axum::body::Body::from(&b"{\"error\":\"db\"}"[..]))
.unwrap()
}
// ---------------------------------------------------------------------------
// main
// ---------------------------------------------------------------------------
#[tokio::main(flavor = "multi_thread")]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let cli = match parse_cli() {
Ok(c) => c,
Err(e) => { eprintln!("axum-server: {e}"); print_usage(); std::process::exit(2); }
};
let addr: SocketAddr = cli.addr.parse()?;
let store = match cli.store {
StoreKind::Memory => Store::Memory(Arc::new(MemoryStore::new())),
StoreKind::Sqlite => Store::Sqlite(open_sqlite(&cli.db_path).await?),
};
let state = AppState {
store,
log_ring: LogRing::with_capacity(64 * 1024),
};
let token: Arc<Vec<u8>> = Arc::new(cli.token.as_bytes().to_vec());
// Build the router. /ping is on the same router; the auth middleware
// is path-aware and lets /ping through (see auth_mw).
let api_router: Router<AppState> = Router::new()
.route("/api/v1/users", get(h_list).post(h_create))
.route("/api/v1/users/:id", get(h_get).put(h_update).delete(h_delete))
.route("/ping", get(h_ping));
let mut app = api_router.with_state(state.clone());
// Apply middleware in inside-out order so the request flow ends up
// matching urus-server: logger (outermost) -> request_id -> auth.
// axum applies middleware nearest-the-handler first, so the layer
// call order is: auth then request_id then logger.
if !cli.no_auth {
app = app.layer(middleware::from_fn_with_state(token, auth_mw));
}
app = app
.layer(middleware::from_fn(request_id_mw))
.layer(middleware::from_fn_with_state(state, logger_mw));
eprintln!(
"axum-server: store={:?} addr={} auth={}",
cli.store, cli.addr, !cli.no_auth
);
let listener = tokio::net::TcpListener::bind(addr).await?;
axum::serve(listener, app).await?;
Ok(())
}