Attempt to fix flaky test_hw_metrics_cancellation test (#9233)

* Fix flaky test_hw_metrics_cancellation test

* Speedup test duration

.config/nextest.toml

@@ -10,3 +10,13 @@ failure-output = "immediate-final"
 # save junit report
 [profile.ci.junit]
 path = "junit.xml"
+
+# The `test_hw_metrics_cancellation` test relies on a search timeout, which might
+# not be scheduled at all if the CI runner's CPU pressure is too high.
+# This is because `spawn_blocking` sometimes doesn't spawn a new worker thread before
+# the test's search timeout is reached.
+# To trick nextest into not running other tests in parallel, we override this test
+# specifically and force it to get all CPUs allocated.
+[[profile.default.overrides]]
+filter = 'package(collection) and test(test_hw_metrics_cancellation)'
+threads-required = "num-test-threads"

lib/collection/src/tests/hw_metrics.rs

@@ -4,8 +4,8 @@ use std::time::Duration;
 use common::budget::ResourceBudget;
 use common::counter::hardware_accumulator::{HwMeasurementAcc, HwSharedDrain};
 use common::save_on_disk::SaveOnDisk;
-use rand::rngs::ThreadRng;
+use rand::rngs::SmallRng;
-use rand::{Rng, rng};
+use rand::{Rng, SeedableRng, rng};
 use segment::data_types::vectors::{NamedQuery, VectorInternal, VectorStructInternal};
 use shard::query::query_enum::QueryEnum;
 use shard::search::CoreSearchRequestBatch;
@@ -27,7 +27,9 @@ use crate::tests::fixtures::create_collection_config_with_dim;
 async fn test_hw_metrics_cancellation() {
     let collection_dir = Builder::new().prefix("test_collection").tempdir().unwrap();
 
-    let mut config = create_collection_config_with_dim(512);
+    const DIM: usize = 2048;
+
+    let mut config = create_collection_config_with_dim(DIM);
     config.optimizer_config.indexing_threshold = None;
 
     let collection_name = "test".to_string();
@@ -55,7 +57,7 @@ async fn test_hw_metrics_cancellation() {
     .await
     .unwrap();
 
-    let upsert_ops = make_random_points_upsert_op(10_000);
+    let upsert_ops = make_random_points_upsert_op(50_000, DIM);
     shard
         .update(
             upsert_ops.into(),
@@ -71,7 +73,7 @@ async fn test_hw_metrics_cancellation() {
         searches: vec![CoreSearchRequest {
             query: QueryEnum::Nearest(NamedQuery {
                 using: None,
-                query: VectorInternal::from(rand_vector(512, &mut rand)),
+                query: VectorInternal::from(rand_vector(DIM, &mut rand)),
             }),
             filter: None,
             params: None,
@@ -91,7 +93,9 @@ async fn test_hw_metrics_cancellation() {
             .do_search(
                 Arc::new(req),
                 &current_runtime,
-                Duration::from_millis(10), // Very short duration to hit timeout before the search finishes
+                // Short but not extremely short timeout to not conflict with os-thread spawning overhead.
+                // (For more information see https://github.com/qdrant/qdrant/pull/9233)
+                Duration::from_millis(350),
                 hw_counter,
             )
             .await;
@@ -119,13 +123,15 @@ async fn test_hw_metrics_cancellation() {
     assert!(outer_hw.get_cpu() > 0);
 }
 
-fn make_random_points_upsert_op(len: usize) -> CollectionUpdateOperations {
+fn make_random_points_upsert_op(len: usize, dim: usize) -> CollectionUpdateOperations {
     let mut points = vec![];
 
-    let mut rand = rng();
+    // ThreadRng is too slow for creating 40k vectors @ 2048 dimensions each.
+    // SmallRng cuts total test duration in half (20s->10s).
+    let mut rand = SmallRng::seed_from_u64(0xC0FFEE);
 
     for i in 0..len as u64 {
-        let rand_vector = rand_vector(512, &mut rand);
+        let rand_vector = rand_vector(dim, &mut rand);
         points.push(PointStructPersisted {
             id: segment::types::ExtendedPointId::NumId(i),
             vector: VectorStructInternal::from(rand_vector).into(),
@@ -138,6 +144,6 @@ fn make_random_points_upsert_op(len: usize) -> CollectionUpdateOperations {
     CollectionUpdateOperations::PointOperation(PointOperations::UpsertPoints(op))
 }
 
-fn rand_vector(size: usize, rand: &mut ThreadRng) -> Vec<f32> {
+fn rand_vector(size: usize, rand: &mut impl Rng) -> Vec<f32> {
     (0..size).map(|_| rand.next_u32() as f32).collect()
 }