Spaces:
Build error
Build error
File size: 5,598 Bytes
84d2a97 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 |
use std::collections::{HashMap, HashSet};
use std::num::NonZeroU32;
use std::sync::Arc;
use common::cpu::CpuBudget;
use segment::types::Distance;
use tempfile::Builder;
use crate::collection::{Collection, RequestShardTransfer};
use crate::config::{CollectionConfigInternal, CollectionParams, WalConfig};
use crate::operations::shared_storage_config::SharedStorageConfig;
use crate::operations::types::{NodeType, VectorsConfig};
use crate::operations::vector_params_builder::VectorParamsBuilder;
use crate::shards::channel_service::ChannelService;
use crate::shards::collection_shard_distribution::CollectionShardDistribution;
use crate::shards::replica_set::{AbortShardTransfer, ChangePeerFromState};
use crate::tests::fixtures::TEST_OPTIMIZERS_CONFIG;
pub fn dummy_on_replica_failure() -> ChangePeerFromState {
Arc::new(move |_peer_id, _shard_id, _from_state| {})
}
pub fn dummy_request_shard_transfer() -> RequestShardTransfer {
Arc::new(move |_transfer| {})
}
pub fn dummy_abort_shard_transfer() -> AbortShardTransfer {
Arc::new(|_transfer, _reason| {})
}
fn init_logger() {
let _ = env_logger::builder().is_test(true).try_init();
}
async fn _test_snapshot_collection(node_type: NodeType) {
let wal_config = WalConfig {
wal_capacity_mb: 1,
wal_segments_ahead: 0,
};
let collection_params = CollectionParams {
vectors: VectorsConfig::Single(VectorParamsBuilder::new(4, Distance::Dot).build()),
shard_number: NonZeroU32::new(4).unwrap(),
replication_factor: NonZeroU32::new(3).unwrap(),
write_consistency_factor: NonZeroU32::new(2).unwrap(),
..CollectionParams::empty()
};
let config = CollectionConfigInternal {
params: collection_params,
optimizer_config: TEST_OPTIMIZERS_CONFIG.clone(),
wal_config,
hnsw_config: Default::default(),
quantization_config: Default::default(),
strict_mode_config: Default::default(),
uuid: None,
};
let snapshots_path = Builder::new().prefix("test_snapshots").tempdir().unwrap();
let collection_dir = Builder::new().prefix("test_collection").tempdir().unwrap();
let collection_name = "test".to_string();
let collection_name_rec = "test_rec".to_string();
let mut shards = HashMap::new();
shards.insert(0, HashSet::from([1]));
shards.insert(1, HashSet::from([1]));
shards.insert(2, HashSet::from([10_000])); // remote shard
shards.insert(3, HashSet::from([1, 20_000, 30_000]));
let storage_config: SharedStorageConfig = SharedStorageConfig {
node_type,
..Default::default()
};
let collection = Collection::new(
collection_name,
1,
collection_dir.path(),
snapshots_path.path(),
&config,
Arc::new(storage_config),
CollectionShardDistribution { shards },
ChannelService::default(),
dummy_on_replica_failure(),
dummy_request_shard_transfer(),
dummy_abort_shard_transfer(),
None,
None,
CpuBudget::default(),
None,
)
.await
.unwrap();
let snapshots_temp_dir = Builder::new().prefix("temp_dir").tempdir().unwrap();
let snapshot_description = collection
.create_snapshot(snapshots_temp_dir.path(), 0)
.await
.unwrap();
assert_eq!(snapshot_description.checksum.unwrap().len(), 64);
{
let recover_dir = Builder::new()
.prefix("test_collection_rec")
.tempdir()
.unwrap();
// Do not recover in local mode if some shards are remote
assert!(Collection::restore_snapshot(
&snapshots_path.path().join(&snapshot_description.name),
recover_dir.path(),
0,
false,
)
.is_err());
}
let recover_dir = Builder::new()
.prefix("test_collection_rec")
.tempdir()
.unwrap();
if let Err(err) = Collection::restore_snapshot(
&snapshots_path.path().join(snapshot_description.name),
recover_dir.path(),
0,
true,
) {
panic!("Failed to restore snapshot: {err}")
}
let recovered_collection = Collection::load(
collection_name_rec,
1,
recover_dir.path(),
snapshots_path.path(),
Default::default(),
ChannelService::default(),
dummy_on_replica_failure(),
dummy_request_shard_transfer(),
dummy_abort_shard_transfer(),
None,
None,
CpuBudget::default(),
None,
)
.await;
{
let shards_holder = &recovered_collection.shards_holder.read().await;
let replica_ser_0 = shards_holder.get_shard(0).unwrap();
assert!(replica_ser_0.is_local().await);
let replica_ser_1 = shards_holder.get_shard(1).unwrap();
assert!(replica_ser_1.is_local().await);
let replica_ser_2 = shards_holder.get_shard(2).unwrap();
assert!(!replica_ser_2.is_local().await);
assert_eq!(replica_ser_2.peers().len(), 1);
let replica_ser_3 = shards_holder.get_shard(3).unwrap();
assert!(replica_ser_3.is_local().await);
assert_eq!(replica_ser_3.peers().len(), 3); // 2 remotes + 1 local
}
}
#[tokio::test(flavor = "multi_thread")]
async fn test_snapshot_collection_normal() {
init_logger();
_test_snapshot_collection(NodeType::Normal).await;
}
#[tokio::test(flavor = "multi_thread")]
async fn test_snapshot_collection_listener() {
init_logger();
_test_snapshot_collection(NodeType::Listener).await;
}
|