use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main};
use std::time::Duration;
use tempfile::TempDir;
use parking_lot::RwLock;
use sqlitegraph::backend::native::v3::btree::BTreeManager;
use sqlitegraph::backend::native::v3::{PageAllocator, PersistentHeaderV3, V3Backend};
use sqlitegraph::backend::{BackendDirection, EdgeSpec, GraphBackend, NeighborQuery, NodeSpec};
use std::sync::Arc;
fn bench_allocator_allocate(c: &mut Criterion) {
let mut group = c.benchmark_group("allocator/allocate");
group.measurement_time(Duration::from_secs(5));
group.sample_size(50);
for size in [100, 1_000, 10_000] {
group.throughput(Throughput::Elements(size as u64));
group.bench_with_input(BenchmarkId::new("sequential", size), &size, |b, &size| {
b.iter(|| {
let header = PersistentHeaderV3::new_v3();
let mut allocator = PageAllocator::new(&header);
for _ in 0..size {
black_box(allocator.allocate().unwrap());
}
});
});
}
group.finish();
}
fn bench_allocator_allocate_deallocate_reuse(c: &mut Criterion) {
let mut group = c.benchmark_group("allocator/reuse");
group.measurement_time(Duration::from_secs(5));
group.bench_function("allocate_1000_dealloc_500_realloc_500", |b| {
b.iter(|| {
let header = PersistentHeaderV3::new_v3();
let mut allocator = PageAllocator::new(&header);
let mut pages: Vec<u64> = Vec::with_capacity(1000);
for _ in 0..1000 {
pages.push(allocator.allocate().unwrap());
}
for i in 0..500 {
allocator.deallocate(pages[i]).unwrap();
}
for _ in 0..500 {
let reused = allocator.allocate().unwrap();
black_box(reused);
}
});
});
group.finish();
}
fn bench_btree_insert_lookup(c: &mut Criterion) {
let mut group = c.benchmark_group("btree/insert_lookup");
group.measurement_time(Duration::from_secs(5));
group.sample_size(30);
let temp = TempDir::new().unwrap();
let db_path = temp.path().join("bench.db");
std::fs::File::create(&db_path).unwrap();
for size in [100, 1_000, 10_000] {
group.throughput(Throughput::Elements(size as u64));
group.bench_with_input(BenchmarkId::new("insert", size), &size, |b, &size| {
b.iter(|| {
let header = PersistentHeaderV3::new_v3();
let allocator = Arc::new(RwLock::new(PageAllocator::new(&header)));
let mut btree = BTreeManager::new(allocator, None, db_path.clone());
for i in 0..size {
black_box(btree.insert(i as i64, (i + 1) as u64).unwrap());
}
});
});
let header = PersistentHeaderV3::new_v3();
let allocator = Arc::new(RwLock::new(PageAllocator::new(&header)));
let mut btree = BTreeManager::new(allocator, None, db_path.clone());
for i in 0..size {
btree.insert(i as i64, (i + 1) as u64).unwrap();
}
group.throughput(Throughput::Elements(size as u64));
group.bench_with_input(BenchmarkId::new("lookup", size), &size, |b, &size| {
b.iter(|| {
for i in 0..size {
black_box(btree.lookup(i as i64).unwrap());
}
});
});
}
group.finish();
}
fn bench_v3_insert_nodes(c: &mut Criterion) {
let mut group = c.benchmark_group("v3/insert_nodes");
group.measurement_time(Duration::from_secs(10));
group.sample_size(20);
for size in [100, 1_000] {
group.throughput(Throughput::Elements(size as u64));
group.bench_with_input(BenchmarkId::new("nodes", size), &size, |b, &size| {
b.iter_batched(
|| {
let temp = TempDir::new().unwrap();
let backend = V3Backend::create(temp.path().join("v3.db")).unwrap();
(backend, temp)
},
|(backend, _temp)| {
for i in 0..size {
black_box(
backend
.insert_node(NodeSpec {
kind: "Node".to_string(),
name: format!("node_{}", i),
file_path: None,
data: serde_json::json!({"id": i}),
})
.unwrap(),
);
}
},
criterion::BatchSize::LargeInput,
);
});
}
group.finish();
}
fn bench_v3_insert_edges(c: &mut Criterion) {
let mut group = c.benchmark_group("v3/insert_edges");
group.measurement_time(Duration::from_secs(10));
group.sample_size(20);
for (name, nodes) in [("small", 50), ("medium", 200)] {
group.throughput(Throughput::Elements(nodes as u64));
group.bench_with_input(BenchmarkId::new("edges", name), &nodes, |b, &nodes| {
b.iter_batched(
|| {
let temp = TempDir::new().unwrap();
let backend = V3Backend::create(temp.path().join("v3.db")).unwrap();
let mut node_ids = Vec::with_capacity(nodes);
for i in 0..nodes {
node_ids.push(
backend
.insert_node(NodeSpec {
kind: "Node".to_string(),
name: format!("n_{}", i),
file_path: None,
data: serde_json::json!({}),
})
.unwrap(),
);
}
(backend, node_ids, temp)
},
|(backend, node_ids, _temp)| {
for i in 0..node_ids.len() - 1 {
black_box(
backend
.insert_edge(EdgeSpec {
from: node_ids[i],
to: node_ids[i + 1],
edge_type: "NEXT".to_string(),
data: serde_json::json!({}),
})
.unwrap(),
);
}
},
criterion::BatchSize::LargeInput,
);
});
}
group.finish();
}
fn bench_v3_get_neighbors(c: &mut Criterion) {
let mut group = c.benchmark_group("v3/get_neighbors");
group.measurement_time(Duration::from_secs(5));
group.sample_size(50);
let temp = TempDir::new().unwrap();
let backend = V3Backend::create(temp.path().join("v3.db")).unwrap();
let center = backend
.insert_node(NodeSpec {
kind: "Node".to_string(),
name: "center".to_string(),
file_path: None,
data: serde_json::json!({}),
})
.unwrap();
let mut leaves = Vec::new();
for i in 0..100 {
let leaf = backend
.insert_node(NodeSpec {
kind: "Node".to_string(),
name: format!("leaf_{}", i),
file_path: None,
data: serde_json::json!({}),
})
.unwrap();
backend
.insert_edge(EdgeSpec {
from: center,
to: leaf,
edge_type: "CONNECTS".to_string(),
data: serde_json::json!({}),
})
.unwrap();
leaves.push(leaf);
}
let snapshot = sqlitegraph::snapshot::SnapshotId::current();
group.bench_function("star_100_outgoing", |b| {
b.iter(|| {
black_box(
backend
.neighbors(
snapshot,
center,
NeighborQuery {
direction: BackendDirection::Outgoing,
edge_type: None,
},
)
.unwrap(),
);
});
});
group.bench_function("leaf_incoming", |b| {
b.iter(|| {
black_box(
backend
.neighbors(
snapshot,
leaves[0],
NeighborQuery {
direction: BackendDirection::Incoming,
edge_type: None,
},
)
.unwrap(),
);
});
});
group.finish();
}
criterion_group!(
benches,
bench_allocator_allocate,
bench_allocator_allocate_deallocate_reuse,
bench_btree_insert_lookup,
bench_v3_insert_nodes,
bench_v3_insert_edges,
bench_v3_get_neighbors,
);
criterion_main!(benches);