use commonware_allocator::mimalloc;
use commonware_runtime::{
Metrics, Name, Supervisor,
telemetry::metrics::{Metric, Registered},
};
use commonware_storage::{
index::{Unordered, ordered, partitioned, unordered},
translator::{Cap, EightCap},
};
use commonware_utils::{TestRng, Widen};
use rand::Rng;
use std::{
hint::black_box,
time::{Duration, Instant},
};
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc::new(mimalloc::Config::DEFAULT);
const DEFAULT_ITEMS: [u64; 2] = [20_000_000, 100_000_000];
const DEFAULT_BATCH: usize = 10_000;
const SEED: u64 = 0;
#[derive(Clone)]
struct DummyMetrics;
impl Supervisor for DummyMetrics {
fn child(&self, _: &'static str) -> Self {
Self
}
fn with_attribute(self, _: &'static str, _: impl std::fmt::Display) -> Self {
Self
}
fn name(&self) -> Name {
Name::default()
}
}
impl Metrics for DummyMetrics {
fn register<N: Into<String>, H: Into<String>, M: Metric>(
&self,
_: N,
_: H,
metric: M,
) -> Registered<M> {
Registered::detached(metric)
}
fn encode(&self) -> String {
String::new()
}
}
fn measure<I: Unordered<Value = u64>>(
mut index: I,
items: u64,
batch: usize,
) -> (Duration, Duration, Duration) {
let mut rng = TestRng::new(SEED);
let start = Instant::now();
for value in 0..items {
index.insert(&rng.next_u64().to_be_bytes(), value);
}
let insert = start.elapsed();
let mut rng = TestRng::new(SEED);
let mut sum = 0u64;
let start = Instant::now();
for _ in 0..items {
for value in index.get(&rng.next_u64().to_be_bytes()) {
sum = sum.wrapping_add(*value);
}
}
let lookup = start.elapsed();
black_box(sum);
let mut rng = TestRng::new(SEED);
let mut keys = Vec::with_capacity(batch);
let mut many = Duration::ZERO;
let mut remaining = items;
while remaining > 0 {
let n = remaining.min(Widen::widen(batch));
keys.clear();
keys.extend((0..n).map(|_| rng.next_u64().to_be_bytes()));
keys.sort_unstable();
let mut sum = 0u64;
let start = Instant::now();
index.get_many(&keys, |_, value| sum = sum.wrapping_add(*value));
many += start.elapsed();
black_box(sum);
remaining -= n;
}
(insert, lookup, many)
}
fn main() {
let argv: Vec<String> = std::env::args().skip(1).collect();
let args: Vec<u64> = argv.iter().filter_map(|a| a.parse().ok()).collect();
let batch =
argv.iter()
.find_map(|a| a.strip_prefix("batch="))
.map_or(DEFAULT_BATCH, |b| match b.parse() {
Ok(batch) if batch > 0 => batch,
_ => panic!("batch size must be a positive integer, got {b:?}"),
});
let only: Vec<String> = argv
.into_iter()
.filter(|a| a.parse::<u64>().is_err() && !a.starts_with('-') && !a.starts_with("batch="))
.collect();
let sizes = if !args.is_empty() {
args
} else if cfg!(huge_bench) {
DEFAULT_ITEMS.to_vec()
} else {
eprintln!(
"index_scale is opt-in; pass key counts (e.g. `-- 500000000`), or build with \
`--cfg huge_bench` for the default 20M/100M tier"
);
return;
};
for items in sizes {
println!("index_scale: items={items}");
macro_rules! run {
($name:literal, $index:expr) => {{
if only.is_empty() || only.iter().any(|v| v.as_str() == $name) {
let (insert, lookup, many) = measure($index, items, batch);
println!(
" {:<32} insert={} ns/op lookup={} ns/op get_many(batch={batch})={} ns/op",
$name,
insert.as_nanos() / u128::from(items),
lookup.as_nanos() / u128::from(items),
many.as_nanos() / u128::from(items),
);
}
}};
}
run!(
"partitioned_ordered_3",
partitioned::ordered::Index::<_, _, 3>::new(DummyMetrics, Cap::<5>::new())
);
run!(
"partitioned_ordered_3_eightcap",
partitioned::ordered::Index::<_, _, 3>::new(DummyMetrics, EightCap)
);
run!("unordered", unordered::Index::new(DummyMetrics, EightCap));
run!(
"partitioned_unordered_1",
partitioned::unordered::Index::<_, _, 1>::new(DummyMetrics, Cap::<7>::new())
);
run!(
"partitioned_unordered_2",
partitioned::unordered::Index::<_, _, 2>::new(DummyMetrics, Cap::<6>::new())
);
run!("ordered", ordered::Index::new(DummyMetrics, EightCap));
}
}