Interleave serverless requests across collections

src/processor.rs

@@ -33,4 +33,15 @@ pub trait Processor {
     }
 
     fn get_batch_size(&self) -> usize;
+
+    /// Number of requests needed for `total_items`. Processors with variable
+    /// request sizes can override this to expose a precomputed schedule.
+    fn request_count(&self, total_items: usize) -> usize {
+        total_items.div_ceil(self.get_batch_size())
+    }
+
+    /// Number of items represented by one request (for progress accounting).
+    fn request_size(&self, _req_id: usize) -> usize {
+        self.get_batch_size()
+    }
 }

src/serverless/distribution.rs

@@ -20,6 +20,56 @@ pub struct CollectionPicker {
     zipf: Option<Zipf<f64>>,
 }
 
+/// One request produced while draining pre-allocated collection budgets.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct BudgetBatch {
+    pub collection: usize,
+    /// Global item offset. Each collection owns one contiguous range, even
+    /// though requests from those ranges are emitted in random order.
+    pub offset: usize,
+    pub count: usize,
+}
+
+/// Randomly drain collection budgets in batches of at most `batch_size`.
+///
+/// Sampling is uniform among collections that still have budget. The desired
+/// uniform/Zipf skew is already represented by `budgets`; applying it again
+/// while draining would skew the workload twice.
+pub fn drain_budgets(budgets: &[usize], batch_size: usize, rng: &mut impl Rng) -> Vec<BudgetBatch> {
+    assert!(batch_size > 0, "batch size must be positive");
+
+    let mut remaining = budgets.to_vec();
+    let mut next_offset = Vec::with_capacity(budgets.len());
+    let mut offset = 0usize;
+    for &budget in budgets {
+        next_offset.push(offset);
+        offset += budget;
+    }
+    let mut active: Vec<_> = remaining
+        .iter()
+        .enumerate()
+        .filter_map(|(i, &budget)| (budget > 0).then_some(i))
+        .collect();
+    let mut batches = Vec::new();
+
+    while !active.is_empty() {
+        let active_idx = rng.random_range(0..active.len());
+        let collection = active[active_idx];
+        let count = remaining[collection].min(batch_size);
+        batches.push(BudgetBatch {
+            collection,
+            offset: next_offset[collection],
+            count,
+        });
+        next_offset[collection] += count;
+        remaining[collection] -= count;
+        if remaining[collection] == 0 {
+            active.swap_remove(active_idx);
+        }
+    }
+    batches
+}
+
 impl CollectionPicker {
     pub fn new(n: usize, distribution: Distribution) -> anyhow::Result<Self> {
         anyhow::ensure!(n > 0, "collections-count must be > 0");
@@ -92,4 +142,21 @@ mod tests {
         // Rank 0 should get strictly more than the last rank on average.
         assert!(counts[0] > counts[9]);
     }
+
+    #[test]
+    fn draining_budgets_preserves_limits_offsets_and_totals() {
+        let mut rng = StdRng::seed_from_u64(7);
+        let batches = drain_budgets(&[5, 0, 7], 3, &mut rng);
+        assert!(batches.iter().all(|batch| batch.count <= 3));
+        let mut totals = [0usize; 3];
+        let mut offsets = [Vec::new(), Vec::new(), Vec::new()];
+        for batch in batches {
+            totals[batch.collection] += batch.count;
+            offsets[batch.collection].push(batch.offset);
+        }
+        assert_eq!(totals, [5, 0, 7]);
+        offsets.iter_mut().for_each(|values| values.sort_unstable());
+        assert_eq!(offsets[0], [0, 3]);
+        assert_eq!(offsets[2], [5, 8, 11]);
+    }
 }

src/serverless/query.rs

@@ -23,7 +23,7 @@ use qdrant_client::serverless::{CollectionConfig, QdrantServerless};
 use super::args::ServerlessQueryArgs;
 use super::client::create_clients;
 use super::collections::list_matching;
-use super::distribution::CollectionPicker;
+use super::distribution::{BudgetBatch, CollectionPicker, drain_budgets};
 use crate::args::Args;
 use crate::client::retry_with_clients;
 use crate::config::examples::{ExampleKind, resolve};
@@ -48,7 +48,7 @@ struct ServerlessQueryProcessor {
     stopped: Arc<AtomicBool>,
     clients: Vec<QdrantServerless>,
     names: Vec<String>,
-    picker: CollectionPicker,
+    batches: Vec<BudgetBatch>,
     generator: ConfigSearchGenerator,
     start_timestamp_millis: f64,
     start_time: Instant,
@@ -97,11 +97,12 @@ impl ServerlessQueryProcessor {
             return Ok(());
         }
 
+        let batch = self.batches[req_id];
         let mut rng = rand::rng();
-        let collection = &self.names[self.picker.pick(&mut rng)];
+        let collection = &self.names[batch.collection];
         let template_idx = self.generator.random_template_idx(&mut rng);
 
-        let queries = (0..self.args.search_batch_size)
+        let queries = (0..batch.count)
             .map(|_| {
                 let generated = self
                     .generator
@@ -183,6 +184,14 @@ impl Processor for ServerlessQueryProcessor {
     fn get_batch_size(&self) -> usize {
         self.args.search_batch_size
     }
+
+    fn request_count(&self, _total_items: usize) -> usize {
+        self.batches.len()
+    }
+
+    fn request_size(&self, req_id: usize) -> usize {
+        self.batches[req_id].count
+    }
 }
 
 /// Derive a search config from a live collection: one random-query template
@@ -292,11 +301,15 @@ pub async fn run(args: &Args, query: ServerlessQueryArgs, stopped: Arc<AtomicBoo
     args.collection_name = format!("{}*", query.collection_prefix);
     let mut results = BenchmarkResults::new(&args, origin);
 
+    let picker = CollectionPicker::new(names.len(), query.distribution.into())?;
+    let query_budgets = picker.allocate(args.num_vectors_or_default(), &mut rand::rng());
+    let batches = drain_budgets(&query_budgets, args.search_batch_size, &mut rand::rng());
+
     let processor = ServerlessQueryProcessor {
         args: args.clone(),
         stopped: stopped.clone(),
         clients,
-        picker: CollectionPicker::new(names.len(), query.distribution.into())?,
+        batches,
         names,
         generator: ConfigSearchGenerator::new(&search_config)?,
         start_timestamp_millis: std::time::SystemTime::now()

src/serverless/upload.rs

@@ -1,6 +1,5 @@
 //! `bfb serverless upload` — spread points across lazily-created collections.
 
-use std::cmp::min;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
@@ -17,7 +16,7 @@ use super::args::ServerlessUploadArgs;
 use super::client::create_clients;
 use super::collections::CollectionRegistry;
 use super::convert::to_serverless_config;
-use super::distribution::CollectionPicker;
+use super::distribution::{CollectionPicker, drain_budgets};
 use crate::args::Args;
 use crate::client::retry_with_clients;
 use crate::config;
@@ -38,28 +37,6 @@ struct Batch {
     count: usize,
 }
 
-/// Split the per-collection allocation into upsert batches. Point ids are
-/// laid out contiguously across collections starting at `offset`, so with a
-/// dataset source every collection gets a different slice of the data.
-fn plan_batches(per_collection: &[usize], batch_size: usize, offset: usize) -> Vec<Batch> {
-    let mut batches = Vec::new();
-    let mut next_id = offset as u64;
-    for (coll_idx, &count) in per_collection.iter().enumerate() {
-        let mut remaining = count;
-        while remaining > 0 {
-            let n = min(batch_size, remaining);
-            batches.push(Batch {
-                coll_idx,
-                start_id: next_id,
-                count: n,
-            });
-            next_id += n as u64;
-            remaining -= n;
-        }
-    }
-    batches
-}
-
 /// Request pacing: `--rps` fires at a fixed interval regardless of how many
 /// requests are in flight (missed ticks are skipped, not burst); otherwise
 /// the regular `--throttle` stream.
@@ -201,7 +178,17 @@ pub async fn run(
         println!("  …");
     }
 
-    let batches = plan_batches(&per_collection, args.batch_size, args.offset);
+    // Budgets determine the distribution; request order is randomized among
+    // collections that still have budget so bounded parallelism interleaves
+    // traffic instead of draining collections sequentially.
+    let batches: Vec<_> = drain_budgets(&per_collection, args.batch_size, &mut rand::rng())
+        .into_iter()
+        .map(|batch| Batch {
+            coll_idx: batch.collection,
+            start_id: (args.offset + batch.offset) as u64,
+            count: batch.count,
+        })
+        .collect();
     let generator: Box<dyn PointGenerator> = Box::new(ConfigGenerator::new(&upload_config)?);
 
     let logger = env_logger::Builder::from_default_env().build();
@@ -302,13 +289,17 @@ pub async fn run(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use rand::SeedableRng;
 
     #[test]
     fn batches_cover_every_point_with_unique_ids() {
-        let batches = plan_batches(&[5, 0, 7], 3, 100);
-        let counts: Vec<_> = batches.iter().map(|b| (b.coll_idx, b.count)).collect();
-        assert_eq!(counts, vec![(0, 3), (0, 2), (2, 3), (2, 3), (2, 1)]);
-        let ids: Vec<_> = batches.iter().map(|b| b.start_id).collect();
-        assert_eq!(ids, vec![100, 103, 105, 108, 111]);
+        let mut rng = rand::rngs::StdRng::seed_from_u64(7);
+        let batches = drain_budgets(&[5, 0, 7], 3, &mut rng);
+        let mut ids: Vec<_> = batches
+            .iter()
+            .flat_map(|b| b.offset..b.offset + b.count)
+            .collect();
+        ids.sort_unstable();
+        assert_eq!(ids, (0..12).collect::<Vec<_>>());
     }
 }

src/stats.rs

@@ -77,8 +77,7 @@ pub async fn process<P: Processor + Sync>(
     processor: P,
 ) -> Result<QueryPhase> {
     let started = Instant::now();
-    let batch_size = processor.get_batch_size();
-    let batch_count = args.num_vectors_or_default().div_ceil(batch_size);
+    let batch_count = processor.request_count(args.num_vectors_or_default());
 
     let multiprogress = MultiProgress::new();
     let progress_bar = multiprogress.add(ProgressBar::new(args.num_vectors_or_default() as u64));
@@ -98,7 +97,6 @@ pub async fn process<P: Processor + Sync>(
             &processor,
             &progress_bar,
             batch_count,
-            batch_size,
             target_rps,
         )
         .await?;
@@ -109,7 +107,6 @@ pub async fn process<P: Processor + Sync>(
             &processor,
             &progress_bar,
             batch_count,
-            batch_size,
         )
         .await?;
     }
@@ -183,13 +180,12 @@ async fn process_with_parallel<P: Processor>(
     processor: &P,
     progress_bar: &ProgressBar,
     batch_count: usize,
-    batch_size: usize,
 ) -> Result<()> {
     let query_stream = (0..batch_count)
         .take_while(|_| !stopped.load(Ordering::Relaxed))
         .map(|n| {
             let future = processor.make_request(n, args, progress_bar);
-            progress_bar.inc(batch_size as u64);
+            progress_bar.inc(processor.request_size(n) as u64);
             future
         });
 
@@ -219,7 +215,6 @@ async fn process_with_rps<P: Processor + Sync>(
     processor: &P,
     progress_bar: &ProgressBar,
     batch_count: usize,
-    batch_size: usize,
     target_rps: f64,
 ) -> Result<()> {
     use futures::stream::FuturesUnordered;
@@ -255,7 +250,7 @@ async fn process_with_rps<P: Processor + Sync>(
 
                 let req_id = requests_sent;
                 requests_sent += 1;
-                progress_bar.inc(batch_size as u64);
+                progress_bar.inc(processor.request_size(req_id) as u64);
 
                 let future = processor.make_request(req_id, args, progress_bar);
                 in_flight.push(future);