Spaces:
Build error
Build error
| import multiprocessing | |
| import pathlib | |
| from time import sleep | |
| from .fixtures import upsert_random_points, create_collection | |
| from .utils import * | |
| N_PEERS = 3 | |
| N_SHARDS = 3 | |
| N_REPLICA = 1 | |
| COLLECTION_NAME = "test_collection" | |
| def update_points_in_loop(peer_url, collection_name, offset=0, throttle=False, duration=None): | |
| start = time.time() | |
| limit = 3 | |
| while True: | |
| upsert_random_points(peer_url, limit, collection_name, offset=offset) | |
| offset += limit | |
| if throttle: | |
| sleep(0.1) | |
| if duration is not None and (time.time() - start) > duration: | |
| break | |
| def run_update_points_in_background(peer_url, collection_name, init_offset=0, throttle=False, duration=None): | |
| p = multiprocessing.Process(target=update_points_in_loop, args=(peer_url, collection_name, init_offset, throttle, duration)) | |
| p.start() | |
| return p | |
| # Transfer shards from one node to another while applying updates in parallel | |
| # Test that data on the both sides is consistent | |
| def test_shard_snapshot_transfer_cancel(tmp_path: pathlib.Path): | |
| assert_project_root() | |
| # seed port to reuse the same port for the restarted nodes | |
| peer_api_uris, peer_dirs, bootstrap_uri = start_cluster(tmp_path, N_PEERS, 20000) | |
| create_collection(peer_api_uris[0], shard_number=N_SHARDS, replication_factor=N_REPLICA) | |
| wait_collection_exists_and_active_on_all_peers( | |
| collection_name=COLLECTION_NAME, | |
| peer_api_uris=peer_api_uris | |
| ) | |
| # Insert some initial number of points | |
| upsert_random_points(peer_api_uris[0], 100) | |
| transfer_collection_cluster_info = get_collection_cluster_info(peer_api_uris[0], COLLECTION_NAME) | |
| receiver_collection_cluster_info = get_collection_cluster_info(peer_api_uris[2], COLLECTION_NAME) | |
| from_peer_id = transfer_collection_cluster_info['peer_id'] | |
| to_peer_id = receiver_collection_cluster_info['peer_id'] | |
| shard_id = transfer_collection_cluster_info['local_shards'][0]['shard_id'] | |
| # Transfer shard from one node to another | |
| # Move shard `shard_id` to peer `target_peer_id` | |
| r = requests.post( | |
| f"{peer_api_uris[0]}/collections/{COLLECTION_NAME}/cluster", json={ | |
| "replicate_shard": { | |
| "shard_id": shard_id, | |
| "from_peer_id": from_peer_id, | |
| "to_peer_id": to_peer_id, | |
| "method": "snapshot" | |
| } | |
| }) | |
| assert_http_ok(r) | |
| # Cancel shard transfer | |
| r = requests.post( | |
| f"{peer_api_uris[0]}/collections/{COLLECTION_NAME}/cluster", json={ | |
| "abort_transfer": { | |
| "shard_id": shard_id, | |
| "from_peer_id": from_peer_id, | |
| "to_peer_id": to_peer_id, | |
| } | |
| }) | |
| assert_http_ok(r) | |
| r = requests.post( | |
| f"{peer_api_uris[0]}/collections/{COLLECTION_NAME}/cluster", json={ | |
| "replicate_shard": { | |
| "shard_id": shard_id, | |
| "from_peer_id": from_peer_id, | |
| "to_peer_id": to_peer_id, | |
| "method": "snapshot" | |
| } | |
| }) | |
| assert_http_ok(r) | |
| # Wait for end of shard transfer | |
| wait_for_collection_shard_transfers_count(peer_api_uris[0], COLLECTION_NAME, 0) | |
| receiver_collection_cluster_info = get_collection_cluster_info(peer_api_uris[2], COLLECTION_NAME) | |
| number_local_shards = len(receiver_collection_cluster_info['local_shards']) | |
| assert number_local_shards == 2 | |
| # Point counts must be consistent across nodes | |
| counts = [] | |
| for uri in peer_api_uris: | |
| r = requests.post( | |
| f"{uri}/collections/{COLLECTION_NAME}/points/count", json={ | |
| "exact": True | |
| } | |
| ) | |
| assert_http_ok(r) | |
| counts.append(r.json()["result"]['count']) | |
| assert counts[0] == counts[1] == counts[2] == 100 | |
| # Transfer shards from one node to another and cancel while applying updates in parallel | |
| # | |
| # A fast stream of updates is sent to the cluster, the queue proxy will not be | |
| # able to keep up with it transferring all updates to the receiving node. If we | |
| # cancel ater 5 seconds during this process, cancellation should still happen | |
| # normally within 3 seconds. | |
| # | |
| # Test that data on the both sides is consistent | |
| def test_shard_snapshot_transfer_cancel_during_updates(tmp_path: pathlib.Path): | |
| assert_project_root() | |
| # seed port to reuse the same port for the restarted nodes | |
| peer_api_uris, peer_dirs, bootstrap_uri = start_cluster(tmp_path, N_PEERS, 20000) | |
| create_collection(peer_api_uris[0], shard_number=N_SHARDS, replication_factor=N_REPLICA) | |
| wait_collection_exists_and_active_on_all_peers( | |
| collection_name=COLLECTION_NAME, | |
| peer_api_uris=peer_api_uris | |
| ) | |
| # Insert some initial number of points | |
| upsert_random_points(peer_api_uris[0], 10000) | |
| # Start pushing points to the cluster | |
| upload_process_1 = run_update_points_in_background(peer_api_uris[0], COLLECTION_NAME, init_offset=100) | |
| upload_process_2 = run_update_points_in_background(peer_api_uris[1], COLLECTION_NAME, init_offset=10000) | |
| upload_process_3 = run_update_points_in_background(peer_api_uris[2], COLLECTION_NAME, init_offset=20000) | |
| transfer_collection_cluster_info = get_collection_cluster_info(peer_api_uris[0], COLLECTION_NAME) | |
| receiver_collection_cluster_info = get_collection_cluster_info(peer_api_uris[2], COLLECTION_NAME) | |
| from_peer_id = transfer_collection_cluster_info['peer_id'] | |
| to_peer_id = receiver_collection_cluster_info['peer_id'] | |
| shard_id = transfer_collection_cluster_info['local_shards'][0]['shard_id'] | |
| # Transfer shard from one node to another | |
| # Move shard `shard_id` to peer `target_peer_id` | |
| r = requests.post( | |
| f"{peer_api_uris[0]}/collections/{COLLECTION_NAME}/cluster", json={ | |
| "move_shard": { | |
| "shard_id": shard_id, | |
| "from_peer_id": from_peer_id, | |
| "to_peer_id": to_peer_id, | |
| "method": "snapshot" | |
| } | |
| }) | |
| assert_http_ok(r) | |
| upload_process_1.kill() | |
| upload_process_2.kill() | |
| upload_process_3.kill() | |
| # Cancel shard transfer | |
| cancel_at = time.time() | |
| r = requests.post( | |
| f"{peer_api_uris[0]}/collections/{COLLECTION_NAME}/cluster", json={ | |
| "abort_transfer": { | |
| "shard_id": shard_id, | |
| "from_peer_id": from_peer_id, | |
| "to_peer_id": to_peer_id, | |
| } | |
| }) | |
| assert_http_ok(r) | |
| # Wait for end of shard transfer | |
| wait_for_collection_shard_transfers_count(peer_api_uris[2], COLLECTION_NAME, 0) | |
| # Cancellation must happen within 3 seconds | |
| assert (time.time() - cancel_at) <= 3 | |
| # Should have only one local shard since transfer was cancelled | |
| receiver_collection_cluster_info = get_collection_cluster_info(peer_api_uris[2], COLLECTION_NAME) | |
| number_local_shards = len(receiver_collection_cluster_info['local_shards']) | |
| assert number_local_shards == 1 | |
| # Point counts must be consistent across nodes | |
| counts = [] | |
| for uri in peer_api_uris: | |
| r = requests.post( | |
| f"{uri}/collections/{COLLECTION_NAME}/points/count", json={ | |
| "exact": True | |
| } | |
| ) | |
| assert_http_ok(r) | |
| counts.append(r.json()["result"]['count']) | |
| assert counts[0] == counts[1] == counts[2] | |