use linkrs_core::types::{EdgeId, Timestamp};
use linkrs_core::{StorageError, StorageResult};
use super::super::{
CsrBase, CsrVariant, ImmutableCsr, MutableCsrTrait, Nbr, INVALID_EDGE_ID, NO_EDGES_STORED_MSG,
};
use super::CsrShardSet;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FreezeBlockReason {
MissingGroup,
NoEdgesStored,
AlreadyFrozen,
}
impl FreezeBlockReason {
pub fn message(self, gid: usize) -> String {
match self {
FreezeBlockReason::MissingGroup => {
format!("group {} missing, nothing to freeze", gid)
}
FreezeBlockReason::NoEdgesStored => NO_EDGES_STORED_MSG.to_string(),
FreezeBlockReason::AlreadyFrozen => {
format!("group {} is already frozen", gid)
}
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FreezeFeasibility {
Ready {
edges: u64,
},
Blocked { reason: FreezeBlockReason },
}
impl FreezeFeasibility {
pub fn is_ready(self) -> bool {
matches!(self, FreezeFeasibility::Ready { .. })
}
}
impl CsrShardSet {
pub fn is_frozen(&self, gid: usize) -> bool {
self.shards
.get(&gid)
.is_some_and(|shard| matches!(shard.variant, CsrVariant::Frozen(_)))
}
pub fn is_group_frozen_for(&self, vid: u32) -> bool {
let Some((gid, _)) = self.route(vid) else {
return false;
};
self.is_frozen(gid)
}
pub fn freeze_group(
&mut self,
gid: usize,
cutoff: Timestamp,
reserve_ratio: f32,
on_edge_removed: &mut dyn FnMut(EdgeId, Timestamp),
) -> StorageResult<u64> {
if let Err(reason) = self.check_freeze_gate(gid) {
return Err(StorageError::invalid_operation(reason.message(gid)));
}
self.compact_group_with_reporting(gid, cutoff, reserve_ratio, on_edge_removed);
let shard = self.shards.get_mut(&gid).ok_or_else(|| {
StorageError::invalid_operation(format!("group {} missing on freeze", gid))
})?;
let frozen = match &shard.variant {
CsrVariant::Multiple(csr) => ImmutableCsr::pack_from_mutable(csr),
CsrVariant::Single(csr) => ImmutableCsr::pack_single_from(csr),
CsrVariant::Pure(csr) => ImmutableCsr::pack_from_pure(csr),
CsrVariant::Bundled(csr) => ImmutableCsr::pack_from_bundled(csr),
CsrVariant::Frozen(_) | CsrVariant::None { .. } => {
return Err(StorageError::invalid_operation(format!(
"group {} changed under freeze",
gid
)));
}
};
let packed_edges = frozen.edge_count();
shard.variant = CsrVariant::Frozen(Box::new(frozen));
shard.mapped = None;
shard.dirty.deleted = true;
for region in shard.regions.iter_mut() {
region.deleted = true;
}
shard.reclaim_hint = false;
self.clear_group_append_log(gid);
Ok(packed_edges)
}
pub fn freeze_feasibility(&self, gid: usize) -> FreezeFeasibility {
match self.check_freeze_gate(gid) {
Ok(()) => {
let edges = self
.shards
.get(&gid)
.map(|shard| shard.variant.edge_count())
.unwrap_or(0);
FreezeFeasibility::Ready { edges }
}
Err(reason) => FreezeFeasibility::Blocked { reason },
}
}
fn check_freeze_gate(&self, gid: usize) -> Result<(), FreezeBlockReason> {
let Some(shard) = self.shards.get(&gid) else {
return Err(FreezeBlockReason::MissingGroup);
};
if shard.variant.is_writable() {
return Ok(());
}
if shard.variant.is_read_only() {
return Err(FreezeBlockReason::AlreadyFrozen);
}
Err(FreezeBlockReason::NoEdgesStored)
}
pub fn unfreeze_group(&mut self, gid: usize) -> StorageResult<u64> {
let frozen = match self.shards.get(&gid) {
Some(shard) => match &shard.variant {
CsrVariant::Frozen(csr) => csr.clone(),
_ => {
return Err(StorageError::invalid_operation(format!(
"group {} is not frozen",
gid
)));
}
},
None => {
return Err(StorageError::invalid_operation(format!(
"group {} missing, nothing to unfreeze",
gid
)));
}
};
let mut variant = self.fresh_variant()?;
let rows = frozen.vertex_capacity();
let mut restored_tombstones = 0usize;
if variant.is_bundled() {
let mut valued_buf: Vec<(Nbr, Option<u64>)> = Vec::new();
for local in 0..rows {
valued_buf.clear();
frozen.visit_physical_with_values(local as u32, |nbr, value| {
valued_buf.push((nbr, value));
true
});
for (nbr, value) in &valued_buf {
if nbr.edge_id == INVALID_EDGE_ID {
continue;
}
variant.insert_edge_with_value(
local as u32,
nbr.to_vertex_id(),
nbr.edge_id,
*value,
)?;
if nbr.delete_ts != Timestamp::MAX {
let deleted =
variant.delete_edge(local as u32, nbr.edge_id, nbr.delete_ts)?;
if !deleted {
return Err(StorageError::invalid_operation(format!(
"unfreeze replay missed edge {:?} in group {}",
nbr.edge_id, gid
)));
}
restored_tombstones += 1;
}
}
}
} else {
let mut row_buf = Vec::new();
for local in 0..rows {
frozen.fill_physical_into(local as u32, &mut row_buf);
for nbr in &row_buf {
if nbr.edge_id == INVALID_EDGE_ID {
continue;
}
variant.insert_edge(
local as u32,
nbr.to_vertex_id(),
nbr.edge_id,
Timestamp::MAX,
)?;
if nbr.delete_ts != Timestamp::MAX {
let deleted =
variant.delete_edge(local as u32, nbr.edge_id, nbr.delete_ts)?;
if !deleted {
return Err(StorageError::invalid_operation(format!(
"unfreeze replay missed edge {:?} in group {}",
nbr.edge_id, gid
)));
}
restored_tombstones += 1;
}
}
}
}
let restored = frozen.edge_count();
let shard = self.shards.get_mut(&gid).ok_or_else(|| {
StorageError::invalid_operation(format!("group {} missing on unfreeze", gid))
})?;
shard.variant = variant;
shard.mapped = None;
shard.dirty.deleted = true;
for region in shard.regions.iter_mut() {
region.deleted = true;
}
shard.dead_entries = restored_tombstones;
shard.reclaim_hint = restored_tombstones > 0;
self.clear_group_append_log(gid);
Ok(restored)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::edge::csr_trait::CsrBase;
use linkrs_core::types::{EdgeStrategy, VertexId};
use super::super::DEFAULT_NODE_GROUP_BITS;
use super::super::RecordForm;
fn sample_set() -> CsrShardSet {
let mut set = CsrShardSet::new(
EdgeStrategy::Multiple,
DEFAULT_NODE_GROUP_BITS,
64,
RecordForm::Columnar,
)
.unwrap();
for src in 0..32u32 {
for k in 0..3u32 {
let dst = VertexId::edge_endpoint_key(src + k + 1, 0);
set.insert_edge(src, dst, EdgeId((src * 10 + k) as u64), 1)
.unwrap();
}
}
set.delete_edge(0, EdgeId(1), 5).unwrap();
set
}
#[test]
fn freeze_rejects_writes_and_unfreeze_restores() {
let mut set = sample_set();
let before: Vec<_> = (0..32u32).map(|src| set.physical_edges_of(src)).collect();
let packed = set
.freeze_group(0, Timestamp::MAX, 0.0, &mut |_, _| {})
.unwrap();
assert_eq!(packed, set.edge_count());
assert!(set.is_frozen(0));
assert!(set.is_group_frozen_for(3));
assert!(!set.is_group_frozen_for(1 << DEFAULT_NODE_GROUP_BITS));
for src in 0..32u32 {
assert_eq!(set.physical_edges_of(src), before[src as usize]);
let expected: Vec<_> = before[src as usize]
.iter()
.filter(|nbr| nbr.is_alive_at(6))
.copied()
.collect();
assert_eq!(set.edges_of(src, 6), expected);
}
assert!(set
.insert_edge(1, VertexId::edge_endpoint_key(99, 0), EdgeId(1000), 9)
.is_err());
assert!(set.delete_edge(1, EdgeId(10), 9).is_err());
assert_eq!(
set.delete_edge_by_dst(1, VertexId::edge_endpoint_key(11, 0), 9),
0
);
let restored = set.unfreeze_group(0).unwrap();
assert_eq!(restored, packed);
assert!(!set.is_frozen(0));
for src in 0..32u32 {
assert_eq!(set.physical_edges_of(src), before[src as usize]);
}
set.insert_edge(1, VertexId::edge_endpoint_key(99, 0), EdgeId(1000), 9)
.unwrap();
assert_eq!(set.edge_count(), packed + 1);
}
#[test]
fn freeze_guards_bad_states() {
let mut set = sample_set();
assert!(set
.freeze_group(77, Timestamp::MAX, 0.0, &mut |_, _| {})
.is_err());
set.freeze_group(0, Timestamp::MAX, 0.0, &mut |_, _| {})
.unwrap();
assert!(set
.freeze_group(0, Timestamp::MAX, 0.0, &mut |_, _| {})
.is_err());
assert!(set.unfreeze_group(77).is_err());
let mut none_set = CsrShardSet::new(
EdgeStrategy::None,
DEFAULT_NODE_GROUP_BITS,
64,
RecordForm::Columnar,
)
.unwrap();
assert!(none_set
.freeze_group(0, Timestamp::MAX, 0.0, &mut |_, _| {})
.is_err());
let mut single = CsrShardSet::new(
EdgeStrategy::Single,
DEFAULT_NODE_GROUP_BITS,
64,
RecordForm::Columnar,
)
.unwrap();
single
.insert_edge(5, VertexId::edge_endpoint_key(6, 0), EdgeId(1), 1)
.unwrap();
let packed = single
.freeze_group(0, Timestamp::MAX, 0.0, &mut |_, _| {})
.unwrap();
assert_eq!(packed, 1);
assert_eq!(
single.edges_of(5, 1),
vec![single.physical_edges_of(5).pop().unwrap()]
);
single.unfreeze_group(0).unwrap();
assert!(!single.is_frozen(0));
assert_eq!(single.edge_count(), 1);
}
#[test]
fn feasibility_matches_freeze_outcome() {
let mut set = sample_set();
let quoted = set.freeze_feasibility(0);
assert!(quoted.is_ready());
let packed = set
.freeze_group(0, Timestamp::MAX, 0.0, &mut |_, _| {})
.unwrap();
assert_eq!(quoted, FreezeFeasibility::Ready { edges: packed });
let blocked = set.freeze_feasibility(0);
assert_eq!(
blocked,
FreezeFeasibility::Blocked {
reason: FreezeBlockReason::AlreadyFrozen
}
);
assert!(!blocked.is_ready());
let err = set
.freeze_group(0, Timestamp::MAX, 0.0, &mut |_, _| {})
.unwrap_err();
assert!(err
.to_string()
.contains(&FreezeBlockReason::AlreadyFrozen.message(0)));
assert_eq!(
set.freeze_feasibility(77),
FreezeFeasibility::Blocked {
reason: FreezeBlockReason::MissingGroup
}
);
assert!(set
.freeze_group(77, Timestamp::MAX, 0.0, &mut |_, _| {})
.is_err());
}
#[test]
fn valued_bundled_groups_freeze_with_values_intact() {
let mut valued = CsrShardSet::new(
EdgeStrategy::Multiple,
DEFAULT_NODE_GROUP_BITS,
64,
RecordForm::Bundled,
)
.unwrap();
valued
.bundled_insert_with_value(0, VertexId::edge_endpoint_key(1, 0), EdgeId(0), 1, Some(7))
.unwrap();
valued
.bundled_insert_with_value(0, VertexId::edge_endpoint_key(2, 0), EdgeId(1), 1, None)
.unwrap();
assert!(valued.freeze_feasibility(0).is_ready());
valued
.freeze_group(0, Timestamp::MAX, 0.0, &mut |_, _| {})
.unwrap();
assert!(valued.is_frozen(0));
let frozen = valued.group_variant(0).expect("frozen group present");
assert_eq!(
frozen.bundled_value_by_edge_id(0, EdgeId(0)),
Some((7, true))
);
assert_eq!(
frozen.bundled_value_by_edge_id(0, EdgeId(1)),
Some((0, false))
);
let mut seen = Vec::new();
valued.visit_physical_with_values(0, |nbr, value| {
seen.push((nbr.edge_id, value));
true
});
assert_eq!(seen, vec![(EdgeId(0), Some(7)), (EdgeId(1), None)]);
valued.unfreeze_group(0).unwrap();
assert!(!valued.is_frozen(0));
let live = valued.group_variant(0).expect("live group present");
assert_eq!(live.bundled_value_by_edge_id(0, EdgeId(0)), Some((7, true)));
assert_eq!(
live.bundled_value_by_edge_id(0, EdgeId(1)),
Some((0, false))
);
let mut nulls = CsrShardSet::new(
EdgeStrategy::Multiple,
DEFAULT_NODE_GROUP_BITS,
64,
RecordForm::Bundled,
)
.unwrap();
nulls
.bundled_insert_with_value(0, VertexId::edge_endpoint_key(1, 0), EdgeId(0), 1, None)
.unwrap();
assert!(nulls.freeze_feasibility(0).is_ready());
nulls
.freeze_group(0, Timestamp::MAX, 0.0, &mut |_, _| {})
.unwrap();
}
}