diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml index 2d6f470f16..fc9cae8be4 100644 --- a/.github/workflows/integration-tests.yml +++ b/.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 diff --git a/Cargo.lock b/Cargo.lock index 67455f2885..40fc7ee84d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5310,6 +5310,7 @@ dependencies = [ "criterion", "dataset", "fnv", + "fs4", "fs_extra", "generic-tests", "geo", diff --git a/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs b/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs index ce6f1daea6..bb94844271 100644 --- a/lib/collection/src/collection_manager/optimizers/segment_optimizer.rs +++ b/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 ----- diff --git a/lib/common/common/src/disk.rs b/lib/common/common/src/disk.rs new file mode 100644 index 0000000000..4a4eb57497 --- /dev/null +++ b/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) -> std::io::Result { + fn dir_size(mut dir: std::fs::ReadDir) -> std::io::Result { + 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())?) +} diff --git a/lib/common/common/src/lib.rs b/lib/common/common/src/lib.rs index 50cd9f5adf..16f510f5bd 100644 --- a/lib/common/common/src/lib.rs +++ b/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; diff --git a/lib/segment/Cargo.toml b/lib/segment/Cargo.toml index b51c346756..d448b187d2 100644 --- a/lib/segment/Cargo.toml +++ b/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" diff --git a/lib/segment/src/common/mod.rs b/lib/segment/src/common/mod.rs index 4344bc289b..84f032b395 100644 --- a/lib/segment/src/common/mod.rs +++ b/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 OperationResult<()> + Send>; - /// Check that the given vector name is part of the segment config. /// /// Returns an error if incompatible. diff --git a/tests/low-disk/create_and_search_items.py b/tests/low-disk/create_and_search_items.py index 164442ab41..07a7c79cc5 100644 --- a/tests/low-disk/create_and_search_items.py +++ b/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") diff --git a/tests/low-disk/low-disk.sh b/tests/low-disk/low-disk.sh index bbdf7f9c20..1f2960a06d 100755 --- a/tests/low-disk/low-disk.sh +++ b/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" \ No newline at end of file