From 20740b9e98aaa59840a70b2f9caae85f301a6a6b Mon Sep 17 00:00:00 2001 From: Di Zhao Date: Tue, 12 Sep 2023 11:21:19 -0700 Subject: [PATCH] Add update_concurrency option to control the parallelism of shard updates (#2599) * add update_concurrency to control the parallelism of updating shards * reformat the files * reformat * add update_concurrency control to shard updates * use unwrap_or * test using buffer_ordered * use buffered (testing) * Pre-allocate space for update futures * use NonZeroUsize * linter * use NonZeroUsize in the test * consensus test for update_concurrency * update config commit --------- Co-authored-by: Di Zhao Co-authored-by: timvisee Co-authored-by: generall --- config/config.yaml | 5 +++ .../src/operations/shared_storage_config.rs | 6 ++- lib/collection/src/shards/replica_set.rs | 37 +++++++++---------- lib/storage/src/types.rs | 4 ++ lib/storage/tests/integration/alias_tests.rs | 4 +- tests/consensus_tests/fixtures.py | 11 ++++++ .../test_update_concurrency.py | 35 ++++++++++++++++++ tests/consensus_tests/utils.py | 17 ++++++--- 8 files changed, 93 insertions(+), 26 deletions(-) create mode 100644 tests/consensus_tests/test_update_concurrency.py diff --git a/config/config.yaml b/config/config.yaml index 93c569e9df..ad1445b972 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -17,6 +17,10 @@ storage: # Note: those payload values that are involved in filtering and are indexed - remain in RAM. on_disk_payload: true + # Maximum number of concurrent updates to shard replicas + # If `null` - maximum concurrency is used. + update_concurrency: null + # Write-ahead-log related configuration wal: # Size of a single WAL segment @@ -109,6 +113,7 @@ storage: # Custom M param for hnsw graph built for payload index. If not set, default M will be used. payload_m: null + service: # Maximum size of POST data in a single request in megabytes diff --git a/lib/collection/src/operations/shared_storage_config.rs b/lib/collection/src/operations/shared_storage_config.rs index 8aaf766989..853257d5ac 100644 --- a/lib/collection/src/operations/shared_storage_config.rs +++ b/lib/collection/src/operations/shared_storage_config.rs @@ -1,3 +1,4 @@ +use std::num::NonZeroUsize; use std::time::Duration; use crate::operations::types::NodeType; @@ -18,6 +19,7 @@ pub struct SharedStorageConfig { pub handle_collection_load_errors: bool, pub recovery_mode: Option, pub search_timeout: Duration, + pub update_concurrency: Option, } impl Default for SharedStorageConfig { @@ -28,6 +30,7 @@ impl Default for SharedStorageConfig { handle_collection_load_errors: false, recovery_mode: None, search_timeout: DEFAULT_SEARCH_TIMEOUT, + update_concurrency: None, } } } @@ -39,18 +42,19 @@ impl SharedStorageConfig { handle_collection_load_errors: bool, recovery_mode: Option, search_timeout: Option, + update_concurrency: Option, ) -> Self { let update_queue_size = update_queue_size.unwrap_or(match node_type { NodeType::Normal => DEFAULT_UPDATE_QUEUE_SIZE, NodeType::Listener => DEFAULT_UPDATE_QUEUE_SIZE_LISTENER, }); - Self { update_queue_size, node_type, handle_collection_load_errors, recovery_mode, search_timeout: search_timeout.unwrap_or(DEFAULT_SEARCH_TIMEOUT), + update_concurrency, } } } diff --git a/lib/collection/src/shards/replica_set.rs b/lib/collection/src/shards/replica_set.rs index 452da156a6..0e0fa1b326 100644 --- a/lib/collection/src/shards/replica_set.rs +++ b/lib/collection/src/shards/replica_set.rs @@ -7,7 +7,7 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::Duration; -use futures::future::{join, join_all, BoxFuture}; +use futures::future::{BoxFuture, Either}; use futures::stream::FuturesUnordered; use futures::{FutureExt, StreamExt}; use itertools::Itertools; @@ -1580,19 +1580,19 @@ impl ShardReplicaSet { ))); } - let mut remote_futures = Vec::new(); + let mut update_futures = Vec::with_capacity(active_remote_shards.len() + 1); for remote in active_remote_shards { - let op = operation.clone(); - remote_futures.push(async move { + let operation = operation.clone(); + update_futures.push(Either::Left(async move { remote - .update(op, wait) + .update(operation, wait) .await .map_err(|err| (remote.peer_id, err)) - }); + })); } - match local.deref() { - Some(local) if self.peer_is_active_or_pending(&this_peer_id) => { + if let Some(local) = local.deref() { + if self.peer_is_active_or_pending(&this_peer_id) { let local_wait = if self.peer_state(&this_peer_id) == Some(ReplicaState::Listener) { false @@ -1611,18 +1611,17 @@ impl ShardReplicaSet { (peer_id, err) }) }; - let remote_updates = join_all(remote_futures); - - // run local and remote shards read concurrently - let (mut remote_res, local_res): ( - Vec>, - _, - ) = join(remote_updates, local_update).await; - // return both remote and local results - remote_res.push(local_res); - remote_res + update_futures.push(Either::Right(local_update)); } - _ => join_all(remote_futures).await, + } + match self.shared_storage_config.update_concurrency { + Some(concurrency) => { + futures::stream::iter(update_futures) + .buffered(concurrency.get()) + .collect::>() + .await + } + _ => futures::future::join_all(update_futures).await, } }; diff --git a/lib/storage/src/types.rs b/lib/storage/src/types.rs index bf983ee2ce..11757c6375 100644 --- a/lib/storage/src/types.rs +++ b/lib/storage/src/types.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::num::NonZeroUsize; use std::time::Duration; use chrono::{DateTime, Utc}; @@ -69,6 +70,8 @@ pub struct StorageConfig { /// Provided value will be used error message for unavailable requests. #[serde(default)] pub recovery_mode: Option, + #[serde(default)] + pub update_concurrency: Option, } impl StorageConfig { @@ -81,6 +84,7 @@ impl StorageConfig { self.performance .search_timeout_sec .map(|x| Duration::from_secs(x as u64)), + self.update_concurrency, ) } } diff --git a/lib/storage/tests/integration/alias_tests.rs b/lib/storage/tests/integration/alias_tests.rs index 3f50d5349b..c91895c278 100644 --- a/lib/storage/tests/integration/alias_tests.rs +++ b/lib/storage/tests/integration/alias_tests.rs @@ -1,4 +1,4 @@ -use std::num::NonZeroU64; +use std::num::{NonZeroU64, NonZeroUsize}; use std::sync::Arc; use collection::operations::types::VectorParams; @@ -55,6 +55,8 @@ fn test_alias_operation() { handle_collection_load_errors: false, recovery_mode: None, async_scorer: false, + update_concurrency: Some(NonZeroUsize::new(2).unwrap()), + // update_concurrency: None, }; let search_runtime = Runtime::new().unwrap(); diff --git a/tests/consensus_tests/fixtures.py b/tests/consensus_tests/fixtures.py index 0b95890643..4938b21623 100644 --- a/tests/consensus_tests/fixtures.py +++ b/tests/consensus_tests/fixtures.py @@ -67,3 +67,14 @@ def search(peer_url, vector, city, collection="test_collection"): r_search = requests.post(f"{peer_url}/collections/{collection}/points/search", json=q) assert_http_ok(r_search) return r_search.json()["result"] + + +def count_counts(peer_url, collection="test_collection"): + r_search = requests.post( + f"{peer_url}/collections/{collection}/points/count", + json={ + "exact": True, + } + ) + assert_http_ok(r_search) + return r_search.json()["result"]["count"] diff --git a/tests/consensus_tests/test_update_concurrency.py b/tests/consensus_tests/test_update_concurrency.py new file mode 100644 index 0000000000..011e776c94 --- /dev/null +++ b/tests/consensus_tests/test_update_concurrency.py @@ -0,0 +1,35 @@ +import logging +import pathlib + +from .fixtures import create_collection, upsert_random_points, count_counts +from .utils import * + +logging.basicConfig(level=logging.DEBUG) +logger = logging.getLogger() + +N_PEERS = 3 +N_SHARDS = 1 +N_REPLICAS = 2 +COLLECTION_NAME = "test_collection" + + +def test_listener_node(tmp_path: pathlib.Path): + + assert_project_root() + + peer_api_uris, peer_dirs, bootstrap_uri = start_cluster( + tmp_path, + N_PEERS, + extra_env={ + "QDRANT__STORAGE__UPDATE_CONCURRENCY": "1", + } + ) + + create_collection(peer_api_uris[0], shard_number=N_SHARDS, replication_factor=N_REPLICAS) + wait_collection_exists_and_active_on_all_peers(collection_name=COLLECTION_NAME, peer_api_uris=peer_api_uris) + + upsert_random_points(peer_api_uris[0], 100) + + for uri in peer_api_uris: + count = count_counts(uri, COLLECTION_NAME) + assert count == 100 diff --git a/tests/consensus_tests/utils.py b/tests/consensus_tests/utils.py index aaad232df6..5b6a9129da 100644 --- a/tests/consensus_tests/utils.py +++ b/tests/consensus_tests/utils.py @@ -95,11 +95,17 @@ def start_peer(peer_dir: Path, log_file: str, bootstrap_uri: str, port=None, ext # Starts a peer and returns its api_uri and p2p_uri -def start_first_peer(peer_dir: Path, log_file: str, port=None) -> Tuple[str, str]: +def start_first_peer(peer_dir: Path, log_file: str, port=None, extra_env=None) -> Tuple[str, str]: + if extra_env is None: + extra_env = {} + p2p_port = get_port() if port is None else port + 0 grpc_port = get_port() if port is None else port + 1 http_port = get_port() if port is None else port + 2 - env = get_env(p2p_port, grpc_port, http_port) + env = { + **get_env(p2p_port, grpc_port, http_port), + **extra_env + } test_log_folder = init_pytest_log_folder() log_file = open(f"{test_log_folder}/{log_file}", "w") bootstrap_uri = get_uri(p2p_port) @@ -110,7 +116,7 @@ def start_first_peer(peer_dir: Path, log_file: str, port=None) -> Tuple[str, str return get_uri(http_port), bootstrap_uri -def start_cluster(tmp_path, num_peers, port_seed=None): +def start_cluster(tmp_path, num_peers, port_seed=None, extra_env=None): assert_project_root() peer_dirs = make_peer_folders(tmp_path, num_peers) @@ -118,7 +124,8 @@ def start_cluster(tmp_path, num_peers, port_seed=None): peer_api_uris = [] # Start bootstrap - (bootstrap_api_uri, bootstrap_uri) = start_first_peer(peer_dirs[0], "peer_0_0.log", port=port_seed) + (bootstrap_api_uri, bootstrap_uri) = start_first_peer(peer_dirs[0], "peer_0_0.log", port=port_seed, + extra_env=extra_env) peer_api_uris.append(bootstrap_api_uri) # Wait for leader @@ -129,7 +136,7 @@ def start_cluster(tmp_path, num_peers, port_seed=None): for i in range(1, len(peer_dirs)): if port_seed is not None: port = port_seed + i * 100 - peer_api_uris.append(start_peer(peer_dirs[i], f"peer_0_{i}.log", bootstrap_uri, port=port)) + peer_api_uris.append(start_peer(peer_dirs[i], f"peer_0_{i}.log", bootstrap_uri, port=port, extra_env=extra_env)) # Wait for cluster wait_for_uniform_cluster_status(peer_api_uris, leader)