use axum::{
body::Body,
http::{Request, StatusCode},
};
use chrono::Utc;
use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
use qrusty::{
api::{ApiServer, StorageApi},
memory_storage::MemoryStorage,
message::{Message, Priority, PriorityOrdering, QueueConfig},
};
use serde_json::{json, Value};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use tower::ServiceExt;
const DEPTHS: &[usize] = &[0, 100, 1_000, 10_000];
const BATCH_SIZES: &[usize] = &[1, 10, 50, 100];
const BATCH_DEPTH: usize = 1_000;
const QUEUE_ALLOW_DUP: &str = "http_mem_allow_dup";
const QUEUE_NO_DUP: &str = "http_mem_no_dup";
const CONSUMER: &str = "bench-consumer";
static COUNTER: AtomicU64 = AtomicU64::new(0);
fn next_id() -> u64 {
COUNTER.fetch_add(1, Ordering::Relaxed)
}
async fn http(
api: &ApiServer,
method: &str,
path: &str,
body: Option<Value>,
) -> (StatusCode, Value) {
let app = api.clone().router();
let builder = Request::builder()
.method(method)
.uri(path)
.header("content-type", "application/json");
let request = match body {
Some(b) => builder.body(Body::from(b.to_string())).unwrap(),
None => builder.body(Body::empty()).unwrap(),
};
let response = app.oneshot(request).await.unwrap();
let status = response.status();
let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.unwrap();
let text = String::from_utf8(bytes.to_vec()).unwrap();
let json: Value = if text.is_empty() {
Value::Null
} else {
serde_json::from_str(&text).unwrap_or(Value::String(text))
};
(status, json)
}
fn make_storage_message(queue: &str, payload: &str, priority: u64) -> Message {
Message {
id: format!("http-mem-bench-{}", next_id()),
queue: queue.to_string(),
priority: Priority::Numeric(priority),
payload: payload.to_string(),
created_at: Utc::now(),
locked_until: None,
locked_by: None,
retry_count: 0,
max_retries: 10,
payload_ref: None,
payload_hash: None,
}
}
fn setup_api_memory(
rt: &tokio::runtime::Runtime,
queue: &str,
allow_duplicates: bool,
depth: usize,
) -> ApiServer {
let storage = Arc::new(MemoryStorage::new()) as Arc<dyn StorageApi>;
rt.block_on(async {
storage
.create_queue(
queue,
QueueConfig {
ordering: PriorityOrdering::MaxFirst,
allow_duplicates,
..Default::default()
},
)
.await;
for i in 0..depth {
let msg = make_storage_message(queue, &format!("background-{}", i), 100);
storage.push(msg).await.expect("pre-populate push");
}
});
ApiServer::new(storage)
}
fn bench_http_mem_publish_allow_duplicates(c: &mut Criterion) {
let rt = tokio::runtime::Runtime::new().unwrap();
let mut group = c.benchmark_group("http_memory/publish/allow_duplicates");
group.sample_size(200);
for &depth in DEPTHS {
let api = setup_api_memory(&rt, QUEUE_ALLOW_DUP, true, depth);
group.bench_with_input(
BenchmarkId::new("queue_depth", depth),
&depth,
|b, _depth| {
b.iter(|| {
let n = next_id();
rt.block_on(http(
&api,
"POST",
"/publish",
Some(json!({
"queue": QUEUE_ALLOW_DUP,
"priority": 50,
"payload": format!("payload-{}", n),
})),
))
});
},
);
}
group.finish();
}
fn bench_http_mem_publish_no_duplicates(c: &mut Criterion) {
let rt = tokio::runtime::Runtime::new().unwrap();
let mut group = c.benchmark_group("http_memory/publish/no_duplicates");
group.sample_size(200);
for &depth in DEPTHS {
let api = setup_api_memory(&rt, QUEUE_NO_DUP, false, depth);
group.bench_with_input(
BenchmarkId::new("queue_depth", depth),
&depth,
|b, _depth| {
b.iter(|| {
let n = next_id();
rt.block_on(http(
&api,
"POST",
"/publish",
Some(json!({
"queue": QUEUE_NO_DUP,
"priority": 50,
"payload": format!("unique-{}", n),
})),
))
});
},
);
}
group.finish();
}
fn bench_http_mem_ack(c: &mut Criterion) {
let rt = tokio::runtime::Runtime::new().unwrap();
let mut group = c.benchmark_group("http_memory/ack");
group.sample_size(200);
for &depth in DEPTHS {
let api = setup_api_memory(&rt, QUEUE_ALLOW_DUP, true, depth);
group.bench_with_input(
BenchmarkId::new("queue_depth", depth),
&depth,
|b, _depth| {
b.iter_batched(
|| {
rt.block_on(async {
let n = next_id();
let (status, body) = http(
&api,
"POST",
"/publish",
Some(json!({
"queue": QUEUE_ALLOW_DUP,
"priority": 1,
"payload": format!("ack-target-{}", n),
})),
)
.await;
assert_eq!(status, StatusCode::OK, "publish failed: {:?}", body);
let (status, body) = http(
&api,
"POST",
&format!("/consume/{}", QUEUE_ALLOW_DUP),
Some(json!({
"consumer_id": CONSUMER,
"timeout_seconds": 3600,
})),
)
.await;
assert_eq!(status, StatusCode::OK, "consume failed: {:?}", body);
body["id"].as_str().unwrap().to_string()
})
},
|msg_id| {
rt.block_on(http(
&api,
"POST",
&format!("/ack/{}/{}", QUEUE_ALLOW_DUP, msg_id),
Some(json!({ "consumer_id": CONSUMER })),
))
},
BatchSize::SmallInput,
);
},
);
}
group.finish();
}
fn bench_http_mem_nack_allow_duplicates(c: &mut Criterion) {
let rt = tokio::runtime::Runtime::new().unwrap();
let mut group = c.benchmark_group("http_memory/nack/allow_duplicates");
group.sample_size(200);
for &depth in DEPTHS {
let api = setup_api_memory(&rt, QUEUE_ALLOW_DUP, true, depth);
group.bench_with_input(
BenchmarkId::new("queue_depth", depth),
&depth,
|b, _depth| {
b.iter_batched(
|| {
rt.block_on(async {
let n = next_id();
let (status, body) = http(
&api,
"POST",
"/publish",
Some(json!({
"queue": QUEUE_ALLOW_DUP,
"priority": 1,
"payload": format!("nack-target-{}", n),
})),
)
.await;
assert_eq!(status, StatusCode::OK, "publish failed: {:?}", body);
let (status, body) = http(
&api,
"POST",
&format!("/consume/{}", QUEUE_ALLOW_DUP),
Some(json!({
"consumer_id": CONSUMER,
"timeout_seconds": 3600,
})),
)
.await;
assert_eq!(status, StatusCode::OK, "consume failed: {:?}", body);
body["id"].as_str().unwrap().to_string()
})
},
|msg_id| {
rt.block_on(http(
&api,
"POST",
&format!("/nack/{}/{}", QUEUE_ALLOW_DUP, msg_id),
Some(json!({ "consumer_id": CONSUMER })),
))
},
BatchSize::SmallInput,
);
},
);
}
group.finish();
}
fn bench_http_mem_nack_no_duplicates(c: &mut Criterion) {
let rt = tokio::runtime::Runtime::new().unwrap();
let mut group = c.benchmark_group("http_memory/nack/no_duplicates");
group.sample_size(200);
for &depth in DEPTHS {
let api = setup_api_memory(&rt, QUEUE_NO_DUP, false, depth);
group.bench_with_input(
BenchmarkId::new("queue_depth", depth),
&depth,
|b, _depth| {
b.iter_batched(
|| {
rt.block_on(async {
let n = next_id();
let (status, body) = http(
&api,
"POST",
"/publish",
Some(json!({
"queue": QUEUE_NO_DUP,
"priority": 1,
"payload": format!("nack-unique-{}", n),
})),
)
.await;
assert_eq!(status, StatusCode::OK, "publish failed: {:?}", body);
let (status, body) = http(
&api,
"POST",
&format!("/consume/{}", QUEUE_NO_DUP),
Some(json!({
"consumer_id": CONSUMER,
"timeout_seconds": 3600,
})),
)
.await;
assert_eq!(status, StatusCode::OK, "consume failed: {:?}", body);
body["id"].as_str().unwrap().to_string()
})
},
|msg_id| {
rt.block_on(http(
&api,
"POST",
&format!("/nack/{}/{}", QUEUE_NO_DUP, msg_id),
Some(json!({ "consumer_id": CONSUMER })),
))
},
BatchSize::SmallInput,
);
},
);
}
group.finish();
}
fn bench_http_mem_batch_ack(c: &mut Criterion) {
let rt = tokio::runtime::Runtime::new().unwrap();
let mut group = c.benchmark_group("http_memory/batch_ack");
group.sample_size(50);
for &batch_size in BATCH_SIZES {
let api = setup_api_memory(&rt, QUEUE_ALLOW_DUP, true, BATCH_DEPTH);
group.bench_with_input(
BenchmarkId::new("batch_size", batch_size),
&batch_size,
|b, &bs| {
b.iter_batched(
|| {
rt.block_on(async {
let mut ids = Vec::with_capacity(bs);
for _ in 0..bs {
let n = next_id();
let (status, body) = http(
&api,
"POST",
"/publish",
Some(json!({
"queue": QUEUE_ALLOW_DUP,
"priority": 1,
"payload": format!("batch-ack-{}", n),
})),
)
.await;
assert_eq!(status, StatusCode::OK, "publish: {:?}", body);
let (status, body) = http(
&api,
"POST",
&format!("/consume/{}", QUEUE_ALLOW_DUP),
Some(json!({
"consumer_id": CONSUMER,
"timeout_seconds": 3600,
})),
)
.await;
assert_eq!(status, StatusCode::OK, "consume: {:?}", body);
ids.push(body["id"].as_str().unwrap().to_string());
}
ids
})
},
|ids| {
rt.block_on(http(
&api,
"POST",
&format!("/ack-batch/{}", QUEUE_ALLOW_DUP),
Some(json!({
"consumer_id": CONSUMER,
"message_ids": ids,
})),
))
},
BatchSize::SmallInput,
);
},
);
}
group.finish();
}
fn bench_http_mem_batch_nack_allow_duplicates(c: &mut Criterion) {
let rt = tokio::runtime::Runtime::new().unwrap();
let mut group = c.benchmark_group("http_memory/batch_nack/allow_duplicates");
group.sample_size(50);
for &batch_size in BATCH_SIZES {
let api = setup_api_memory(&rt, QUEUE_ALLOW_DUP, true, BATCH_DEPTH);
group.bench_with_input(
BenchmarkId::new("batch_size", batch_size),
&batch_size,
|b, &bs| {
b.iter_batched(
|| {
rt.block_on(async {
let mut ids = Vec::with_capacity(bs);
for _ in 0..bs {
let n = next_id();
let (status, body) = http(
&api,
"POST",
"/publish",
Some(json!({
"queue": QUEUE_ALLOW_DUP,
"priority": 1,
"payload": format!("batch-nack-dup-{}", n),
})),
)
.await;
assert_eq!(status, StatusCode::OK, "publish: {:?}", body);
let (status, body) = http(
&api,
"POST",
&format!("/consume/{}", QUEUE_ALLOW_DUP),
Some(json!({
"consumer_id": CONSUMER,
"timeout_seconds": 3600,
})),
)
.await;
assert_eq!(status, StatusCode::OK, "consume: {:?}", body);
ids.push(body["id"].as_str().unwrap().to_string());
}
ids
})
},
|ids| {
rt.block_on(http(
&api,
"POST",
&format!("/nack-batch/{}", QUEUE_ALLOW_DUP),
Some(json!({
"consumer_id": CONSUMER,
"message_ids": ids,
})),
))
},
BatchSize::SmallInput,
);
},
);
}
group.finish();
}
fn bench_http_mem_batch_nack_no_duplicates(c: &mut Criterion) {
let rt = tokio::runtime::Runtime::new().unwrap();
let mut group = c.benchmark_group("http_memory/batch_nack/no_duplicates");
group.sample_size(50);
for &batch_size in BATCH_SIZES {
let api = setup_api_memory(&rt, QUEUE_NO_DUP, false, BATCH_DEPTH);
group.bench_with_input(
BenchmarkId::new("batch_size", batch_size),
&batch_size,
|b, &bs| {
b.iter_batched(
|| {
rt.block_on(async {
let mut ids = Vec::with_capacity(bs);
for _ in 0..bs {
let n = next_id();
let (status, body) = http(
&api,
"POST",
"/publish",
Some(json!({
"queue": QUEUE_NO_DUP,
"priority": 1,
"payload": format!("batch-nack-nodup-{}", n),
})),
)
.await;
assert_eq!(status, StatusCode::OK, "publish: {:?}", body);
let (status, body) = http(
&api,
"POST",
&format!("/consume/{}", QUEUE_NO_DUP),
Some(json!({
"consumer_id": CONSUMER,
"timeout_seconds": 3600,
})),
)
.await;
assert_eq!(status, StatusCode::OK, "consume: {:?}", body);
ids.push(body["id"].as_str().unwrap().to_string());
}
ids
})
},
|ids| {
rt.block_on(http(
&api,
"POST",
&format!("/nack-batch/{}", QUEUE_NO_DUP),
Some(json!({
"consumer_id": CONSUMER,
"message_ids": ids,
})),
))
},
BatchSize::SmallInput,
);
},
);
}
group.finish();
}
criterion_group!(
benches,
bench_http_mem_publish_allow_duplicates,
bench_http_mem_publish_no_duplicates,
bench_http_mem_ack,
bench_http_mem_nack_allow_duplicates,
bench_http_mem_nack_no_duplicates,
bench_http_mem_batch_ack,
bench_http_mem_batch_nack_allow_duplicates,
bench_http_mem_batch_nack_no_duplicates,
);
criterion_main!(benches);