use std::collections::{BTreeMap, HashMap, HashSet};
use std::io::{Cursor, Seek};
use std::path::Path;
use crate::common::types::PointOffsetType;
use fs_err as fs;
use itertools::Itertools;
use rand::prelude::*;
use tempfile::Builder;
use uuid::Uuid;
use super::MutableIdTracker;
use super::change::{MappingChange, read_entry, write_entry};
use super::mappings_storage::{load_mappings, mappings_path, read_mappings};
use super::versions_storage::{
VERSION_ELEMENT_SIZE, load_versions, store_version_changes, versions_path,
};
use crate::segment::id_tracker::compressed::compressed_point_mappings::CompressedPointMappings;
use crate::segment::id_tracker::in_memory_id_tracker::InMemoryIdTracker;
use crate::segment::id_tracker::{IdTracker, IdTrackerRead};
use crate::segment::types::{PointIdType, SeqNumberType};
const RAND_SEED: u64 = 42;
const DEFAULT_VERSION: SeqNumberType = 42;
#[test]
fn test_iterator() {
let segment_dir = Builder::new().prefix("segment_dir").tempdir().unwrap();
let mut id_tracker = MutableIdTracker::open(segment_dir.path(), None).unwrap();
id_tracker.set_link(200.into(), 0).unwrap();
id_tracker.set_link(100.into(), 1).unwrap();
id_tracker.set_link(150.into(), 2).unwrap();
id_tracker.set_link(120.into(), 3).unwrap();
id_tracker.set_link(180.into(), 4).unwrap();
id_tracker.set_link(110.into(), 5).unwrap();
id_tracker.set_link(115.into(), 6).unwrap();
id_tracker.set_link(190.into(), 7).unwrap();
id_tracker.set_link(177.into(), 8).unwrap();
id_tracker.set_link(118.into(), 9).unwrap();
let first_four = id_tracker
.point_mappings()
.iter_from(None)
.take(4)
.collect_vec();
assert_eq!(first_four.len(), 4);
assert_eq!(first_four[0].0, 100.into());
let last = id_tracker
.point_mappings()
.iter_from(Some(first_four[3].0))
.collect_vec();
assert_eq!(last.len(), 7);
}
pub const TEST_POINTS: &[PointIdType] = &[
PointIdType::NumId(100),
PointIdType::Uuid(Uuid::from_u128(123_u128)),
PointIdType::Uuid(Uuid::from_u128(156_u128)),
PointIdType::NumId(150),
PointIdType::NumId(120),
PointIdType::Uuid(Uuid::from_u128(12_u128)),
PointIdType::NumId(180),
PointIdType::NumId(110),
PointIdType::NumId(115),
PointIdType::Uuid(Uuid::from_u128(673_u128)),
PointIdType::NumId(190),
PointIdType::NumId(177),
PointIdType::Uuid(Uuid::from_u128(971_u128)),
];
#[test]
fn test_mixed_types_iterator() {
let segment_dir = Builder::new().prefix("segment_dir").tempdir().unwrap();
let id_tracker = make_mutable_tracker(segment_dir.path());
let sorted_from_tracker = id_tracker
.point_mappings()
.iter_from(None)
.map(|(k, _)| k)
.collect_vec();
let mut values = TEST_POINTS.to_vec();
values.sort();
assert_eq!(sorted_from_tracker, values);
}
#[test]
fn test_load_store() {
let segment_dir = Builder::new().prefix("segment_dir").tempdir().unwrap();
let (old_mappings, old_versions) = {
let id_tracker = make_mutable_tracker(segment_dir.path());
(id_tracker.mappings, id_tracker.internal_to_version)
};
let mut loaded_id_tracker = MutableIdTracker::open(segment_dir.path(), None).unwrap();
assert_eq!(
old_versions.len(),
loaded_id_tracker.internal_to_version.len(),
);
for i in 0..old_versions.len() {
assert_eq!(
old_versions.get(i),
loaded_id_tracker.internal_to_version.get(i),
"Version mismatch at index {i}",
);
}
assert_eq!(old_mappings, loaded_id_tracker.mappings);
loaded_id_tracker.drop(PointIdType::NumId(180)).unwrap();
}
#[test]
fn test_store_load_mutated() {
let mut rng = SmallRng::seed_from_u64(RAND_SEED);
let segment_dir = Builder::new().prefix("segment_dir").tempdir().unwrap();
let (dropped_points, custom_version) = {
let mut id_tracker = make_mutable_tracker(segment_dir.path());
let mut dropped_points = HashSet::new();
let mut custom_version = HashMap::new();
for (index, point) in TEST_POINTS.iter().enumerate() {
if index % 2 == 0 {
continue;
}
if index % 3 == 0 {
id_tracker.drop(*point).unwrap();
dropped_points.insert(*point);
continue;
}
if index % 5 == 0 {
let new_version = rng.next_u64();
id_tracker
.set_internal_version(index as PointOffsetType, new_version)
.unwrap();
custom_version.insert(index as PointOffsetType, new_version);
}
}
id_tracker.mapping_flusher()().unwrap();
id_tracker.versions_flusher()().unwrap();
(dropped_points, custom_version)
};
let id_tracker = MutableIdTracker::open(segment_dir.path(), None).unwrap();
for (index, point) in TEST_POINTS.iter().enumerate() {
let internal_id = index as PointOffsetType;
if dropped_points.contains(point) {
assert!(id_tracker.is_deleted_point(internal_id));
assert_eq!(id_tracker.external_id(internal_id), None);
assert!(
id_tracker
.mappings
.internal_id_with_behavior(point, common::types::DeferredBehavior::VisibleOnly)
.is_none()
);
continue;
}
let expect_version = custom_version
.get(&internal_id)
.copied()
.unwrap_or(DEFAULT_VERSION);
assert_eq!(
id_tracker.internal_version(internal_id),
Some(expect_version),
);
assert_eq!(
id_tracker.external_id(index as PointOffsetType),
Some(*point),
);
}
}
#[test]
fn test_all_points_have_version() {
let segment_dir = Builder::new().prefix("segment_dir").tempdir().unwrap();
let id_tracker = make_mutable_tracker(segment_dir.path());
for i in id_tracker.point_mappings().iter_internal() {
assert!(id_tracker.internal_version(i).is_some());
}
}
#[test]
fn test_point_deletion_correctness() {
let segment_dir = Builder::new().prefix("segment_dir").tempdir().unwrap();
let mut id_tracker = make_mutable_tracker(segment_dir.path());
let deleted_points = id_tracker.total_point_count() - id_tracker.available_point_count();
let point_to_delete = PointIdType::NumId(100);
assert!(
id_tracker
.point_mappings()
.iter_external()
.contains(&point_to_delete)
);
assert_eq!(
id_tracker.internal_id_with_behavior(
point_to_delete,
common::types::DeferredBehavior::VisibleOnly
),
Some(0)
);
id_tracker.drop(point_to_delete).unwrap();
let point_exists = id_tracker
.internal_id_with_behavior(
point_to_delete,
crate::common::types::DeferredBehavior::VisibleOnly,
)
.is_some()
&& id_tracker
.point_mappings()
.iter_external()
.contains(&point_to_delete)
&& id_tracker
.point_mappings()
.iter_from(None)
.any(|i| i.0 == point_to_delete);
assert!(!point_exists);
let new_deleted_points = id_tracker.total_point_count() - id_tracker.available_point_count();
assert_eq!(new_deleted_points, deleted_points + 1);
}
#[test]
fn test_point_deletion_persists_reload() {
let segment_dir = Builder::new().prefix("segment_dir").tempdir().unwrap();
let point_to_delete = PointIdType::NumId(100);
let old_mappings = {
let mut id_tracker = make_mutable_tracker(segment_dir.path());
let intetrnal_id = id_tracker
.internal_id_with_behavior(
point_to_delete,
crate::common::types::DeferredBehavior::VisibleOnly,
)
.expect("Point to delete exists.");
assert!(!id_tracker.is_deleted_point(intetrnal_id));
id_tracker.drop(point_to_delete).unwrap();
id_tracker.mapping_flusher()().unwrap();
id_tracker.versions_flusher()().unwrap();
id_tracker.mappings
};
let id_tracker = MutableIdTracker::open(segment_dir.path(), None).unwrap();
assert_eq!(
id_tracker.internal_id_with_behavior(
point_to_delete,
common::types::DeferredBehavior::VisibleOnly
),
None
);
old_mappings
.iter_internal_raw()
.zip(id_tracker.mappings.iter_internal_raw())
.for_each(
|((old_internal, old_external), (new_internal, new_external))| {
assert_eq!(old_internal, new_internal);
assert_eq!(old_external, new_external);
},
);
}
#[test]
fn test_point_mappings_de_serialization_single() {
let mut rng = SmallRng::seed_from_u64(RAND_SEED);
const SIZE: usize = 400_000;
let mappings = CompressedPointMappings::random(&mut rng, SIZE as u32);
for i in 0..SIZE {
let mut buf = vec![];
let internal_id = i as PointOffsetType;
let expected_external = mappings.external_id(internal_id).unwrap();
let change = MappingChange::Insert(expected_external, internal_id);
write_entry(&mut buf, change).unwrap();
let (got_change, bytes_read) = read_entry(&mut buf.as_slice()).unwrap();
assert_eq!(change, got_change);
assert!(bytes_read == 13 || bytes_read == 21);
}
}
#[test]
fn test_point_mappings_deserializing_special() {
let buf = Cursor::new(b"");
assert_eq!(read_mappings(buf, None).unwrap().total_point_count(), 0);
let buf = Cursor::new(b"\x00");
assert!(
read_mappings(buf, None)
.unwrap_err()
.to_string()
.contains("Corrupted ID tracker mapping storage")
);
let buf = Cursor::new(b"malformed!");
assert!(read_mappings(buf, None).is_err());
let buf = Cursor::new(b"\x01\x01\x00\x00\x00\x00");
assert_eq!(read_mappings(buf, None).unwrap().total_point_count(), 0);
let buf = Cursor::new(b"\x01\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00");
assert_eq!(
read_mappings(buf, None).unwrap().internal_id_with_behavior(
&PointIdType::NumId(1),
common::types::DeferredBehavior::VisibleOnly,
),
Some(2)
);
let buf = Cursor::new(b"\x01\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x01\x00");
assert!(
read_mappings(buf, None)
.unwrap_err()
.to_string()
.contains("Corrupted ID tracker mapping storage")
);
let buf = Cursor::new(b"\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x03\x01\x00");
let mappings = read_mappings(buf, None).unwrap();
assert_eq!(mappings.total_point_count(), 1);
assert_eq!(
mappings.internal_id_with_behavior(
&PointIdType::NumId(1),
common::types::DeferredBehavior::VisibleOnly
),
Some(0)
);
}
#[test]
fn test_operation_size() {
let items = [
(
b"\x01\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00".as_slice(),
MappingChange::Insert(PointIdType::NumId(1), 2),
),
(
b"\x02\x10\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00"
.as_slice(),
MappingChange::Insert(
PointIdType::Uuid(Uuid::parse_str("10000000-0000-0000-0000-000000000000").unwrap()),
2,
),
),
(
b"\x03\x01\x00\x00\x00\x00\x00\x00\x00".as_slice(),
MappingChange::Delete(PointIdType::NumId(1)),
),
(
b"\x04\x10\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00".as_slice(),
MappingChange::Delete(PointIdType::Uuid(
Uuid::parse_str("10000000-0000-0000-0000-000000000000").unwrap(),
)),
),
];
for (buf, expected_change) in items {
let mut buf = Cursor::new(buf);
let (entry, bytes_read) = read_entry(&mut buf).unwrap();
assert_eq!(entry, expected_change);
assert_eq!(
bytes_read,
buf.stream_position().unwrap(),
"read different number of bytes from stream than returned",
);
}
}
#[test]
fn test_point_mappings_truncation() {
let segment_dir = Builder::new().prefix("segment_dir").tempdir().unwrap();
let mappings_path = mappings_path(segment_dir.path());
fs::write(
&mappings_path,
b"\x01\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00",
)
.unwrap();
assert_eq!(fs::metadata(&mappings_path).unwrap().len(), 13);
assert_eq!(
load_mappings(&mappings_path, None)
.unwrap()
.0
.internal_id_with_behavior(
&PointIdType::NumId(1),
common::types::DeferredBehavior::VisibleOnly,
),
Some(2)
);
assert_eq!(fs::metadata(&mappings_path).unwrap().len(), 13);
fs::write(
&mappings_path,
b"\x01\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x01",
)
.unwrap();
assert_eq!(fs::metadata(&mappings_path).unwrap().len(), 14);
assert_eq!(
load_mappings(&mappings_path, None)
.unwrap()
.0
.internal_id_with_behavior(
&PointIdType::NumId(1),
common::types::DeferredBehavior::VisibleOnly,
),
Some(2)
);
assert_eq!(fs::metadata(&mappings_path).unwrap().len(), 13);
fs::write(
&mappings_path,
b"\x01\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x01\x02\x00",
)
.unwrap();
assert_eq!(fs::metadata(&mappings_path).unwrap().len(), 16);
assert_eq!(
load_mappings(&mappings_path, None)
.unwrap()
.0
.internal_id_with_behavior(
&PointIdType::NumId(1),
common::types::DeferredBehavior::VisibleOnly,
),
Some(2)
);
assert_eq!(fs::metadata(&mappings_path).unwrap().len(), 13);
fs::write(
&mappings_path,
b"\x01\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x01\x02\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x02\x00",
).unwrap();
assert_eq!(fs::metadata(&mappings_path).unwrap().len(), 28);
assert_eq!(
load_mappings(&mappings_path, None)
.unwrap()
.0
.internal_id_with_behavior(
&PointIdType::NumId(1),
common::types::DeferredBehavior::VisibleOnly,
),
Some(2)
);
assert_eq!(fs::metadata(&mappings_path).unwrap().len(), 26);
}
fn make_in_memory_tracker_from_memory() -> InMemoryIdTracker {
let mut id_tracker = InMemoryIdTracker::new();
for value in TEST_POINTS.iter() {
let internal_id = id_tracker.total_point_count() as PointOffsetType;
id_tracker.set_link(*value, internal_id).unwrap();
id_tracker
.set_internal_version(internal_id, DEFAULT_VERSION)
.unwrap()
}
id_tracker
}
fn make_mutable_tracker(path: &Path) -> MutableIdTracker {
let mut id_tracker =
MutableIdTracker::open(path, None).expect("failed to open mutable ID tracker");
for value in TEST_POINTS.iter() {
let internal_id = id_tracker.total_point_count() as PointOffsetType;
id_tracker.set_link(*value, internal_id).unwrap();
id_tracker
.set_internal_version(internal_id, DEFAULT_VERSION)
.unwrap()
}
id_tracker.mapping_flusher()().expect("failed to flush ID tracker mappings");
id_tracker.versions_flusher()().expect("failed to flush ID tracker versions");
id_tracker
}
#[test]
fn test_id_tracker_equal() {
let in_memory_id_tracker = make_in_memory_tracker_from_memory();
let mutable_id_tracker_dir = Builder::new()
.prefix("segment_dir_mutable")
.tempdir()
.unwrap();
let mutable_id_tracker = make_mutable_tracker(mutable_id_tracker_dir.path());
assert_eq!(
in_memory_id_tracker.available_point_count(),
mutable_id_tracker.available_point_count(),
);
assert_eq!(
in_memory_id_tracker.total_point_count(),
mutable_id_tracker.total_point_count(),
);
for (internal, external) in TEST_POINTS.iter().enumerate() {
let internal = internal as PointOffsetType;
assert_eq!(
in_memory_id_tracker
.internal_id_with_behavior(*external, common::types::DeferredBehavior::VisibleOnly),
mutable_id_tracker
.internal_id_with_behavior(*external, common::types::DeferredBehavior::VisibleOnly),
);
assert_eq!(
in_memory_id_tracker
.internal_version(internal)
.unwrap_or_default(),
mutable_id_tracker
.internal_version(internal)
.unwrap_or_default(),
);
assert_eq!(
in_memory_id_tracker.external_id(internal),
mutable_id_tracker.external_id(internal),
);
}
}
#[test]
#[cfg(false)]
fn simple_id_tracker_vs_mutable_tracker_congruence() {
use crate::segment::common::rocksdb_wrapper::{DB_VECTOR_CF, open_db};
let segment_dir = Builder::new().prefix("segment_dir").tempdir().unwrap();
let db = open_db(segment_dir.path(), &[DB_VECTOR_CF]).unwrap();
let mut mutable_id_tracker = MutableIdTracker::open(segment_dir.path(), None).unwrap();
let mut simple_id_tracker = SimpleIdTracker::open(db).unwrap();
let num_points = 200;
let mut rng = SmallRng::seed_from_u64(RAND_SEED);
for _ in 0..num_points {
let point_id = PointIdType::NumId(rng.random_range(0..num_points as u64));
let version = rng.random_range(0..1000);
let internal_id_mmap = mutable_id_tracker.total_point_count() as PointOffsetType;
let internal_id_simple = simple_id_tracker.total_point_count() as PointOffsetType;
assert_eq!(internal_id_mmap, internal_id_simple);
if mutable_id_tracker
.internal_id_with_behavior(point_id, crate::common::types::DeferredBehavior::VisibleOnly)
.is_some()
{
mutable_id_tracker.drop(point_id).unwrap();
}
mutable_id_tracker
.set_link(point_id, internal_id_mmap)
.unwrap();
mutable_id_tracker
.set_internal_version(internal_id_mmap, version)
.unwrap();
if simple_id_tracker
.internal_id_with_behavior(point_id, crate::common::types::DeferredBehavior::VisibleOnly)
.is_some()
{
simple_id_tracker.drop(point_id).unwrap();
}
simple_id_tracker
.set_link(point_id, internal_id_simple)
.unwrap();
simple_id_tracker
.set_internal_version(internal_id_simple, version)
.unwrap();
}
fn check_trackers(a: &SimpleIdTracker, b: &MutableIdTracker) {
for (external_id, internal_id) in a.point_mappings().iter_from(None) {
assert_eq!(
a.internal_version(internal_id).unwrap(),
b.internal_version(internal_id).unwrap()
);
assert_eq!(a.external_id(internal_id), b.external_id(internal_id));
assert_eq!(external_id, b.external_id(internal_id).unwrap());
assert_eq!(
a.external_id(internal_id).unwrap(),
b.external_id(internal_id).unwrap()
);
}
for (external_id, internal_id) in b.point_mappings().iter_from(None) {
assert_eq!(
a.internal_version(internal_id).unwrap(),
b.internal_version(internal_id).unwrap()
);
assert_eq!(a.external_id(internal_id), b.external_id(internal_id));
assert_eq!(external_id, a.external_id(internal_id).unwrap());
assert_eq!(
a.external_id(internal_id).unwrap(),
b.external_id(internal_id).unwrap()
);
}
}
check_trackers(&simple_id_tracker, &mutable_id_tracker);
mutable_id_tracker.mapping_flusher()().unwrap();
mutable_id_tracker.versions_flusher()().unwrap();
drop(mutable_id_tracker);
let mutable_id_tracker = MutableIdTracker::open(segment_dir.path(), None).unwrap();
check_trackers(&simple_id_tracker, &mutable_id_tracker);
}
#[test]
fn test_load_versions_partial_tail_ignored() {
let dir = Builder::new().prefix("segment_dir").tempdir().unwrap();
let path = versions_path(dir.path());
let mut data = Vec::new();
for v in [10u64, 20, 30] {
data.extend_from_slice(&v.to_le_bytes());
}
data.extend_from_slice(&[0xFF; 5]);
fs::write(&path, &data).unwrap();
assert_eq!(fs::metadata(&path).unwrap().len(), 29);
let versions = load_versions(&path).unwrap();
assert_eq!(versions, vec![10, 20, 30]);
assert_eq!(fs::metadata(&path).unwrap().len(), 29);
}
#[test]
fn test_store_versions_truncates_partial_tail() {
let dir = Builder::new().prefix("segment_dir").tempdir().unwrap();
let path = versions_path(dir.path());
let mut data = Vec::new();
for v in [10u64, 20] {
data.extend_from_slice(&v.to_le_bytes());
}
data.extend_from_slice(&[0xFF; 3]);
fs::write(&path, &data).unwrap();
assert_eq!(fs::metadata(&path).unwrap().len(), 19);
let mut changes = BTreeMap::new();
changes.insert(1u32, 99u64);
store_version_changes(&path, &changes).unwrap();
let file_len = fs::metadata(&path).unwrap().len();
assert_eq!(file_len % VERSION_ELEMENT_SIZE, 0);
let versions = load_versions(&path).unwrap();
assert_eq!(versions.len(), 2);
assert_eq!(versions[0], 10); assert_eq!(versions[1], 99); }
#[test]
fn test_store_versions_extends_without_set_len() {
let dir = Builder::new().prefix("segment_dir").tempdir().unwrap();
let path = versions_path(dir.path());
let mut changes = BTreeMap::new();
changes.insert(0u32, 100u64);
changes.insert(1u32, 200u64);
store_version_changes(&path, &changes).unwrap();
let versions = load_versions(&path).unwrap();
assert_eq!(versions, vec![100, 200]);
let mut changes = BTreeMap::new();
changes.insert(5u32, 500u64);
store_version_changes(&path, &changes).unwrap();
let versions = load_versions(&path).unwrap();
assert_eq!(versions.len(), 6);
assert_eq!(versions[0], 100);
assert_eq!(versions[1], 200);
assert_eq!(versions[2], 0);
assert_eq!(versions[3], 0);
assert_eq!(versions[4], 0);
assert_eq!(versions[5], 500);
}
#[test]
fn shadow_visible_head_survives_mapping_flush_reload() {
use crate::common::types::DeferredBehavior;
let segment_dir = Builder::new().prefix("segment_dir").tempdir().unwrap();
let cutoff = Some(5 as PointOffsetType);
let p7 = PointIdType::NumId(7);
{
let mut id_tracker = MutableIdTracker::open(segment_dir.path(), cutoff).unwrap();
id_tracker.set_link(p7, 2).unwrap();
id_tracker.set_internal_version(2, 1).unwrap();
id_tracker.set_link(p7, 9).unwrap();
id_tracker.set_internal_version(9, 2).unwrap();
assert_eq!(
id_tracker.internal_id_with_behavior(p7, DeferredBehavior::VisibleOnly),
Some(2),
"live: the visible (active) head should resolve",
);
assert_eq!(
id_tracker.internal_id_with_behavior(p7, DeferredBehavior::WithDeferred),
Some(9),
"live: the visible (active) head should resolve",
);
id_tracker.mapping_flusher()().unwrap();
id_tracker.versions_flusher()().unwrap();
}
let id_tracker = MutableIdTracker::open(segment_dir.path(), cutoff).unwrap();
assert_eq!(
id_tracker.internal_id_with_behavior(p7, DeferredBehavior::VisibleOnly),
Some(2)
);
assert_eq!(
id_tracker.internal_id_with_behavior(p7, DeferredBehavior::WithDeferred),
Some(9),
);
assert_eq!(id_tracker.external_id(2), Some(p7));
assert!(!id_tracker.is_deleted_point(2));
}