use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::fmt;
use serde::{Deserialize, Serialize};
use crate::cluster::{
partition_for_key, ClusterEpoch, ClusterMember, ClusterNodeId, PartitionId,
RendezvousClusterOwnership,
};
use crate::grid::hardening::{
ReplicatedValueRecord, ReplicatedValueStore, ValueStoreError, ValueVersion, WriteWatermark,
REPLICATED_VALUE_RECORD_FORMAT_VERSION,
};
use crate::grid::{ClusterReplicationStrategy, EffectiveReplicationMap, Replicas};
use crate::invalidation_bus::CACHE_INVALIDATION_FRAME_VERSION;
pub const NODE_TOPOLOGY_REGION_METADATA_KEY: &str = "hydracache.topology.region";
pub const NODE_TOPOLOGY_ZONE_METADATA_KEY: &str = "hydracache.topology.zone";
pub const CONTROL_PLANE_SNAPSHOT_FORMAT_VERSION: u32 = 1;
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct RegionId(String);
impl RegionId {
pub fn new(value: impl Into<String>) -> Self {
Self(value.into())
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl From<&str> for RegionId {
fn from(value: &str) -> Self {
Self::new(value)
}
}
impl From<String> for RegionId {
fn from(value: String) -> Self {
Self::new(value)
}
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct ZoneId(String);
impl ZoneId {
pub fn new(value: impl Into<String>) -> Self {
Self(value.into())
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl From<&str> for ZoneId {
fn from(value: &str) -> Self {
Self::new(value)
}
}
impl From<String> for ZoneId {
fn from(value: String) -> Self {
Self::new(value)
}
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct NodeTopology {
pub region: RegionId,
pub zone: ZoneId,
}
impl NodeTopology {
pub fn new(region: impl Into<RegionId>, zone: impl Into<ZoneId>) -> Self {
Self {
region: region.into(),
zone: zone.into(),
}
}
pub fn same_zone(&self, other: &Self) -> bool {
self.region == other.region && self.zone == other.zone
}
pub fn same_region(&self, other: &Self) -> bool {
self.region == other.region
}
fn default_single_zone() -> Self {
Self::new("default", "default")
}
}
pub fn topology_from_member_metadata(member: &ClusterMember) -> Option<NodeTopology> {
let region = member.metadata.get(NODE_TOPOLOGY_REGION_METADATA_KEY)?;
let zone = member.metadata.get(NODE_TOPOLOGY_ZONE_METADATA_KEY)?;
Some(NodeTopology::new(region.clone(), zone.clone()))
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct TopologyAuthority {
epoch: ClusterEpoch,
committed: BTreeMap<ClusterNodeId, NodeTopology>,
observed_gossip: BTreeMap<ClusterNodeId, NodeTopology>,
}
impl TopologyAuthority {
pub fn new() -> Self {
Self::default()
}
pub fn observe_gossip(&mut self, node: impl Into<ClusterNodeId>, topology: NodeTopology) {
self.observed_gossip.insert(node.into(), topology);
}
pub fn commit_topology(
&mut self,
node: impl Into<ClusterNodeId>,
topology: NodeTopology,
epoch: ClusterEpoch,
) {
if epoch >= self.epoch {
self.epoch = epoch;
self.committed.insert(node.into(), topology);
}
}
pub fn epoch(&self) -> ClusterEpoch {
self.epoch
}
pub fn topology(&self, node: &ClusterNodeId) -> Option<&NodeTopology> {
self.committed.get(node)
}
pub fn committed_map(&self) -> BTreeMap<ClusterNodeId, NodeTopology> {
self.committed.clone()
}
pub fn gossip_map(&self) -> BTreeMap<ClusterNodeId, NodeTopology> {
self.observed_gossip.clone()
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ZoneAwareReplicaSet {
pub replicas: Replicas,
pub topology: BTreeMap<ClusterNodeId, NodeTopology>,
pub placement_zone_underspread: bool,
}
impl ZoneAwareReplicaSet {
pub fn all_nodes(&self) -> Vec<ClusterNodeId> {
self.replicas.all_nodes()
}
pub fn zone_count(&self) -> usize {
self.topology
.values()
.map(|topology| (topology.region.clone(), topology.zone.clone()))
.collect::<BTreeSet<_>>()
.len()
}
pub fn single_zone_loss_keeps_write_quorum(&self, write_quorum: usize) -> bool {
let write_quorum = write_quorum.max(1);
let zones = self
.topology
.values()
.map(|topology| (topology.region.clone(), topology.zone.clone()))
.collect::<BTreeSet<_>>();
if zones.len() <= 1 {
return self.replicas.copy_count() >= write_quorum;
}
zones.into_iter().all(|lost_zone| {
self.all_nodes()
.into_iter()
.filter(|node| {
self.topology
.get(node)
.map(|topology| {
(topology.region.clone(), topology.zone.clone()) != lost_zone
})
.unwrap_or(true)
})
.count()
>= write_quorum
})
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ZonePlacementReadiness {
pub placement_zone_underspread: bool,
pub single_zone_loss_keeps_write_quorum: bool,
pub zone_count: usize,
}
impl ZonePlacementReadiness {
pub fn is_ready(&self) -> bool {
!self.placement_zone_underspread && self.single_zone_loss_keeps_write_quorum
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ZoneAwareReplicationStrategy {
topology: BTreeMap<ClusterNodeId, NodeTopology>,
replication_factor: usize,
min_zones: usize,
}
impl ZoneAwareReplicationStrategy {
pub fn new(
topology: BTreeMap<ClusterNodeId, NodeTopology>,
replication_factor: usize,
min_zones: usize,
) -> Self {
Self {
topology,
replication_factor: replication_factor.max(1),
min_zones: min_zones.max(1),
}
}
pub fn replication_factor(&self) -> usize {
self.replication_factor
}
pub fn min_zones(&self) -> usize {
self.min_zones
}
pub fn zone_replicas_for_key(
&self,
key: &str,
members: &[ClusterMember],
) -> Option<ZoneAwareReplicaSet> {
let flat = RendezvousClusterOwnership.replicas_for_key(
key,
members,
members.len().max(self.replication_factor),
)?;
let ranked = flat.all_nodes();
let primary = ranked.first()?.clone();
let mut selected = vec![primary.clone()];
let mut used_zones = BTreeSet::new();
if let Some(topology) = self.topology_for_node(&primary) {
used_zones.insert((topology.region.clone(), topology.zone.clone()));
}
for node in ranked.iter().skip(1) {
if selected.len() >= self.replication_factor {
break;
}
let Some(topology) = self.topology_for_node(node) else {
continue;
};
let zone_key = (topology.region.clone(), topology.zone.clone());
if used_zones.insert(zone_key) {
selected.push(node.clone());
}
}
for node in ranked.iter().skip(1) {
if selected.len() >= self.replication_factor {
break;
}
if !selected.contains(node) {
selected.push(node.clone());
}
}
let mut selected_iter = selected.into_iter();
let primary = selected_iter.next()?;
let backups = selected_iter.collect::<Vec<_>>();
let replicas = Replicas::new(primary, backups);
let topology = replicas
.all_nodes()
.into_iter()
.map(|node| {
let topology = self
.topology_for_node(&node)
.unwrap_or_else(NodeTopology::default_single_zone);
(node, topology)
})
.collect::<BTreeMap<_, _>>();
let required_zones = self.min_zones.min(self.replication_factor);
Some(ZoneAwareReplicaSet {
placement_zone_underspread: topology
.values()
.map(|topology| (topology.region.clone(), topology.zone.clone()))
.collect::<BTreeSet<_>>()
.len()
< required_zones,
replicas,
topology,
})
}
pub fn zone_replicas_for_key_in_regions(
&self,
key: &str,
members: &[ClusterMember],
allowed_regions: &BTreeSet<RegionId>,
min_replicas_in_policy: usize,
) -> Option<ZoneAwareReplicaSet> {
let required = self.replication_factor.max(min_replicas_in_policy.max(1));
let filtered = members
.iter()
.filter(|member| {
self.topology_for_node(&member.node_id)
.map(|topology| allowed_regions.contains(&topology.region))
.unwrap_or(false)
})
.cloned()
.collect::<Vec<_>>();
if filtered.len() < required {
return None;
}
let mut strategy = self.clone();
strategy.replication_factor = required;
let replicas = strategy.zone_replicas_for_key(key, &filtered)?;
if replicas.replicas.copy_count() < required {
return None;
}
if replicas
.topology
.values()
.all(|topology| allowed_regions.contains(&topology.region))
{
Some(replicas)
} else {
None
}
}
pub fn readiness_for_key(
&self,
key: &str,
members: &[ClusterMember],
write_quorum: usize,
) -> Option<ZonePlacementReadiness> {
let replicas = self.zone_replicas_for_key(key, members)?;
Some(ZonePlacementReadiness {
placement_zone_underspread: replicas.placement_zone_underspread,
single_zone_loss_keeps_write_quorum: replicas
.single_zone_loss_keeps_write_quorum(write_quorum),
zone_count: replicas.zone_count(),
})
}
fn topology_for_node(&self, node: &ClusterNodeId) -> Option<NodeTopology> {
self.topology.get(node).cloned()
}
}
impl ClusterReplicationStrategy for ZoneAwareReplicationStrategy {
fn name(&self) -> &'static str {
"zone-aware"
}
fn replicas_for_key(
&self,
key: &str,
members: &[ClusterMember],
replication_factor: usize,
) -> Option<Replicas> {
let mut strategy = self.clone();
strategy.replication_factor = replication_factor.max(1);
strategy
.zone_replicas_for_key(key, members)
.map(|replicas| replicas.replicas)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MovePhase {
Prepare,
Backfill,
Commit,
Cleanup,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PartitionMove {
pub partition: PartitionId,
pub from: ClusterNodeId,
pub to: ClusterNodeId,
pub phase: MovePhase,
pub backfilled_bytes: u64,
pub total_bytes: u64,
}
impl PartitionMove {
pub fn new(
partition: PartitionId,
from: impl Into<ClusterNodeId>,
to: impl Into<ClusterNodeId>,
total_bytes: u64,
) -> Self {
Self {
partition,
from: from.into(),
to: to.into(),
phase: MovePhase::Prepare,
backfilled_bytes: 0,
total_bytes: total_bytes.max(1),
}
}
pub fn write_targets(&self) -> Vec<ClusterNodeId> {
match self.phase {
MovePhase::Prepare | MovePhase::Backfill => {
vec![self.from.clone(), self.to.clone()]
}
MovePhase::Commit | MovePhase::Cleanup => vec![self.to.clone()],
}
}
pub fn read_owner(&self) -> ClusterNodeId {
match self.phase {
MovePhase::Prepare | MovePhase::Backfill => self.from.clone(),
MovePhase::Commit | MovePhase::Cleanup => self.to.clone(),
}
}
pub fn record_backfill(&mut self, bytes: u64) {
self.backfilled_bytes = self
.backfilled_bytes
.saturating_add(bytes)
.min(self.total_bytes);
}
pub fn progress_ratio(&self) -> f32 {
(self.backfilled_bytes as f32 / self.total_bytes as f32).min(1.0)
}
pub fn advance(&mut self) {
self.phase = match self.phase {
MovePhase::Prepare => MovePhase::Backfill,
MovePhase::Backfill if self.backfilled_bytes >= self.total_bytes => MovePhase::Commit,
MovePhase::Backfill => MovePhase::Backfill,
MovePhase::Commit => MovePhase::Cleanup,
MovePhase::Cleanup => MovePhase::Cleanup,
};
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ReshardPlan {
pub epoch: ClusterEpoch,
pub moves: Vec<PartitionMove>,
pub max_concurrent: usize,
}
impl ReshardPlan {
pub fn new(epoch: ClusterEpoch, mut moves: Vec<PartitionMove>, max_concurrent: usize) -> Self {
moves.sort_by_key(|movement| movement.partition.value());
Self {
epoch,
moves,
max_concurrent: max_concurrent.max(1),
}
}
pub fn active_moves(&self) -> Vec<&PartitionMove> {
self.moves
.iter()
.filter(|movement| movement.phase != MovePhase::Cleanup)
.take(self.max_concurrent)
.collect()
}
pub fn write_targets_for_partition(
&self,
partition: PartitionId,
) -> Option<Vec<ClusterNodeId>> {
self.moves
.iter()
.find(|movement| movement.partition == partition)
.map(PartitionMove::write_targets)
}
pub fn record_backfill(&mut self, partition: PartitionId, bytes: u64) {
if let Some(movement) = self
.moves
.iter_mut()
.find(|movement| movement.partition == partition)
{
movement.record_backfill(bytes);
}
}
pub fn snapshot(&self) -> Self {
self.clone()
}
pub fn resume_from(snapshot: Self) -> Self {
snapshot
}
pub fn drain_node(
epoch: ClusterEpoch,
node: impl Into<ClusterNodeId>,
targets: impl IntoIterator<Item = (PartitionId, ClusterNodeId, u64)>,
max_concurrent: usize,
) -> Self {
let node = node.into();
let moves = targets
.into_iter()
.map(|(partition, target, bytes)| {
PartitionMove::new(partition, node.clone(), target, bytes)
})
.collect();
Self::new(epoch, moves, max_concurrent)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ReshardPlanError {
message: String,
}
impl ReshardPlanError {
pub fn new(message: impl Into<String>) -> Self {
Self {
message: message.into(),
}
}
}
impl fmt::Display for ReshardPlanError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.message)
}
}
impl std::error::Error for ReshardPlanError {}
pub fn validate_move_preserves_zone_quorum(
candidate: &ZoneAwareReplicaSet,
write_quorum: usize,
) -> Result<(), ReshardPlanError> {
if !candidate.placement_zone_underspread
&& candidate.single_zone_loss_keeps_write_quorum(write_quorum)
{
Ok(())
} else {
Err(ReshardPlanError::new(
"reshard move would violate zone-spread write quorum",
))
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ReplicaSelection {
NearestZone,
LowestLatency,
RoundRobin,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ReplicaObservation {
pub node: ClusterNodeId,
pub topology: NodeTopology,
pub healthy: bool,
pub ewma_rtt_ms: u64,
pub version: ValueVersion,
pub epoch: ClusterEpoch,
}
impl ReplicaObservation {
pub fn healthy(
node: impl Into<ClusterNodeId>,
topology: NodeTopology,
ewma_rtt_ms: u64,
version: ValueVersion,
epoch: ClusterEpoch,
) -> Self {
Self {
node: node.into(),
topology,
healthy: true,
ewma_rtt_ms,
version,
epoch,
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ReplicaScorer {
observations: BTreeMap<ClusterNodeId, ReplicaObservation>,
}
impl ReplicaScorer {
pub fn new() -> Self {
Self::default()
}
pub fn observe(&mut self, observation: ReplicaObservation) {
self.observations
.insert(observation.node.clone(), observation);
}
pub fn order(
&self,
replicas: &[ClusterNodeId],
local: &NodeTopology,
selection: ReplicaSelection,
) -> Vec<ClusterNodeId> {
let mut ranked = replicas.to_vec();
ranked.sort_by(|left, right| {
let left_obs = self.observations.get(left);
let right_obs = self.observations.get(right);
let left_key = replica_score_key(left_obs, local, left, selection);
let right_key = replica_score_key(right_obs, local, right, selection);
left_key.cmp(&right_key)
});
ranked
}
pub fn observations_for(&self, nodes: &[ClusterNodeId]) -> Vec<ReplicaObservation> {
nodes
.iter()
.filter_map(|node| self.observations.get(node).cloned())
.collect()
}
}
fn replica_score_key(
observation: Option<&ReplicaObservation>,
local: &NodeTopology,
node: &ClusterNodeId,
selection: ReplicaSelection,
) -> (u8, u8, u64, String) {
let healthy = observation.map(|obs| obs.healthy).unwrap_or(false);
let distance = observation
.map(|obs| {
if obs.topology.same_zone(local) {
0
} else if obs.topology.same_region(local) {
1
} else {
2
}
})
.unwrap_or(3);
let latency = observation.map(|obs| obs.ewma_rtt_ms).unwrap_or(u64::MAX);
match selection {
ReplicaSelection::NearestZone => (!healthy as u8, distance, latency, node.to_string()),
ReplicaSelection::LowestLatency => (!healthy as u8, 0, latency, node.to_string()),
ReplicaSelection::RoundRobin => (0, 0, 0, node.to_string()),
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct HedgePolicy {
pub percentile: u8,
pub max_extra: usize,
pub min_delay_ms: u64,
}
impl HedgePolicy {
pub fn new(percentile: u8, max_extra: usize, min_delay_ms: u64) -> Self {
Self {
percentile: percentile.clamp(1, 100),
max_extra,
min_delay_ms,
}
}
pub fn delay_ms(self, observed_rtts: &[u64]) -> u64 {
if observed_rtts.is_empty() {
return self.min_delay_ms;
}
let mut sorted = observed_rtts.to_vec();
sorted.sort_unstable();
let index = ((sorted.len() - 1) * self.percentile as usize) / 100;
sorted[index].max(self.min_delay_ms)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct HedgedReadPlan {
pub primary: Option<ClusterNodeId>,
pub hedges: Vec<ClusterNodeId>,
pub required_acks: usize,
}
pub fn plan_hedged_read(
ordered_replicas: &[ClusterNodeId],
required_acks: usize,
elapsed_ms: u64,
observed_rtts: &[u64],
policy: HedgePolicy,
) -> HedgedReadPlan {
let primary = ordered_replicas.first().cloned();
let hedge_delay = policy.delay_ms(observed_rtts);
let hedges = if elapsed_ms >= hedge_delay {
ordered_replicas
.iter()
.skip(1)
.take(policy.max_extra)
.cloned()
.collect()
} else {
Vec::new()
};
HedgedReadPlan {
primary,
hedges,
required_acks: required_acks.max(1),
}
}
pub fn hedge_winner(
responses: impl IntoIterator<Item = ReplicatedValueRecord>,
) -> Option<ReplicatedValueRecord> {
responses
.into_iter()
.max_by_key(|record| (record.version, record.epoch))
}
#[derive(Debug, Clone)]
pub struct TieredValueStore<S> {
cold: S,
hot: BTreeMap<String, ReplicatedValueRecord>,
order: VecDeque<String>,
max_hot_bytes: u64,
promotions_total: u64,
demotions_total: u64,
}
impl<S> TieredValueStore<S>
where
S: ReplicatedValueStore,
{
pub fn new(cold: S, max_hot_bytes: u64) -> Self {
Self {
cold,
hot: BTreeMap::new(),
order: VecDeque::new(),
max_hot_bytes: max_hot_bytes.max(1),
promotions_total: 0,
demotions_total: 0,
}
}
pub fn cold(&self) -> &S {
&self.cold
}
pub fn cold_mut(&mut self) -> &mut S {
&mut self.cold
}
pub fn hot_contains(&self, key: &str) -> bool {
self.hot.contains_key(key)
}
pub fn hot_bytes(&self) -> u64 {
self.hot
.values()
.map(ReplicatedValueRecord::approx_bytes)
.sum()
}
pub fn hot_ratio(&self) -> f32 {
let hot = self.hot.len() as f32;
if hot == 0.0 {
return 0.0;
}
hot / (hot + 1.0)
}
pub fn promotions_total(&self) -> u64 {
self.promotions_total
}
pub fn demotions_total(&self) -> u64 {
self.demotions_total
}
pub fn get_promote(
&mut self,
key: &str,
) -> Result<Option<ReplicatedValueRecord>, ValueStoreError> {
let record = self.get(key)?;
if let Some(record) = record.clone() {
self.promote_hot(key.to_owned(), record)?;
}
Ok(record)
}
fn promote_hot(
&mut self,
key: String,
record: ReplicatedValueRecord,
) -> Result<(), ValueStoreError> {
if record.approx_bytes() > self.max_hot_bytes {
return Ok(());
}
let merged = self
.hot
.remove(&key)
.map(|current| current.merge(record.clone()))
.unwrap_or(record);
self.hot.insert(key.clone(), merged);
self.order.retain(|existing| existing != &key);
self.order.push_back(key);
self.promotions_total = self.promotions_total.saturating_add(1);
self.enforce_hot_budget()
}
fn enforce_hot_budget(&mut self) -> Result<(), ValueStoreError> {
while self.hot_bytes() > self.max_hot_bytes {
let Some(key) = self.order.pop_front() else {
break;
};
if let Some(record) = self.hot.remove(&key) {
self.cold.upsert(key, record)?;
self.demotions_total = self.demotions_total.saturating_add(1);
}
}
Ok(())
}
}
impl<S> ReplicatedValueStore for TieredValueStore<S>
where
S: ReplicatedValueStore,
{
fn upsert(
&mut self,
key: impl Into<String>,
rec: ReplicatedValueRecord,
) -> Result<(), ValueStoreError> {
let key = key.into();
self.cold.upsert(key.clone(), rec.clone())?;
self.promote_hot(key, rec)
}
fn get(&self, key: &str) -> Result<Option<ReplicatedValueRecord>, ValueStoreError> {
let hot = self.hot.get(key).cloned();
let cold = self.cold.get(key)?;
Ok(match (hot, cold) {
(Some(hot), Some(cold)) => Some(hot.merge(cold)),
(Some(hot), None) => Some(hot),
(None, Some(cold)) => Some(cold),
(None, None) => None,
})
}
fn tombstone(
&mut self,
key: impl Into<String>,
partition: PartitionId,
version: ValueVersion,
epoch: ClusterEpoch,
) -> Result<(), ValueStoreError> {
self.upsert(
key,
ReplicatedValueRecord::tombstone(partition, version, epoch, None),
)
}
fn scan_owned(
&self,
map: &EffectiveReplicationMap,
) -> Result<Vec<(String, ReplicatedValueRecord)>, ValueStoreError> {
let mut merged = BTreeMap::new();
for (key, record) in self.cold.scan_owned(map)? {
merged.insert(key, record);
}
for (key, record) in &self.hot {
let record = merged
.remove(key)
.map(|cold| cold.merge(record.clone()))
.unwrap_or_else(|| record.clone());
merged.insert(key.clone(), record);
}
Ok(merged.into_iter().collect())
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct InvalidateBatch {
pub partition: PartitionId,
pub keys: Vec<String>,
pub version: ValueVersion,
pub epoch: ClusterEpoch,
}
impl InvalidateBatch {
pub fn try_new(
keys: impl IntoIterator<Item = impl Into<String>>,
partition_count: u32,
version: ValueVersion,
epoch: ClusterEpoch,
) -> Result<Self, AtomicInvalidationError> {
let keys = keys.into_iter().map(Into::into).collect::<Vec<_>>();
if keys.is_empty() {
return Err(AtomicInvalidationError::new("invalidate batch is empty"));
}
let partition = partition_for_key(&keys[0], partition_count);
if keys
.iter()
.any(|key| partition_for_key(key, partition_count) != partition)
{
return Err(AtomicInvalidationError::new(
"cross-partition invalidation batch rejected; use InvalidationSaga",
));
}
Ok(Self {
partition,
keys,
version,
epoch,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AtomicInvalidationError {
message: String,
}
impl AtomicInvalidationError {
pub fn new(message: impl Into<String>) -> Self {
Self {
message: message.into(),
}
}
}
impl fmt::Display for AtomicInvalidationError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.message)
}
}
impl std::error::Error for AtomicInvalidationError {}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct BatchInvalidationState {
applied: BTreeMap<String, WriteWatermark>,
}
impl BatchInvalidationState {
pub fn apply_batch(&mut self, batch: &InvalidateBatch) {
let watermark = WriteWatermark::new(batch.partition, batch.version, batch.epoch);
for key in &batch.keys {
self.applied.insert(key.clone(), watermark);
}
}
pub fn apply_single(&mut self, key: impl Into<String>, watermark: WriteWatermark) {
let key = key.into();
let replace = self
.applied
.get(&key)
.map(|current| (watermark.version, watermark.epoch) >= (current.version, current.epoch))
.unwrap_or(true);
if replace {
self.applied.insert(key, watermark);
}
}
pub fn watermark(&self, key: &str) -> Option<WriteWatermark> {
self.applied.get(key).copied()
}
pub fn batch_is_all_or_nothing(&self, batch: &InvalidateBatch) -> bool {
let expected = WriteWatermark::new(batch.partition, batch.version, batch.epoch);
batch
.keys
.iter()
.all(|key| self.applied.get(key) == Some(&expected))
}
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub struct InvalidationTarget {
pub partition: PartitionId,
pub key: String,
}
impl InvalidationTarget {
pub fn new(partition: PartitionId, key: impl Into<String>) -> Self {
Self {
partition,
key: key.into(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct InvalidationSaga {
pub unit_id: String,
pub targets: Vec<InvalidationTarget>,
applied: BTreeSet<String>,
}
impl InvalidationSaga {
pub fn new(unit_id: impl Into<String>, targets: Vec<InvalidationTarget>) -> Self {
Self {
unit_id: unit_id.into(),
targets,
applied: BTreeSet::new(),
}
}
pub fn dispatch_target(&mut self, target: &InvalidationTarget) -> bool {
self.applied.insert(self.idempotency_key(target))
}
pub fn pending(&self) -> usize {
self.targets
.iter()
.filter(|target| !self.applied.contains(&self.idempotency_key(target)))
.count()
}
pub fn is_complete(&self) -> bool {
self.pending() == 0
}
pub fn idempotency_key(&self, target: &InvalidationTarget) -> String {
format!(
"{}:{}:{}",
self.unit_id,
target.partition.value(),
target.key
)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RepairMode {
Advisory,
Active,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct AutoRepairPolicy {
pub mode: RepairMode,
pub debt_threshold: u64,
pub lag_threshold: u64,
pub max_concurrent_repairs: usize,
}
impl AutoRepairPolicy {
pub fn new(
mode: RepairMode,
debt_threshold: u64,
lag_threshold: u64,
max_concurrent_repairs: usize,
) -> Self {
Self {
mode,
debt_threshold,
lag_threshold,
max_concurrent_repairs: max_concurrent_repairs.max(1),
}
}
pub fn evaluate(&self, debt: u64, lag: u64) -> AutoRepairDecision {
let should_repair = debt > self.debt_threshold || lag > self.lag_threshold;
let recommended = if should_repair {
vec![RepairAction::AntiEntropy]
} else {
Vec::new()
};
let scheduled = if should_repair && self.mode == RepairMode::Active {
recommended
.iter()
.copied()
.take(self.max_concurrent_repairs)
.collect()
} else {
Vec::new()
};
AutoRepairDecision {
mode: self.mode,
recommended,
scheduled,
capped_at: self.max_concurrent_repairs,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RepairAction {
AntiEntropy,
ReReplicate,
MovePartition,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AutoRepairDecision {
pub mode: RepairMode,
pub recommended: Vec<RepairAction>,
pub scheduled: Vec<RepairAction>,
pub capped_at: usize,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ControlPlaneSnapshot {
pub format_version: u32,
pub epoch: ClusterEpoch,
pub topology: BTreeMap<ClusterNodeId, NodeTopology>,
pub ownership: BTreeMap<PartitionId, Replicas>,
pub tombstone_versions: BTreeMap<String, ValueVersion>,
}
impl ControlPlaneSnapshot {
pub fn new(epoch: ClusterEpoch) -> Self {
Self {
format_version: CONTROL_PLANE_SNAPSHOT_FORMAT_VERSION,
epoch,
..Self::default()
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotError {
message: String,
}
impl SnapshotError {
pub fn new(message: impl Into<String>) -> Self {
Self {
message: message.into(),
}
}
}
impl fmt::Display for SnapshotError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.message)
}
}
impl std::error::Error for SnapshotError {}
pub trait SnapshotSink: Send + Sync {
fn put(&mut self, snapshot: ControlPlaneSnapshot) -> Result<(), SnapshotError>;
fn latest(&self) -> Result<Option<ControlPlaneSnapshot>, SnapshotError>;
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct InMemorySnapshotSink {
snapshots: Vec<ControlPlaneSnapshot>,
}
impl InMemorySnapshotSink {
pub fn new() -> Self {
Self::default()
}
pub fn len(&self) -> usize {
self.snapshots.len()
}
pub fn is_empty(&self) -> bool {
self.snapshots.is_empty()
}
}
impl SnapshotSink for InMemorySnapshotSink {
fn put(&mut self, snapshot: ControlPlaneSnapshot) -> Result<(), SnapshotError> {
if snapshot.format_version > CONTROL_PLANE_SNAPSHOT_FORMAT_VERSION {
return Err(SnapshotError::new(
"control-plane snapshot format is newer than this binary",
));
}
self.snapshots.push(snapshot);
Ok(())
}
fn latest(&self) -> Result<Option<ControlPlaneSnapshot>, SnapshotError> {
Ok(self.snapshots.last().cloned())
}
}
pub fn restore_topology_from_snapshot(
snapshot: &ControlPlaneSnapshot,
) -> Result<TopologyAuthority, SnapshotError> {
if snapshot.format_version > CONTROL_PLANE_SNAPSHOT_FORMAT_VERSION {
return Err(SnapshotError::new(
"control-plane snapshot format is newer than this binary",
));
}
let mut authority = TopologyAuthority::new();
for (node, topology) in &snapshot.topology {
authority.commit_topology(node.clone(), topology.clone(), snapshot.epoch);
}
Ok(authority)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub struct CompatVersion {
pub major: u16,
pub minor: u16,
pub patch: u16,
}
impl CompatVersion {
pub const fn new(major: u16, minor: u16, patch: u16) -> Self {
Self {
major,
minor,
patch,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct UpgradeStep {
pub from: CompatVersion,
pub to: CompatVersion,
pub raft_log_format: u32,
pub value_record_format: u32,
pub wire_frame_version: u16,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct UpgradeGuard {
pub min_minor: u16,
pub max_minor: u16,
pub raft_log_format: u32,
pub value_record_format: u32,
pub wire_frame_version: u16,
}
impl UpgradeGuard {
pub fn current() -> Self {
Self {
min_minor: 42,
max_minor: 43,
raft_log_format: 1,
value_record_format: REPLICATED_VALUE_RECORD_FORMAT_VERSION,
wire_frame_version: CACHE_INVALIDATION_FRAME_VERSION,
}
}
pub fn check(&self, step: UpgradeStep) -> Result<(), UpgradeGuardError> {
if step.from.major != 0 || step.to.major != 0 {
return Err(UpgradeGuardError::new(
"only 0.x HydraCache versions are supported",
));
}
if step.from.minor < self.min_minor || step.to.minor > self.max_minor {
return Err(UpgradeGuardError::new(
"upgrade step is outside the registered compatibility window",
));
}
if step.to.minor.saturating_sub(step.from.minor) > 1 {
return Err(UpgradeGuardError::new(
"rolling upgrade step skips a supported minor version",
));
}
if step.raft_log_format != self.raft_log_format
|| step.value_record_format != self.value_record_format
|| step.wire_frame_version != self.wire_frame_version
{
return Err(UpgradeGuardError::new(
"upgrade step uses an incompatible persisted or wire format",
));
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpgradeGuardError {
message: String,
}
impl UpgradeGuardError {
pub fn new(message: impl Into<String>) -> Self {
Self {
message: message.into(),
}
}
}
impl fmt::Display for UpgradeGuardError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.message)
}
}
impl std::error::Error for UpgradeGuardError {}