simple integration test for tenant promotion (#7503)

* simple integration test for tenant promotion

* add cleanup after promotion

* When switching to ReadActive state, set it for the correct peer ID

* fix test

---------

Co-authored-by: timvisee <tim@visee.me>

lib/collection/src/shards/shard_holder/mod.rs

@@ -572,6 +572,8 @@ impl ShardHolder {
                 let (shard_ids_to_query, used_shard_key) =
                     self.route_with_fallback_for_read(key)?;
 
+                log::trace!("Search routing with fallback: {used_shard_key:?}");
+
                 for shard_id in shard_ids_to_query {
                     if let Some(replica_set) = self.shards.get(&shard_id) {
                         res.push((replica_set, Some(used_shard_key)));

lib/collection/src/shards/transfer/mod.rs

@@ -283,6 +283,7 @@ pub trait ShardTransferConsensus: Send + Sync {
 
             result = self
                 .set_shard_replica_set_state(
+                    Some(remote_shard.peer_id),
                     collection_id.clone(),
                     remote_shard.id,
                     ReplicaState::ActiveRead,
@@ -480,12 +481,15 @@ pub trait ShardTransferConsensus: Send + Sync {
 
     /// Set the shard replica state on this peer through consensus
     ///
+    /// If the peer ID is not provided, this will set the replica state for the current peer.
+    ///
     /// # Warning
     ///
     /// This only submits a proposal to consensus. Calling this does not guarantee that consensus
     /// will actually apply the operation across the cluster.
     async fn set_shard_replica_set_state(
         &self,
+        peer_id: Option<PeerId>,
         collection_id: CollectionId,
         shard_id: ShardId,
         state: ReplicaState,
@@ -498,6 +502,7 @@ pub trait ShardTransferConsensus: Send + Sync {
     /// operation.
     async fn set_shard_replica_set_state_confirm_and_retry(
         &self,
+        peer_id: Option<PeerId>,
         collection_id: &CollectionId,
         shard_id: ShardId,
         state: ReplicaState,
@@ -515,7 +520,13 @@ pub trait ShardTransferConsensus: Send + Sync {
 
             log::trace!("Propose and confirm set shard replica set state");
             result = self
-                .set_shard_replica_set_state(collection_id.clone(), shard_id, state, from_state)
+                .set_shard_replica_set_state(
+                    peer_id,
+                    collection_id.clone(),
+                    shard_id,
+                    state,
+                    from_state,
+                )
                 .await;
 
             match &result {

lib/storage/src/content_manager/toc/transfer.rs

@@ -125,6 +125,7 @@ impl ShardTransferConsensus for TocDispatcher {
 
     async fn set_shard_replica_set_state(
         &self,
+        peer_id: Option<PeerId>,
         collection_id: CollectionId,
         shard_id: ShardId,
         state: ReplicaState,
@@ -133,7 +134,7 @@ impl ShardTransferConsensus for TocDispatcher {
         let operation = ConsensusOperations::set_replica_state(
             collection_id,
             shard_id,
-            self.this_peer_id(),
+            peer_id.unwrap_or_else(|| self.this_peer_id()),
             state,
             from_state,
         );

tests/consensus_tests/custom_sharding.py

@@ -0,0 +1,59 @@
+from .utils import *
+from .fixtures import *
+
+DEFAULT_COLLECTION_NAME = "test_collection"
+
+
+def create_collection_with_custom_sharding(
+        peer_url,
+        collection=DEFAULT_COLLECTION_NAME,
+        shard_number=1,
+        replication_factor=1,
+        write_consistency_factor=1,
+        timeout=10
+):
+    create_collection(
+        peer_url,
+        collection=collection,
+        shard_number=shard_number,
+        replication_factor=replication_factor,
+        write_consistency_factor=write_consistency_factor,
+        timeout=timeout,
+        sharding_method="custom",
+    )
+
+
+def create_shard(
+        peer_url,
+        collection,
+        shard_key,
+        shard_number=1,
+        replication_factor=1,
+        placement=None,
+        initial_state=None,
+        timeout=10
+):
+    r_batch = requests.put(
+        f"{peer_url}/collections/{collection}/shards?timeout={timeout}", json={
+            "shard_key": shard_key,
+            "shards_number": shard_number,
+            "replication_factor": replication_factor,
+            "placement": placement,
+            "initial_state": initial_state,
+        })
+    assert_http_ok(r_batch)
+
+
+def delete_shard(
+        peer_url,
+        collection,
+        shard_key,
+        timeout=10
+):
+    r_batch = requests.post(
+        f"{peer_url}/collections/{collection}/shards/delete?timeout={timeout}",
+        json={
+            "shard_key": shard_key,
+        }
+    )
+    assert_http_ok(r_batch)

tests/consensus_tests/fixtures.py

@@ -93,7 +93,9 @@ def upsert_random_points(
     shard_key=None,
     num_cities=None,
     headers={},
+    extra_payload=None,
 ):
+    extra_payload = extra_payload or {}
 
     def get_vector():
         # Create points in first peer's collection
@@ -115,7 +117,10 @@ def upsert_random_points(
                     {
                         "id": i + offset,
                         "vector": get_vector(),
-                        "payload": {"city": random.choice(CITIES[:num_cities]) if num_cities is not None else random.choice(CITIES)},
+                        "payload": {
+                            **extra_payload,
+                            "city": random.choice(CITIES[:num_cities]) if num_cities is not None else random.choice(CITIES)
+                        },
                     }
                     for i in range(size)
                 ],

tests/consensus_tests/test_custom_sharding.py

@@ -1,8 +1,8 @@
 import concurrent.futures
 import pathlib
 import threading
-import time
 
+from .custom_sharding import create_collection_with_custom_sharding, create_shard, delete_shard
 from .fixtures import *
 from .utils import *
 
@@ -12,66 +12,6 @@ N_REPLICAS = 1
 
 COLLECTION_NAME = "test_collection"
 
-
-def create_collection_with_custom_sharding(
-        peer_url,
-        collection=COLLECTION_NAME,
-        shard_number=1,
-        replication_factor=1,
-        write_consistency_factor=1,
-        timeout=10
-):
-    # Create collection in peer_url
-    r_batch = requests.put(
-        f"{peer_url}/collections/{collection}?timeout={timeout}", json={
-            "vectors": {
-                "size": 4,
-                "distance": "Dot"
-            },
-            "shard_number": shard_number,
-            "replication_factor": replication_factor,
-            "write_consistency_factor": write_consistency_factor,
-            "sharding_method": "custom",
-        })
-    assert_http_ok(r_batch)
-
-
-def create_shard(
-        peer_url,
-        collection,
-        shard_key,
-        shard_number=1,
-        replication_factor=1,
-        placement=None,
-        initial_state=None,
-        timeout=10
-):
-    r_batch = requests.put(
-        f"{peer_url}/collections/{collection}/shards?timeout={timeout}", json={
-            "shard_key": shard_key,
-            "shards_number": shard_number,
-            "replication_factor": replication_factor,
-            "placement": placement,
-            "initial_state": initial_state,
-        })
-    assert_http_ok(r_batch)
-
-
-def delete_shard(
-        peer_url,
-        collection,
-        shard_key,
-        timeout=10
-):
-    r_batch = requests.post(
-        f"{peer_url}/collections/{collection}/shards/delete?timeout={timeout}",
-        json={
-            "shard_key": shard_key,
-        }
-    )
-    assert_http_ok(r_batch)
-
-
 def test_shard_consistency(tmp_path: pathlib.Path):
     assert_project_root()
 

tests/consensus_tests/test_tenant_promotion.py

@@ -0,0 +1,196 @@
+import concurrent.futures
+import pathlib
+import threading
+import time
+
+from .custom_sharding import create_collection_with_custom_sharding, create_shard
+from .fixtures import *
+from .utils import *
+
+COLLECTION_NAME = "tenant_collection"
+
+N_PEERS = 3
+N_SHARDS = 1
+N_REPLICAS = 1
+
+
+def match_results(results1: List[dict], results2: List[dict]):
+    # Compare ids and payloads of search results
+    if len(results1) != len(results2):
+        return False, f"Results length differ: {len(results1)} != {len(results2)}"
+
+    for idx, (r1, r2) in enumerate(zip(results1, results2)):
+        if r1['id'] != r2['id']:
+            return False, f"Result IDs differ: {r1['id']} != {r2['id']} at index {idx}"
+        if r1.get('payload') != r2.get('payload'):
+            return False, f"Result payloads differ for ID {r1['id']}: {r1.get('payload')} != {r2.get('payload')}, at index {idx}"
+
+    return True, ""
+
+
+def search_points(peer_url):
+    query_vector = [0.1, 0.2, 0.3, 0.4]
+    resp = requests.post(
+        f"{peer_url}/collections/{COLLECTION_NAME}/points/search",
+        json={
+            "vector": query_vector,
+            "top": 5,
+            "filter": {
+                "must": [
+                    {
+                        "key": "tenant",
+                        "match": {
+                            "value": "tenant1"
+                        }
+                    }
+                ]
+            },
+            "shard_key": {
+                "fallback": "default",
+                "target": "tenant1"
+            }
+        }
+    )
+    assert_http_ok(resp)
+
+    results = resp.json()["result"]
+    return results
+
+
+def test_tenant_promotion_simple(tmp_path: pathlib.Path):
+    assert_project_root()
+
+    peer_api_uris, peer_dirs, bootstrap_uri = start_cluster(tmp_path, N_PEERS)
+
+    create_collection_with_custom_sharding(peer_api_uris[0], collection=COLLECTION_NAME, shard_number=N_SHARDS,
+                                           replication_factor=N_REPLICAS)
+
+    create_shard(
+        peer_url=peer_api_uris[0],
+        collection=COLLECTION_NAME,
+        shard_key="default",
+        shard_number=N_SHARDS,
+        replication_factor=N_REPLICAS,
+        # initial_state="Partial",
+    )
+
+    n_points = 1_000
+    upsert_random_points(
+        peer_api_uris[0],
+        n_points,
+        collection_name=COLLECTION_NAME,
+        shard_key={
+            "fallback": "default",
+            "target": "tenant1"
+        },
+        extra_payload={
+            "tenant": "tenant1"
+        }
+    )
+
+    n_points = 100
+    upsert_random_points(
+        peer_api_uris[0],
+        n_points,
+        collection_name=COLLECTION_NAME,
+        shard_key={
+            "fallback": "default",
+            "target": "tenant2"
+        },
+        extra_payload={
+            "tenant": "tenant2"
+        }
+    )
+
+    # Search point with shard key and filter
+    results = search_points(peer_api_uris[0])
+    assert len(results) > 0, "No results found for tenant1"
+
+    create_shard(
+        peer_url=peer_api_uris[0],
+        collection=COLLECTION_NAME,
+        shard_key="tenant1",
+        shard_number=N_SHARDS,
+        replication_factor=N_REPLICAS,
+        initial_state="Partial",
+    )
+
+    # With new shard we should get the same results
+    result2 = search_points(peer_api_uris[0])
+    match, msg = match_results(results, result2)
+    assert match, f"Results differ after tenant1 shard creation: {msg}"
+
+    n_points = 100
+    upsert_random_points(
+        peer_api_uris[0],
+        n_points,
+        collection_name=COLLECTION_NAME,
+        shard_key={
+            "fallback": "default",
+            "target": "tenant1"
+        },
+        extra_payload={
+            "tenant": "tenant1"
+        }
+    )
+
+    results = search_points(peer_api_uris[0])
+
+    # Promote to dedicated shard
+    resp = requests.post(
+        f"{peer_api_uris[0]}/collections/{COLLECTION_NAME}/cluster",
+        json={
+            "replicate_points": {
+                "filter": {
+                    "must": {
+                        "key": "tenant",
+                        "match": {
+                            "value": "tenant1"
+                        }
+                    }
+                },
+                "from_shard_key": "default",
+                "to_shard_key": "tenant1"
+            }
+        }
+    )
+    assert_http_ok(resp)
+
+    results2 = search_points(peer_api_uris[0])
+    match, msg = match_results(results, results2)
+    assert match, f"Results differ after tenant1 promotion start: {msg}"
+
+    # Wait all transfers complete on all peers
+    for peer_url in peer_api_uris:
+        wait_for_collection_shard_transfers_count(peer_url, COLLECTION_NAME, 0)
+        # Check that new shard have active state
+        info = get_collection_cluster_info(peer_api_uris[0], COLLECTION_NAME)
+        for shard in info["local_shards"]:
+            assert shard['state'] == 'Active'
+        for shard in info["remote_shards"]:
+            assert shard['state'] == 'Active'
+
+    results3 = search_points(peer_api_uris[0])
+    match, msg = match_results(results, results3)
+    assert match, f"Results differ after tenant1 promotion complete: {msg}"
+
+    # Delete records from the old shard
+    resp = requests.post(
+        f"{peer_api_uris[0]}/collections/{COLLECTION_NAME}/points/delete?wait=true",
+        json={
+            "filter": {
+                "must": {
+                    "key": "tenant",
+                    "match": {
+                        "value": "tenant1"
+                    }
+                }
+            },
+            "shard_key": "default"
+        }
+    )
+    assert_http_ok(resp)
+
+    results4 = search_points(peer_api_uris[0])
+    match, msg = match_results(results, results4)
+    assert match, f"Results differ after tenant1 deletion from old shard: {msg}"