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 <diz@twitter.com>
Co-authored-by: timvisee <tim@visee.me>
Co-authored-by: generall <andrey@vasnetsov.com>

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

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<String>,
     pub search_timeout: Duration,
+    pub update_concurrency: Option<NonZeroUsize>,
 }
 
 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<String>,
         search_timeout: Option<Duration>,
+        update_concurrency: Option<NonZeroUsize>,
     ) -> 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,
         }
     }
 }

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<Result<UpdateResult, (PeerId, CollectionError)>>,
-                        _,
-                    ) = 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::<Vec<_>>()
+                        .await
+                }
+                _ => futures::future::join_all(update_futures).await,
             }
         };
 

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<String>,
+    #[serde(default)]
+    pub update_concurrency: Option<NonZeroUsize>,
 }
 
 impl StorageConfig {
@@ -81,6 +84,7 @@ impl StorageConfig {
             self.performance
                 .search_timeout_sec
                 .map(|x| Duration::from_secs(x as u64)),
+            self.update_concurrency,
         )
     }
 }

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();

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"]

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

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)