use std::collections::BTreeSet;
use std::fmt;
use serde::{Deserialize, Serialize};
use crate::cluster::{ClusterEpoch, ClusterNodeId, PartitionId};
use crate::grid::elasticity::{
restore_topology_from_snapshot, ControlPlaneSnapshot, RegionId, SnapshotError,
TopologyAuthority, CONTROL_PLANE_SNAPSHOT_FORMAT_VERSION,
};
use crate::grid::hardening::{
anti_entropy_repair, PromotionFreezeWindow, ReplicatedValueRecord, ReplicatedValueStore,
ValueStoreError,
};
use crate::grid::{EffectiveReplicationMap, PromotionPhase, Replicas};
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RegionState {
#[default]
Up,
Suspect,
Down,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RegionObservation {
pub region: RegionId,
pub missed_heartbeats: u64,
pub quorum_reachable: bool,
pub operator_declared_down: bool,
pub split_brain_risk: bool,
}
impl RegionObservation {
pub fn healthy(region: impl Into<RegionId>) -> Self {
Self {
region: region.into(),
missed_heartbeats: 0,
quorum_reachable: true,
operator_declared_down: false,
split_brain_risk: false,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct RegionStateDetector {
suspect_after_missed: u64,
}
impl RegionStateDetector {
pub fn new(suspect_after_missed: u64) -> Self {
Self {
suspect_after_missed: suspect_after_missed.max(1),
}
}
pub fn classify(&self, observation: &RegionObservation) -> RegionState {
if observation.operator_declared_down && !observation.split_brain_risk {
return RegionState::Down;
}
if observation.operator_declared_down
|| observation.split_brain_risk
|| !observation.quorum_reachable
|| observation.missed_heartbeats >= self.suspect_after_missed
{
return RegionState::Suspect;
}
RegionState::Up
}
}
impl Default for RegionStateDetector {
fn default() -> Self {
Self::new(3)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RegionFailoverError {
message: String,
}
impl RegionFailoverError {
fn new(message: impl Into<String>) -> Self {
Self {
message: message.into(),
}
}
}
impl fmt::Display for RegionFailoverError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.message)
}
}
impl std::error::Error for RegionFailoverError {}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RegionPromotion {
pub partitions: Vec<PartitionId>,
pub from_home: RegionId,
pub to_home: RegionId,
pub epoch: ClusterEpoch,
pub phase: PromotionPhase,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RegionPromotionReport {
pub promotion: RegionPromotion,
pub snapshot: ControlPlaneSnapshot,
pub degraded_partitions: Vec<PartitionId>,
pub freeze: PromotionFreezeWindow,
}
impl RegionPromotionReport {
pub fn is_fully_replicated(&self) -> bool {
self.degraded_partitions.is_empty()
}
}
pub fn promote_region_home(
snapshot: &ControlPlaneSnapshot,
from_home: RegionId,
to_home: RegionId,
from_state: RegionState,
freeze: PromotionFreezeWindow,
) -> Result<RegionPromotionReport, RegionFailoverError> {
if from_home == to_home {
return Err(RegionFailoverError::new(
"from_home and to_home must be different regions",
));
}
if from_state != RegionState::Down {
return Err(RegionFailoverError::new(
"region must be explicitly declared down before promotion",
));
}
if !freeze.is_bounded() {
return Err(RegionFailoverError::new(
"region promotion freeze window exceeded its bound",
));
}
let target_nodes = nodes_in_region(snapshot, &to_home);
let mut next = snapshot.clone();
next.epoch = next_epoch(snapshot.epoch);
let mut promoted = Vec::new();
let mut degraded = Vec::new();
for (partition, replicas) in &snapshot.ownership {
if primary_region(snapshot, replicas) != Some(&from_home) {
continue;
}
promoted.push(*partition);
let Some(primary) = target_nodes.first().cloned() else {
degraded.push(*partition);
continue;
};
let desired_copies = replicas.copy_count();
let mut backups = surviving_backups(snapshot, replicas, &from_home, &primary);
for candidate in target_nodes.iter().skip(1) {
if 1 + backups.len() >= desired_copies {
break;
}
if !backups.contains(candidate) {
backups.push(candidate.clone());
}
}
let promoted_replicas = Replicas::new(primary, backups);
if promoted_replicas.copy_count() < desired_copies {
degraded.push(*partition);
}
next.ownership.insert(*partition, promoted_replicas);
}
Ok(RegionPromotionReport {
promotion: RegionPromotion {
partitions: promoted,
from_home,
to_home,
epoch: next.epoch,
phase: PromotionPhase::Finalize,
},
snapshot: next,
degraded_partitions: degraded,
freeze,
})
}
fn next_epoch(epoch: ClusterEpoch) -> ClusterEpoch {
ClusterEpoch::new(epoch.value().saturating_add(1))
}
fn nodes_in_region(snapshot: &ControlPlaneSnapshot, region: &RegionId) -> Vec<ClusterNodeId> {
snapshot
.topology
.iter()
.filter(|(_, topology)| &topology.region == region)
.map(|(node, _)| node.clone())
.collect()
}
fn primary_region<'a>(
snapshot: &'a ControlPlaneSnapshot,
replicas: &Replicas,
) -> Option<&'a RegionId> {
snapshot
.topology
.get(&replicas.primary)
.map(|topology| &topology.region)
}
fn surviving_backups(
snapshot: &ControlPlaneSnapshot,
replicas: &Replicas,
down_region: &RegionId,
new_primary: &ClusterNodeId,
) -> Vec<ClusterNodeId> {
replicas
.all_nodes()
.into_iter()
.filter(|node| node != new_primary)
.filter(|node| {
snapshot
.topology
.get(node)
.map(|topology| &topology.region != down_region)
.unwrap_or(false)
})
.collect()
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RejoiningRegionDecision {
AcceptAuthority,
FenceLowerEpoch {
current_epoch: ClusterEpoch,
rejoining_epoch: ClusterEpoch,
},
}
pub fn rejoining_region_authority(
current_epoch: ClusterEpoch,
rejoining_epoch: ClusterEpoch,
) -> RejoiningRegionDecision {
if rejoining_epoch < current_epoch {
RejoiningRegionDecision::FenceLowerEpoch {
current_epoch,
rejoining_epoch,
}
} else {
RejoiningRegionDecision::AcceptAuthority
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RegionRestore<S> {
snapshot: ControlPlaneSnapshot,
values: S,
}
impl<S> RegionRestore<S> {
pub fn new(snapshot: ControlPlaneSnapshot, values: S) -> Self {
Self { snapshot, values }
}
pub fn snapshot(&self) -> &ControlPlaneSnapshot {
&self.snapshot
}
pub fn values(&self) -> &S {
&self.values
}
pub fn values_mut(&mut self) -> &mut S {
&mut self.values
}
}
impl<S> RegionRestore<S>
where
S: ReplicatedValueStore,
{
pub fn backfill_from(
&mut self,
records: impl IntoIterator<Item = (String, ReplicatedValueRecord)>,
) -> Result<u64, RegionRestoreError> {
anti_entropy_repair(&mut self.values, records).map_err(RegionRestoreError::from)
}
pub fn restore(self) -> Result<RegionRestoreOutcome<S>, RegionRestoreError> {
if self.snapshot.format_version > CONTROL_PLANE_SNAPSHOT_FORMAT_VERSION {
return Err(RegionRestoreError::from(SnapshotError::new(
"control-plane snapshot format is newer than this binary",
)));
}
let authority =
restore_topology_from_snapshot(&self.snapshot).map_err(RegionRestoreError::from)?;
let restored_value_count = restored_value_count(&self.snapshot, &self.values)?;
let report = RegionRestoreReport {
epoch: self.snapshot.epoch,
topology_node_count: self.snapshot.topology.len(),
partition_count: self.snapshot.ownership.len(),
restored_value_count,
};
Ok(RegionRestoreOutcome {
authority,
values: self.values,
report,
})
}
}
fn restored_value_count<S>(
snapshot: &ControlPlaneSnapshot,
values: &S,
) -> Result<usize, RegionRestoreError>
where
S: ReplicatedValueStore,
{
let mut keys = BTreeSet::new();
for replicas in snapshot.ownership.values() {
let map = EffectiveReplicationMap::new(replicas.clone());
for (key, _) in values.scan_owned(&map).map_err(RegionRestoreError::from)? {
keys.insert(key);
}
}
Ok(keys.len())
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RegionRestoreOutcome<S> {
pub authority: TopologyAuthority,
pub values: S,
pub report: RegionRestoreReport,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RegionRestoreReport {
pub epoch: ClusterEpoch,
pub topology_node_count: usize,
pub partition_count: usize,
pub restored_value_count: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RegionRestoreError {
message: String,
}
impl RegionRestoreError {
fn new(message: impl Into<String>) -> Self {
Self {
message: message.into(),
}
}
}
impl From<SnapshotError> for RegionRestoreError {
fn from(error: SnapshotError) -> Self {
Self::new(error.to_string())
}
}
impl From<ValueStoreError> for RegionRestoreError {
fn from(error: ValueStoreError) -> Self {
Self::new(error.to_string())
}
}
impl fmt::Display for RegionRestoreError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.message)
}
}
impl std::error::Error for RegionRestoreError {}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct RegionFailoverMetrics {
pub region_state: RegionState,
pub region_promotion_total: u64,
pub region_restore_duration_ms: u64,
}