Add memory reporting

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

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

src/args/mod.rs

@@ -43,6 +43,10 @@ pub struct Args {
     #[clap(long, default_value = "http://localhost:6334", global = true)]
     pub uri: Vec<String>,
 
+    /// 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<String>,

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,60 @@ async fn run_wait_index(args: &Args, stopped: Arc<AtomicBool>) -> Result<IndexPh
     Ok(IndexPhase { wait_secs })
 }
 
+/// 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<std::time::Duration> {
+    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
+/// failing a run that already has its headline numbers.
+async fn fetch_memory(args: &Args) -> Option<memory::MemoryReport> {
+    if args.skip_server_stats {
+        return None;
+    }
+    // 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<String> = 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 || {
+        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;
+
+    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 +156,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 +187,7 @@ async fn run_benchmark(args: Args, stopped: Arc<AtomicBool>) -> 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 {

src/memory.rs

@@ -0,0 +1,261 @@
+//! 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 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;
+/// `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<VectorUsage>,
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub sparse_vectors: Vec<VectorUsage>,
+    pub payload: Usage,
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub payload_index: Vec<NamedUsage>,
+    /// Everything else the collection stores, e.g. `id_tracker`.
+    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
+    pub other: BTreeMap<String, Usage>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub process: Option<ProcessMemory>,
+}
+
+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 (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!(
+                "{kind} {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<T> {
+    result: T,
+}
+
+#[derive(Debug, Deserialize)]
+struct CollectionMemory {
+    total: Usage,
+    #[serde(default)]
+    vectors: Vec<VectorUsage>,
+    #[serde(default)]
+    sparse_vectors: Vec<VectorUsage>,
+    #[serde(default)]
+    payload: Usage,
+    #[serde(default)]
+    payload_index: Vec<NamedUsage>,
+    #[serde(default)]
+    other: BTreeMap<String, Usage>,
+}
+
+#[derive(Debug, Deserialize)]
+struct Telemetry {
+    memory: Option<ProcessMemory>,
+}
+
+// ------------------------- 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>,
+    timeout: Option<Duration>,
+) -> Result<MemoryReport> {
+    let base = rest_url.trim_end_matches('/');
+    let memory: CollectionMemory = get(
+        &format!("{base}/collections/{collection}/memory"),
+        api_key,
+        timeout,
+    )?;
+    let process = get::<Telemetry>(
+        &format!("{base}/telemetry?details_level=1"),
+        api_key,
+        timeout,
+    )
+    .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<T: serde::de::DeserializeOwned>(
+    url: &str,
+    api_key: Option<&str>,
+    timeout: Option<Duration>,
+) -> Result<T> {
+    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);
+    }
+
+    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<T> =
+        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<CollectionMemory> = 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");
+    }
+}