From 9633038a9ce9198cbed001bfabe1ac49739b2d9d Mon Sep 17 00:00:00 2001 From: tellet-q Date: Mon, 13 Jul 2026 10:16:55 +0200 Subject: [PATCH 1/2] feat: memory + on-disk size reporting Co-Authored-By: Claude Opus 4.8 (1M context) --- DEVELOPMENT.md | 1 + README.md | 23 +++ src/args/mod.rs | 5 + src/main.rs | 41 +++++ src/memory.rs | 232 ++++++++++++++++++++++++++ src/results.rs | 5 + tests/fixtures/collection_memory.json | 55 ++++++ 7 files changed, 362 insertions(+) create mode 100644 src/memory.rs create mode 100644 tests/fixtures/collection_memory.json diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index 511f21c..cac6169 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -13,6 +13,7 @@ src/ ├── upsert.rs # upsert request execution + timing ├── scroll.rs # scroll benchmarking processor ├── query.rs # search/scroll entry points +├── memory.rs # memory + on-disk size (REST /collections/{c}/memory, /telemetry) ├── processor.rs # Processor trait + Timing measurement type ├── results.rs # unified `{config, results}` document written by --json ├── stats.rs # benchmark run loops (parallel/RPS), stats output, throttling diff --git a/README.md b/README.md index 281af5f..a3332eb 100644 --- a/README.md +++ b/README.md @@ -152,6 +152,29 @@ and search configs. The CLI still controls *how* the benchmark runs (`-n`, `-p`, The flag-driven path (`bfb --scroll --keywords 100 …`) is still available when you prefer flat CLI flags over a YAML file. +### Memory & disk reporting + +Sampled automatically after indexing — in `bfb upload` and legacy runs — and +reported under `results.memory` in `--json`: + +``` +--- Memory & disk --- +total: 154.1 MiB disk, 6.2 MiB ram, 3.9 MiB cached (of 72.8 MiB expected) +vector (default): storage 66.0 MiB disk / 0.0 MiB ram, index 6.8 MiB disk / 0.0 MiB ram +payload: 67.1 MiB disk, 0.0 MiB ram +payload index color: 12.2 MiB disk, 2.7 MiB ram +process: 74.8 MiB resident, 49.0 MiB allocated (server-wide) +``` + +`cached` is what the page cache currently holds of the on-disk bytes, against +the `expected` a component wants cached to answer queries without touching disk +— the gap is how cold the cache is. The `process` line covers the whole server, +not just this collection. + +Read over Qdrant's **REST** API, on the port derived from `--uri` (6334 → 6333). +A failure here warns rather than failing a run that already has its headline +number; `--skip-server-stats` turns it off. + ### `schema` — print the upload-config file schema Print an annotated YAML reference enumerating every option accepted by an diff --git a/src/args/mod.rs b/src/args/mod.rs index eaac486..b08cc6f 100644 --- a/src/args/mod.rs +++ b/src/args/mod.rs @@ -43,6 +43,11 @@ pub struct Args { #[clap(long, default_value = "http://localhost:6334", global = true)] pub uri: Vec, + /// Skip the REST-only memory/disk report + #[clap(long, default_value_t = false, global = true)] + pub skip_server_stats: bool, + + /// Source of data to upload - fbin file. Random if not specified #[clap(long)] pub fbin: Option, diff --git a/src/main.rs b/src/main.rs index 9638084..5324c69 100644 --- a/src/main.rs +++ b/src/main.rs @@ -15,6 +15,7 @@ mod config; mod dataset; mod fbin_reader; mod generators; +mod memory; mod processor; mod query; mod results; @@ -33,6 +34,44 @@ async fn run_wait_index(args: &Args, stopped: Arc) -> Result String { + memory::rest_url_from_grpc(args.uri.first().expect("clap guarantees one --uri")) +} + +/// Sample memory and disk usage. Best-effort: a REST failure warns rather than +/// failing a run that already has its headline numbers. +async fn fetch_memory(args: &Args) -> Option { + if args.skip_server_stats { + return None; + } + let (rest_uri, collection) = (rest_uri(args), args.collection_name.clone()); + let api_key = std::env::var("QDRANT_API_KEY").ok(); + + // `ureq` is blocking; keep it off the async worker threads. + let fetched = tokio::task::spawn_blocking(move || { + memory::fetch(&rest_uri, &collection, api_key.as_deref()) + }) + .await; + + match fetched { + Ok(Ok(report)) => { + report.print(); + Some(report) + } + Ok(Err(err)) => { + eprintln!("Warning: could not read memory usage: {err:#}"); + eprintln!(" (pass --skip-server-stats to silence)"); + None + } + Err(err) => { + eprintln!("Warning: memory task failed: {err}"); + None + } + } +} + /// `bfb scroll --file config.yaml`: YAML-config-driven scroll. async fn run_scroll( args: Args, @@ -101,6 +140,7 @@ async fn run_upload( if !args.skip_wait_index && !args.skip_setup { results.results.index = Some(run_wait_index(&args, stopped.clone()).await?); + results.results.memory = fetch_memory(&args).await; } results.write_if_requested(&args) @@ -131,6 +171,7 @@ async fn run_benchmark(args: Args, stopped: Arc) -> Result<()> { if !args.skip_wait_index && !args.skip_setup { results.results.index = Some(run_wait_index(&args, stopped.clone()).await?); + results.results.memory = fetch_memory(&args).await; } if args.search || args.search_quality { diff --git a/src/memory.rs b/src/memory.rs new file mode 100644 index 0000000..5b5dae4 --- /dev/null +++ b/src/memory.rs @@ -0,0 +1,232 @@ +//! Memory and on-disk size, from `GET /collections/{c}/memory` and +//! `GET /telemetry` — the REST endpoints gRPC does not expose. + +use std::collections::BTreeMap; + +use anyhow::{Context, Result}; +use serde::{Deserialize, Serialize}; + +/// Bytes held by one component, on disk and in RAM. +/// +/// `cached_bytes` is what the page cache currently holds of the on-disk part; +/// `expected_cache_bytes` is what the component would like cached to serve +/// queries without hitting disk. A gap between them means a cold cache. +#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq)] +pub struct Usage { + pub disk_bytes: u64, + pub ram_bytes: u64, + pub cached_bytes: u64, + pub expected_cache_bytes: u64, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct NamedUsage { + pub name: String, + pub usage: Usage, +} + +/// A vector's storage and its index, sized separately. +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)] +pub struct VectorUsage { + pub name: String, + pub storage: Usage, + pub index: Usage, +} + +/// Process-wide allocator stats from `/telemetry`. Not per-collection: on a +/// server with other collections these cover all of them. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)] +pub struct ProcessMemory { + pub active_bytes: u64, + pub allocated_bytes: u64, + pub metadata_bytes: u64, + pub resident_bytes: u64, + pub retained_bytes: u64, +} + +/// Where the collection's bytes live, at the moment it was sampled. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct MemoryReport { + pub total: Usage, + pub vectors: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub sparse_vectors: Vec, + pub payload: Usage, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub payload_index: Vec, + /// Everything else the collection stores, e.g. `id_tracker`. + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + pub other: BTreeMap, + #[serde(skip_serializing_if = "Option::is_none")] + pub process: Option, +} + +impl MemoryReport { + pub fn print(&self) { + println!("--- Memory & disk ---"); + println!( + "total: {} disk, {} ram, {} cached (of {} expected)", + mib(self.total.disk_bytes), + mib(self.total.ram_bytes), + mib(self.total.cached_bytes), + mib(self.total.expected_cache_bytes), + ); + for vector in self.vectors.iter().chain(&self.sparse_vectors) { + let name = if vector.name.is_empty() { + "(default)" + } else { + &vector.name + }; + println!( + "vector {name}: storage {} disk / {} ram, index {} disk / {} ram", + mib(vector.storage.disk_bytes), + mib(vector.storage.ram_bytes), + mib(vector.index.disk_bytes), + mib(vector.index.ram_bytes), + ); + } + println!( + "payload: {} disk, {} ram", + mib(self.payload.disk_bytes), + mib(self.payload.ram_bytes) + ); + for field in &self.payload_index { + println!( + "payload index {}: {} disk, {} ram", + field.name, + mib(field.usage.disk_bytes), + mib(field.usage.ram_bytes) + ); + } + if let Some(process) = self.process { + println!( + "process: {} resident, {} allocated (server-wide)", + mib(process.resident_bytes), + mib(process.allocated_bytes) + ); + } + } +} + +fn mib(bytes: u64) -> String { + format!("{:.1} MiB", bytes as f64 / (1024.0 * 1024.0)) +} + +// ------------------------- wire types ------------------------------------- + +#[derive(Debug, Deserialize)] +struct RestEnvelope { + result: T, +} + +#[derive(Debug, Deserialize)] +struct CollectionMemory { + total: Usage, + #[serde(default)] + vectors: Vec, + #[serde(default)] + sparse_vectors: Vec, + #[serde(default)] + payload: Usage, + #[serde(default)] + payload_index: Vec, + #[serde(default)] + other: BTreeMap, +} + +#[derive(Debug, Deserialize)] +struct Telemetry { + memory: Option, +} + +// ------------------------- fetching --------------------------------------- + +/// Sample memory and disk usage for `collection`. +/// +/// The process-wide `/telemetry` half is best-effort: it is unavailable when the +/// server runs without jemalloc, and absent then rather than failing the sample. +pub fn fetch(rest_url: &str, collection: &str, api_key: Option<&str>) -> Result { + let base = rest_url.trim_end_matches('/'); + let memory: CollectionMemory = + get(&format!("{base}/collections/{collection}/memory"), api_key)?; + let process = get::(&format!("{base}/telemetry?details_level=1"), api_key) + .ok() + .and_then(|telemetry| telemetry.memory); + + Ok(MemoryReport { + total: memory.total, + vectors: memory.vectors, + sparse_vectors: memory.sparse_vectors, + payload: memory.payload, + payload_index: memory.payload_index, + other: memory.other, + process, + }) +} + +/// Derive the REST base URL from a gRPC `--uri`. +/// +/// Qdrant serves gRPC on 6334 and REST on 6333, so this is a one-port shift. +/// A `--uri` on any other port is returned unchanged. +pub fn rest_url_from_grpc(uri: &str) -> String { + let trimmed = uri.trim_end_matches('/'); + match trimmed.rsplit_once(":6334") { + Some((head, "")) => format!("{head}:6333"), + _ => trimmed.to_string(), + } +} + +fn get(url: &str, api_key: Option<&str>) -> Result { + let agent = ureq::Agent::new_with_defaults(); + let mut request = agent.get(url); + if let Some(key) = api_key { + request = request.header("api-key", key); + } + + let body = request + .call() + .with_context(|| format!("failed to GET {url}"))? + .into_body() + .read_to_string() + .with_context(|| format!("failed to read response from {url}"))?; + + let envelope: RestEnvelope = + serde_json::from_str(&body).with_context(|| format!("failed to parse {url}"))?; + Ok(envelope.result) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_real_server_response() { + let text = include_str!("../tests/fixtures/collection_memory.json"); + let envelope: RestEnvelope = serde_json::from_str(text).unwrap(); + let memory = envelope.result; + + assert_eq!(memory.total.disk_bytes, 161554912); + assert_eq!(memory.vectors.len(), 1); + assert_eq!(memory.vectors[0].name, ""); + assert_eq!(memory.vectors[0].index.disk_bytes, 7124194); + assert_eq!(memory.payload_index[0].name, "color"); + assert_eq!(memory.payload_index[0].usage.ram_bytes, 2817048); + assert!(memory.other.contains_key("id_tracker")); + } + + #[test] + fn derives_the_rest_url_from_the_grpc_uri() { + assert_eq!( + rest_url_from_grpc("http://localhost:6334"), + "http://localhost:6333" + ); + assert_eq!( + rest_url_from_grpc("https://xyz.cloud.qdrant.io:6334/"), + "https://xyz.cloud.qdrant.io:6333" + ); + // A non-default port is left alone rather than silently rewritten. + assert_eq!(rest_url_from_grpc("http://host:7000"), "http://host:7000"); + // A port that merely *contains* 6334 must not match. + assert_eq!(rest_url_from_grpc("http://host:63340"), "http://host:63340"); + } +} diff --git a/src/results.rs b/src/results.rs index b64fff7..a726179 100644 --- a/src/results.rs +++ b/src/results.rs @@ -11,6 +11,7 @@ use anyhow::{Context, Result}; use serde::{Deserialize, Serialize}; use crate::args::Args; +use crate::memory::MemoryReport; use crate::processor::Timing; /// Top-level document: `{ config: {...}, results: {...} }`. @@ -76,6 +77,9 @@ pub struct PhaseResults { pub search: Option, #[serde(skip_serializing_if = "Option::is_none")] pub scroll: Option, + /// Memory and disk sampled after the run, not a phase of it. + #[serde(skip_serializing_if = "Option::is_none")] + pub memory: Option, } /// M2: upload wall time and throughput. @@ -334,6 +338,7 @@ mod tests { index: Some(IndexPhase { wait_secs: 0.5 }), search: Some(phase(1.0)), scroll: None, + memory: None, }); doc.populate_legacy_fields(); diff --git a/tests/fixtures/collection_memory.json b/tests/fixtures/collection_memory.json new file mode 100644 index 0000000..dced30a --- /dev/null +++ b/tests/fixtures/collection_memory.json @@ -0,0 +1,55 @@ +{ + "result": { + "total": { + "disk_bytes": 161554912, + "ram_bytes": 6454560, + "cached_bytes": 4120576, + "expected_cache_bytes": 76330422 + }, + "vectors": [ + { + "name": "", + "storage": { + "disk_bytes": 69206228, + "ram_bytes": 0, + "cached_bytes": 4120576, + "expected_cache_bytes": 69206228 + }, + "index": { + "disk_bytes": 7124194, + "ram_bytes": 0, + "cached_bytes": 0, + "expected_cache_bytes": 7124194 + } + } + ], + "sparse_vectors": [], + "payload": { + "disk_bytes": 70320706, + "ram_bytes": 0, + "cached_bytes": 0, + "expected_cache_bytes": 0 + }, + "payload_index": [ + { + "name": "color", + "usage": { + "disk_bytes": 12791272, + "ram_bytes": 2817048, + "cached_bytes": 0, + "expected_cache_bytes": 0 + } + } + ], + "other": { + "id_tracker": { + "disk_bytes": 2112512, + "ram_bytes": 3637512, + "cached_bytes": 0, + "expected_cache_bytes": 0 + } + } + }, + "status": "ok", + "time": 0.018762371 +} From f398de7491f7ce48e0db3a680bd8701691fd6e70 Mon Sep 17 00:00:00 2001 From: tellet-q Date: Thu, 16 Jul 2026 08:44:05 +0200 Subject: [PATCH 2/2] Address review --- src/args/mod.rs | 1 - src/main.rs | 28 ++++++++++++++++++++++------ src/memory.rs | 49 +++++++++++++++++++++++++++++++++++++++---------- 3 files changed, 61 insertions(+), 17 deletions(-) diff --git a/src/args/mod.rs b/src/args/mod.rs index b08cc6f..d98e1b9 100644 --- a/src/args/mod.rs +++ b/src/args/mod.rs @@ -47,7 +47,6 @@ pub struct Args { #[clap(long, default_value_t = false, global = true)] pub skip_server_stats: bool, - /// Source of data to upload - fbin file. Random if not specified #[clap(long)] pub fbin: Option, diff --git a/src/main.rs b/src/main.rs index 5324c69..bc460e3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -34,10 +34,11 @@ async fn run_wait_index(args: &Args, stopped: Arc) -> Result String { - memory::rest_url_from_grpc(args.uri.first().expect("clap guarantees one --uri")) +/// The bound applied to the best-effort memory REST call: the CLI `--timeout` if +/// set, otherwise `memory`'s own fallback. +fn memory_http_timeout(args: &Args) -> Option { + args.timeout + .map(|secs| std::time::Duration::from_secs(secs as u64)) } /// Sample memory and disk usage. Best-effort: a REST failure warns rather than @@ -46,12 +47,27 @@ async fn fetch_memory(args: &Args) -> Option { if args.skip_server_stats { return None; } - let (rest_uri, collection) = (rest_uri(args), args.collection_name.clone()); + // Map every `--uri` to its REST port and try each in turn: the client spreads + // work across all URIs, so the first may be down while another is healthy. + let rest_uris: Vec = args + .uri + .iter() + .map(|u| memory::rest_url_from_grpc(u)) + .collect(); + let collection = args.collection_name.clone(); let api_key = std::env::var("QDRANT_API_KEY").ok(); + let timeout = memory_http_timeout(args); // `ureq` is blocking; keep it off the async worker threads. let fetched = tokio::task::spawn_blocking(move || { - memory::fetch(&rest_uri, &collection, api_key.as_deref()) + let mut last_err = None; + for rest_uri in &rest_uris { + match memory::fetch(rest_uri, &collection, api_key.as_deref(), timeout) { + Ok(report) => return Ok(report), + Err(err) => last_err = Some(err), + } + } + Err(last_err.expect("clap guarantees at least one --uri")) }) .await; diff --git a/src/memory.rs b/src/memory.rs index 5b5dae4..1f40d35 100644 --- a/src/memory.rs +++ b/src/memory.rs @@ -2,10 +2,14 @@ //! `GET /telemetry` — the REST endpoints gRPC does not expose. use std::collections::BTreeMap; +use std::time::Duration; use anyhow::{Context, Result}; use serde::{Deserialize, Serialize}; +/// Bound the REST calls when `--timeout` is unset; `ureq` applies none by default. +const DEFAULT_HTTP_TIMEOUT: Duration = Duration::from_secs(30); + /// Bytes held by one component, on disk and in RAM. /// /// `cached_bytes` is what the page cache currently holds of the on-disk part; @@ -71,14 +75,19 @@ impl MemoryReport { mib(self.total.cached_bytes), mib(self.total.expected_cache_bytes), ); - for vector in self.vectors.iter().chain(&self.sparse_vectors) { + for (kind, vector) in self + .vectors + .iter() + .map(|v| ("vector", v)) + .chain(self.sparse_vectors.iter().map(|v| ("sparse vector", v))) + { let name = if vector.name.is_empty() { "(default)" } else { &vector.name }; println!( - "vector {name}: storage {} disk / {} ram, index {} disk / {} ram", + "{kind} {name}: storage {} disk / {} ram, index {} disk / {} ram", mib(vector.storage.disk_bytes), mib(vector.storage.ram_bytes), mib(vector.index.disk_bytes), @@ -145,13 +154,25 @@ struct Telemetry { /// /// The process-wide `/telemetry` half is best-effort: it is unavailable when the /// server runs without jemalloc, and absent then rather than failing the sample. -pub fn fetch(rest_url: &str, collection: &str, api_key: Option<&str>) -> Result { +pub fn fetch( + rest_url: &str, + collection: &str, + api_key: Option<&str>, + timeout: Option, +) -> Result { let base = rest_url.trim_end_matches('/'); - let memory: CollectionMemory = - get(&format!("{base}/collections/{collection}/memory"), api_key)?; - let process = get::(&format!("{base}/telemetry?details_level=1"), api_key) - .ok() - .and_then(|telemetry| telemetry.memory); + let memory: CollectionMemory = get( + &format!("{base}/collections/{collection}/memory"), + api_key, + timeout, + )?; + let process = get::( + &format!("{base}/telemetry?details_level=1"), + api_key, + timeout, + ) + .ok() + .and_then(|telemetry| telemetry.memory); Ok(MemoryReport { total: memory.total, @@ -176,8 +197,16 @@ pub fn rest_url_from_grpc(uri: &str) -> String { } } -fn get(url: &str, api_key: Option<&str>) -> Result { - let agent = ureq::Agent::new_with_defaults(); +fn get( + url: &str, + api_key: Option<&str>, + timeout: Option, +) -> Result { + let agent = ureq::Agent::new_with_config( + ureq::config::Config::builder() + .timeout_global(Some(timeout.unwrap_or(DEFAULT_HTTP_TIMEOUT))) + .build(), + ); let mut request = agent.get(url); if let Some(key) = api_key { request = request.header("api-key", key);