Add test for OOD during indexing (#4267)

* Add test for OOD during indexing

* Start indexing right after 1st OOD message

* Only send search request after insert loop

* Fail early when encountering out-of-storage during optimization (#4578)

* fs4@0.8.4

* fail early on low storage

* move `dir_size` to `common`

* move the ood bailout to `SegmentOptimizer::optimized_segment_builder`

* drop dead code

* move dir_size to common subcrate

---------

Co-authored-by: generall <andrey@vasnetsov.com>

---------

Co-authored-by: xhjkl <xhjkl@users.noreply.github.com>
Co-authored-by: generall <andrey@vasnetsov.com>

.github/workflows/integration-tests.yml

@@ -195,10 +195,14 @@ jobs:
         working-directory: ./tests/low-ram
         shell: bash
         run: ./low-ram.sh
-      - name: Run low Disk test
+      - name: Run low Disk test - search
         working-directory: ./tests/low-disk
         shell: bash
         run: ./low-disk.sh
+      - name: Run low Disk test - indexing
+        working-directory: ./tests/low-disk
+        shell: bash
+        run: ./low-disk.sh indexing
 
   test-snapshot-operations-s3-minio:
     runs-on: ubuntu-latest

Cargo.lock

@@ -5310,6 +5310,7 @@ dependencies = [
  "criterion",
  "dataset",
  "fnv",
+ "fs4",
  "fs_extra",
  "generic-tests",
  "geo",

lib/collection/src/collection_manager/optimizers/segment_optimizer.rs

@@ -4,6 +4,7 @@ use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
 use common::cpu::CpuPermit;
+use common::disk::dir_size;
 use io::storage_version::StorageVersion;
 use itertools::Itertools;
 use parking_lot::{Mutex, RwLock, RwLockUpgradableReadGuard};
@@ -118,6 +119,10 @@ pub trait SegmentOptimizer {
         // }
         let mut bytes_count_by_vector_name = HashMap::new();
 
+        // Counting up how much space do the segments being optimized actually take on the fs.
+        // If there was at least one error while reading the size, this will be `None`.
+        let mut space_occupied = Some(0u64);
+
         for segment in optimizing_segments {
             let segment = match segment {
                 LockedSegment::Original(segment) => segment,
@@ -134,6 +139,28 @@ pub trait SegmentOptimizer {
                 let size = bytes_count_by_vector_name.entry(vector_name).or_insert(0);
                 *size += vector_size;
             }
+
+            space_occupied = space_occupied
+                .and_then(|x| dir_size(locked_segment.data_path()).ok().map(|y| x + y));
+        }
+
+        let space_needed = space_occupied.map(|x| 2 * x);
+        let space_available = fs4::available_space(self.temp_path()).ok();
+
+        match (space_available, space_needed) {
+            (Some(space_available), Some(space_needed)) => {
+                if space_available < space_needed {
+                    return Err(CollectionError::service_error(
+                        "Not enough space available for optimization".to_string(),
+                    ));
+                }
+            }
+            _ => {
+                log::warn!(
+                    "Could not estimate available storage space in `{}`; will try optimizing anyway",
+                    self.name()
+                );
+            }
         }
 
         // Example: maximal_vector_store_size_bytes = 10200 * dim * VECTOR_ELEMENT_SIZE
@@ -511,9 +538,8 @@ pub trait SegmentOptimizer {
             proxy_ids
         };
 
-        check_process_stopped(stopped).map_err(|error| {
+        check_process_stopped(stopped).inspect_err(|_| {
             self.handle_cancellation(&segments, &proxy_ids, &tmp_segment);
-            error
         })?;
 
         // ---- SLOW PART -----

lib/common/common/src/disk.rs

@@ -0,0 +1,17 @@
+use std::path::PathBuf;
+
+/// How many bytes a directory takes.
+pub fn dir_size(path: impl Into<PathBuf>) -> std::io::Result<u64> {
+    fn dir_size(mut dir: std::fs::ReadDir) -> std::io::Result<u64> {
+        dir.try_fold(0, |acc, file| {
+            let file = file?;
+            let size = match file.metadata()? {
+                data if data.is_dir() => dir_size(std::fs::read_dir(file.path())?)?,
+                data => data.len(),
+            };
+            Ok(acc + size)
+        })
+    }
+
+    dir_size(std::fs::read_dir(path.into())?)
+}

lib/common/common/src/lib.rs

@@ -1,5 +1,6 @@
 pub mod cpu;
 pub mod defaults;
+pub mod disk;
 pub mod fixed_length_priority_queue;
 pub mod math;
 pub mod panic;

lib/segment/Cargo.toml

@@ -87,6 +87,7 @@ memory = { path = "../common/memory" }
 sparse = { path = "../sparse" }
 
 tracing = { workspace = true, optional = true }
+fs4 = "0.8.4"
 macro_rules_attribute = "0.2.0"
 generic-tests = { workspace = true }
 nom = "7.1.3"

lib/segment/src/common/mod.rs

@@ -20,7 +20,6 @@ use crate::data_types::vectors::{QueryVector, VectorRef};
 use crate::types::{SegmentConfig, SparseVectorDataConfig, VectorDataConfig};
 
 pub type Flusher = Box<dyn FnOnce() -> OperationResult<()> + Send>;
-
 /// Check that the given vector name is part of the segment config.
 ///
 /// Returns an error if incompatible.

tests/low-disk/create_and_search_items.py

@@ -1,10 +1,14 @@
 import argparse
 import random
 import uuid
+from time import sleep
 
 import requests
 
 
+VECTOR_SIZE = 4
+
+
 def generate_points(amount: int):
     for _ in range(amount):
         result_item = {"points": []}
@@ -12,33 +16,55 @@ def generate_points(amount: int):
             result_item["points"].append(
                 {
                     "id": str(uuid.uuid4()),
-                    "vector": [round(random.uniform(0, 1), 2) for _ in range(4)],
+                    "vector": [round(random.uniform(0, 1), 2) for _ in range(VECTOR_SIZE)],
                     "payload": {"city": ["Berlin", "London"]}
                 }
             )
         yield result_item
 
 
-def create_collection(qdrant_host, collection_name):
+def create_collection(qdrant_host, collection_name, collection_json):
     resp = requests.put(
-        f"{qdrant_host}/collections/{collection_name}?timeout=60", json={
-            "vectors": {
-                "size": 4,
-                "distance": "Cosine"
-            }
-        })
+        f"{qdrant_host}/collections/{collection_name}?timeout=60", json=collection_json)
     if resp.status_code != 200:
         print(f"Collection creation failed with response body:\n{resp.json()}")
         exit(-1)
 
 
-def insert_points(qdrant_host, collection_name, batch_json):
+def update_collection(qdrant_host, collection_name, collection_json: dict):
+    resp = requests.patch(
+        f"{qdrant_host}/collections/{collection_name}?timeout=60", json=collection_json)
+    if resp.status_code != 200:
+        print(f"Collection patching failed with response body:\n{resp.json()}")
+        exit(-1)
+
+
+def wait_for_status(qdrant_host, collection_name, status: str):
+    curr_status = ""
+    for i in range(30):
+        resp = requests.get(
+            f"{qdrant_host}/collections/{collection_name}")
+        curr_status = resp.json()["result"]["status"]
+        if resp.status_code != 200:
+            print(f"Collection info fetching failed with response body:\n{resp.json()}")
+            exit(-1)
+        if curr_status == status:
+            print(f"Status {status}: OK")
+            return "ok"
+        sleep(1)
+        print(f"Wait for status {status}")
+    print(f"After 30s status is not {status}, found: {curr_status}. Stop waiting.")
+
+
+def insert_points(qdrant_host, collection_name, batch_json, quit_on_ood: bool = False):
     resp = requests.put(
         f"{qdrant_host}/collections/{collection_name}/points?wait=true", json=batch_json
     )
-    EXPECTED_ERROR_MESSAGE = "No space left on device"
+    expected_error_message = "No space left on device"
     if resp.status_code != 200:
-        if resp.status_code == 500 and EXPECTED_ERROR_MESSAGE in resp.text:
+        if resp.status_code == 500 and expected_error_message in resp.text:
+            if quit_on_ood:
+                return "ood"
             requests.put(f"{qdrant_host}/collections/{collection_name}/points?wait=true", json=batch_json)
         else:
             error_response = resp.json()
@@ -48,7 +74,7 @@ def insert_points(qdrant_host, collection_name, batch_json):
 
 def search_point(qdrant_host, collection_name):
     query = {
-        "vector": [round(random.uniform(0, 1), 2) for _ in range(4)],
+        "vector": [round(random.uniform(0, 1), 2) for _ in range(VECTOR_SIZE)],
         "top": 10,
         "filters": {
             "city": {
@@ -67,22 +93,78 @@ def search_point(qdrant_host, collection_name):
     return resp.json()
 
 
-def initialize_qdrant(qdrant_host, collection_name, points_amount):
-    create_collection(qdrant_host, collection_name)
+def insert_points_and_search(qdrant_host, collection_name, points_amount):
+    collection_json = {
+        "vectors": {
+            "size": VECTOR_SIZE,
+            "distance": "Cosine"
+        }
+    }
+    create_collection(qdrant_host, collection_name, collection_json)
     for points_batch in generate_points(points_amount):
         insert_points(qdrant_host, collection_name, points_batch)
-        search_point(qdrant_host, collection_name)
+
+    search_point(qdrant_host, collection_name)
+
+
+def insert_points_then_index(qdrant_host, collection_name, points_amount):
+    """
+    Disable indexing, create collection, insert points up to
+    a point when there is no space left on disk, then start indexing.
+    Wait for indexing then send a search request.
+    @param qdrant_host:
+    @param collection_name:
+    @param points_amount:
+    @return:
+    """
+    collection_json = {
+        "vectors": {
+            "size": VECTOR_SIZE,
+            "distance": "Cosine"
+        },
+        "optimizers_config": {
+            "indexing_threshold": 0,
+            "default_segment_number": 2
+        },
+        "wal_config": {
+            "wal_capacity_mb": 1
+        }
+    }
+    create_collection(qdrant_host, collection_name, collection_json)
+    for points_batch in generate_points(points_amount):
+        res = insert_points(qdrant_host, collection_name, points_batch, quit_on_ood=True)
+        if res == "ood":
+            break
+
+    # start indexing
+    collection_json = {
+        "optimizers_config": {
+            "indexing_threshold": 10
+        }
+    }
+    update_collection(qdrant_host, collection_name, collection_json)
+    wait_for_status(qdrant_host, collection_name, "yellow")
+    wait_for_status(qdrant_host, collection_name, "green")
+    search_point(qdrant_host, collection_name)
 
 
 def main():
     parser = argparse.ArgumentParser("Create all required test items")
+    parser.add_argument("test_item")
     parser.add_argument("collection_name")
     parser.add_argument("points_amount", type=int)
     parser.add_argument("ports", type=int, nargs="+")
     args = parser.parse_args()
 
     qdrant_host = f"http://127.0.0.1:{args.ports[0]}"
-    initialize_qdrant(qdrant_host, args.collection_name, args.points_amount)
+    if args.test_item == "search":
+        print("run search test")
+        insert_points_and_search(qdrant_host, args.collection_name, args.points_amount)
+    elif args.test_item == "indexing":
+        print("run indexing test")
+        insert_points_then_index(qdrant_host, args.collection_name, args.points_amount)
+    else:
+        raise ValueError("Invalid test_item value, allowed: search|indexing")
     print("SUCCESS")
 
 

tests/low-disk/low-disk.sh

@@ -3,18 +3,20 @@
 # and verify that it doesn't crash when disk
 # is running low during points insertion.
 
+# possible values are search|indexing
+declare TEST=${1:-"search"}
+
 set -xeuo pipefail
 
 cd "$(dirname "${BASH_SOURCE[0]}")"
 
-#declare DOCKER_IMAGE_NAME=qdrant-recovery
-declare DOCKER_IMAGE_NAME=qdrant/qdrant
+declare DOCKER_IMAGE_NAME=qdrant-recovery
 
 docker buildx build --build-arg=PROFILE=ci --load ../../ --tag=$DOCKER_IMAGE_NAME
 
-declare OOD_CONTAINER_NAME=qdrant-ood
+declare OOD_CONTAINER_NAME=qdrant-ood-$TEST
 
-docker rm -f ${OOD_CONTAINER_NAME} || true
+docker rm -f "${OOD_CONTAINER_NAME}" || true
 
 declare container && container=$(
     docker run -d \
@@ -45,7 +47,7 @@ done
 
 # check that low disk is handled OK during points insertion
 # this also does search after each insertion
-python3 create_and_search_items.py low-disk 2000 6333
+python3 create_and_search_items.py "$TEST" low-disk 2000 6333
 
 sleep 5
 
@@ -57,6 +59,6 @@ if (! docker logs "$container" 2>&1 | grep "$OUT_OF_DISK_MSG") ; then
     exit 9
 fi
 
-printf 'Insertion: OK\n\n'
+printf '%s: OK\n\n' "${TEST}"
 
 echo "Success"