use std::collections::BTreeMap;
use std::fmt;
use serde::{Deserialize, Serialize};
use crate::cluster::{ClusterEpoch, ClusterNodeId, PartitionId};
use crate::grid::{EffectiveReplicationMap, ReplicatedSlot, ReplicationConfig};
pub type ValueVersion = u64;
pub type SealedBytes = Vec<u8>;
pub const REPLICATED_VALUE_RECORD_FORMAT_VERSION: u32 = 1;
pub const REPLICATED_VALUE_RECORD_CHECKSUM_FORMAT_VERSION: u32 = 1;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ReplicatedValueRecord {
pub partition: PartitionId,
pub version: ValueVersion,
pub epoch: ClusterEpoch,
pub state: ReplicatedSlot<SealedBytes>,
}
impl ReplicatedValueRecord {
pub fn value(
partition: PartitionId,
version: ValueVersion,
epoch: ClusterEpoch,
sealed: impl Into<SealedBytes>,
) -> Self {
Self {
partition,
version,
epoch,
state: ReplicatedSlot::Value {
value: sealed.into(),
version,
},
}
}
pub fn tombstone(
partition: PartitionId,
version: ValueVersion,
epoch: ClusterEpoch,
gc_eligible_after: Option<ClusterEpoch>,
) -> Self {
Self {
partition,
version,
epoch,
state: ReplicatedSlot::Tombstone {
version,
gc_eligible_after,
},
}
}
pub fn is_tombstone(&self) -> bool {
self.state.is_tombstone()
}
pub fn approx_bytes(&self) -> u64 {
match &self.state {
ReplicatedSlot::Value { value, .. } => value.len().max(1) as u64,
ReplicatedSlot::Tombstone { .. } => 1,
}
}
pub fn artifact_checksum(&self) -> u64 {
let mut checksum = ArtifactChecksum::new();
checksum.u32(REPLICATED_VALUE_RECORD_FORMAT_VERSION);
checksum.u32(self.partition.value());
checksum.u64(self.version);
checksum.u64(self.epoch.value());
match &self.state {
ReplicatedSlot::Value { value, version } => {
checksum.u8(1);
checksum.u64(*version);
checksum.bytes(value);
}
ReplicatedSlot::Tombstone {
version,
gc_eligible_after,
} => {
checksum.u8(2);
checksum.u64(*version);
checksum.u64(
gc_eligible_after
.map(ClusterEpoch::value)
.unwrap_or(u64::MAX),
);
}
}
checksum.finish()
}
pub fn merge(self, other: Self) -> Self {
match (self.version, self.epoch).cmp(&(other.version, other.epoch)) {
std::cmp::Ordering::Greater => self,
std::cmp::Ordering::Less => other,
std::cmp::Ordering::Equal if self.is_tombstone() => self,
std::cmp::Ordering::Equal => other,
}
}
}
#[derive(Debug, Clone, Copy)]
struct ArtifactChecksum(u64);
impl ArtifactChecksum {
const OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
const PRIME: u64 = 0x0000_0100_0000_01b3;
fn new() -> Self {
Self(Self::OFFSET)
}
fn u8(&mut self, value: u8) {
self.0 ^= u64::from(value);
self.0 = self.0.wrapping_mul(Self::PRIME);
}
fn u32(&mut self, value: u32) {
self.bytes(&value.to_le_bytes());
}
fn u64(&mut self, value: u64) {
self.bytes(&value.to_le_bytes());
}
fn bytes(&mut self, bytes: &[u8]) {
for byte in bytes {
self.u8(*byte);
}
}
fn finish(self) -> u64 {
self.0
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ChecksummedReplicatedValueRecord {
pub checksum_format: u32,
pub record: ReplicatedValueRecord,
pub checksum: u64,
}
impl ChecksummedReplicatedValueRecord {
pub fn seal(record: ReplicatedValueRecord) -> Self {
let checksum = record.artifact_checksum();
Self {
checksum_format: REPLICATED_VALUE_RECORD_CHECKSUM_FORMAT_VERSION,
record,
checksum,
}
}
pub fn from_parts(checksum_format: u32, record: ReplicatedValueRecord, checksum: u64) -> Self {
Self {
checksum_format,
record,
checksum,
}
}
pub fn verify(&self) -> Result<(), ScrubErrorKind> {
if self.checksum_format != REPLICATED_VALUE_RECORD_CHECKSUM_FORMAT_VERSION {
return Err(ScrubErrorKind::UnsupportedChecksumFormat {
found: self.checksum_format,
});
}
let actual = self.record.artifact_checksum();
if actual != self.checksum {
return Err(ScrubErrorKind::ChecksumMismatch {
expected: self.checksum,
actual,
});
}
Ok(())
}
pub fn verified_record(
&self,
key: impl Into<String>,
) -> Result<&ReplicatedValueRecord, ScrubError> {
self.verify().map_err(|kind| ScrubError {
key: key.into(),
kind,
})?;
Ok(&self.record)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ScrubErrorKind {
UnsupportedChecksumFormat { found: u32 },
ChecksumMismatch { expected: u64, actual: u64 },
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ScrubError {
pub key: String,
pub kind: ScrubErrorKind,
}
impl fmt::Display for ScrubError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
formatter,
"scrub error for key '{}': {:?}",
self.key, self.kind
)
}
}
impl std::error::Error for ScrubError {}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ScrubReport {
pub checked: usize,
pub corrupt: usize,
pub repaired: usize,
pub unrepairable: usize,
pub errors: Vec<ScrubError>,
}
impl ScrubReport {
pub fn is_ok(&self) -> bool {
self.unrepairable == 0
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Scrubber {
max_records_per_run: usize,
}
impl Scrubber {
pub fn new(max_records_per_run: usize) -> Self {
Self {
max_records_per_run: max_records_per_run.max(1),
}
}
pub fn scrub_replicated_values(
&self,
primary: &mut BTreeMap<String, ChecksummedReplicatedValueRecord>,
peers: &[BTreeMap<String, ChecksummedReplicatedValueRecord>],
) -> ScrubReport {
let keys = primary
.keys()
.take(self.max_records_per_run)
.cloned()
.collect::<Vec<_>>();
let mut report = ScrubReport::default();
for key in keys {
report.checked = report.checked.saturating_add(1);
let Some(record) = primary.get(&key) else {
continue;
};
if record.verify().is_ok() {
continue;
}
report.corrupt = report.corrupt.saturating_add(1);
if let Some(repair) = peers.iter().find_map(|peer| {
peer.get(&key)
.filter(|candidate| candidate.verify().is_ok())
.cloned()
}) {
primary.insert(key, repair);
report.repaired = report.repaired.saturating_add(1);
} else {
report.unrepairable = report.unrepairable.saturating_add(1);
let kind = record.verify().expect_err("record is known corrupt");
report.errors.push(ScrubError { key, kind });
}
}
report
}
pub fn verified_get(
records: &BTreeMap<String, ChecksummedReplicatedValueRecord>,
key: &str,
) -> Result<Option<ReplicatedValueRecord>, ScrubError> {
records
.get(key)
.map(|record| record.verified_record(key.to_owned()).cloned())
.transpose()
}
}
impl Default for Scrubber {
fn default() -> Self {
Self::new(128)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ValueStoreError {
message: String,
}
impl ValueStoreError {
pub fn new(message: impl Into<String>) -> Self {
Self {
message: message.into(),
}
}
}
impl fmt::Display for ValueStoreError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.message)
}
}
impl std::error::Error for ValueStoreError {}
pub trait ReplicatedValueStore: Send + Sync {
fn upsert(
&mut self,
key: impl Into<String>,
rec: ReplicatedValueRecord,
) -> Result<(), ValueStoreError>;
fn get(&self, key: &str) -> Result<Option<ReplicatedValueRecord>, ValueStoreError>;
fn tombstone(
&mut self,
key: impl Into<String>,
partition: PartitionId,
version: ValueVersion,
epoch: ClusterEpoch,
) -> Result<(), ValueStoreError>;
fn scan_owned(
&self,
map: &EffectiveReplicationMap,
) -> Result<Vec<(String, ReplicatedValueRecord)>, ValueStoreError>;
fn scan_all(&self) -> Result<Vec<(String, ReplicatedValueRecord)>, ValueStoreError>;
fn remove(&mut self, key: &str) -> Result<(), ValueStoreError>;
fn compact(&mut self) -> Result<u64, ValueStoreError> {
Ok(0)
}
fn total_bytes(&self) -> Result<u64, ValueStoreError>;
fn rejected_total(&self) -> u64;
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct InMemoryReplicatedValueStore {
records: BTreeMap<String, ReplicatedValueRecord>,
max_total_bytes: u64,
rejected_total: u64,
}
impl InMemoryReplicatedValueStore {
pub fn with_budget(max_total_bytes: u64) -> Self {
Self {
records: BTreeMap::new(),
max_total_bytes: max_total_bytes.max(1),
rejected_total: 0,
}
}
pub fn total_bytes(&self) -> u64 {
self.records
.values()
.map(ReplicatedValueRecord::approx_bytes)
.sum()
}
pub fn rejected_total(&self) -> u64 {
self.rejected_total
}
pub fn snapshot(&self) -> BTreeMap<String, ReplicatedValueRecord> {
self.records.clone()
}
pub fn reopen_from_snapshot(
max_total_bytes: u64,
records: BTreeMap<String, ReplicatedValueRecord>,
) -> Self {
Self {
records,
max_total_bytes: max_total_bytes.max(1),
rejected_total: 0,
}
}
fn would_fit(&self, key: &str, rec: &ReplicatedValueRecord) -> bool {
let existing = self
.records
.get(key)
.map(ReplicatedValueRecord::approx_bytes)
.unwrap_or_default();
self.total_bytes()
.saturating_sub(existing)
.saturating_add(rec.approx_bytes())
<= self.max_total_bytes
}
}
impl Default for InMemoryReplicatedValueStore {
fn default() -> Self {
Self::with_budget(u64::MAX)
}
}
impl ReplicatedValueStore for InMemoryReplicatedValueStore {
fn upsert(
&mut self,
key: impl Into<String>,
rec: ReplicatedValueRecord,
) -> Result<(), ValueStoreError> {
let key = key.into();
if !self.would_fit(&key, &rec) {
self.rejected_total = self.rejected_total.saturating_add(1);
return Err(ValueStoreError::new(
"replicated value store total byte budget exceeded",
));
}
let merged = self
.records
.remove(&key)
.map(|current| current.merge(rec.clone()))
.unwrap_or(rec);
self.records.insert(key, merged);
Ok(())
}
fn get(&self, key: &str) -> Result<Option<ReplicatedValueRecord>, ValueStoreError> {
Ok(self.records.get(key).cloned())
}
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> {
if map.reading.is_empty() {
return Ok(Vec::new());
}
self.scan_all()
}
fn scan_all(&self) -> Result<Vec<(String, ReplicatedValueRecord)>, ValueStoreError> {
Ok(self
.records
.iter()
.map(|(key, record)| (key.clone(), record.clone()))
.collect())
}
fn remove(&mut self, key: &str) -> Result<(), ValueStoreError> {
self.records.remove(key);
Ok(())
}
fn total_bytes(&self) -> Result<u64, ValueStoreError> {
Ok(InMemoryReplicatedValueStore::total_bytes(self))
}
fn rejected_total(&self) -> u64 {
InMemoryReplicatedValueStore::rejected_total(self)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct AdaptiveWindow {
in_flight: usize,
max_in_flight: usize,
floor: usize,
ceil: usize,
}
impl AdaptiveWindow {
pub fn new(floor: usize, initial: usize, ceil: usize) -> Self {
let floor = floor.max(1);
let ceil = ceil.max(floor);
let max_in_flight = initial.max(floor).min(ceil);
Self {
in_flight: 0,
max_in_flight,
floor,
ceil,
}
}
pub fn admit(&self) -> bool {
self.in_flight < self.max_in_flight
}
pub fn try_acquire(&mut self) -> bool {
if !self.admit() {
return false;
}
self.in_flight = self.in_flight.saturating_add(1);
true
}
pub fn on_ack(&mut self, rtt_ok: bool) {
self.in_flight = self.in_flight.saturating_sub(1);
if rtt_ok {
self.max_in_flight = self.max_in_flight.saturating_add(1).min(self.ceil);
} else {
self.max_in_flight = (self.max_in_flight / 2).max(self.floor);
}
}
pub fn in_flight(&self) -> usize {
self.in_flight
}
pub fn max_in_flight(&self) -> usize {
self.max_in_flight
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum QuorumPosture {
Strong,
DegradedSessionRyow,
}
impl ReplicationConfig {
pub fn is_strong_ryow(&self) -> bool {
self.read_quorum.saturating_add(self.write_quorum) > self.replication_factor
}
pub fn quorum_posture(&self) -> QuorumPosture {
if self.is_strong_ryow() {
QuorumPosture::Strong
} else {
QuorumPosture::DegradedSessionRyow
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct WriteWatermark {
pub partition: PartitionId,
pub version: ValueVersion,
pub epoch: ClusterEpoch,
}
impl WriteWatermark {
pub fn new(partition: PartitionId, version: ValueVersion, epoch: ClusterEpoch) -> Self {
Self {
partition,
version,
epoch,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ReadConsistency {
QuorumReadYourWrites,
Eventual,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct QuorumReadDecision {
pub record: Option<ReplicatedValueRecord>,
pub requires_primary_fallback: bool,
}
pub fn quorum_read_your_writes(
watermark: WriteWatermark,
replicas: impl IntoIterator<Item = ReplicatedValueRecord>,
read_quorum: usize,
) -> QuorumReadDecision {
let mut candidates = replicas
.into_iter()
.filter(|record| record.partition == watermark.partition)
.filter(|record| (record.version, record.epoch) >= (watermark.version, watermark.epoch))
.collect::<Vec<_>>();
candidates.sort_by_key(|record| (record.version, record.epoch));
if candidates.len() < read_quorum.max(1) {
return QuorumReadDecision {
record: None,
requires_primary_fallback: true,
};
}
QuorumReadDecision {
record: candidates.pop(),
requires_primary_fallback: false,
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct SplitBrainReport {
pub winner_epoch: ClusterEpoch,
pub loser_epoch: ClusterEpoch,
pub merged_entries: u64,
pub discarded_entries: u64,
pub unresolved_conflicts: u64,
}
pub trait MergePolicy: Send + Sync {
fn merge(
&self,
winner: Option<&ReplicatedValueRecord>,
loser: &ReplicatedValueRecord,
) -> Option<ReplicatedValueRecord>;
}
#[derive(Debug, Clone, Copy, Default)]
pub struct HigherVersionWins;
impl MergePolicy for HigherVersionWins {
fn merge(
&self,
winner: Option<&ReplicatedValueRecord>,
loser: &ReplicatedValueRecord,
) -> Option<ReplicatedValueRecord> {
Some(
winner
.cloned()
.map(|winner| winner.merge(loser.clone()))
.unwrap_or_else(|| loser.clone()),
)
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct PutIfAbsent;
impl MergePolicy for PutIfAbsent {
fn merge(
&self,
winner: Option<&ReplicatedValueRecord>,
loser: &ReplicatedValueRecord,
) -> Option<ReplicatedValueRecord> {
winner.cloned().or_else(|| Some(loser.clone()))
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct DiscardLoser;
impl MergePolicy for DiscardLoser {
fn merge(
&self,
_winner: Option<&ReplicatedValueRecord>,
_loser: &ReplicatedValueRecord,
) -> Option<ReplicatedValueRecord> {
None
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ClusterMergeOutcome {
pub records: BTreeMap<String, ReplicatedValueRecord>,
pub report: SplitBrainReport,
}
pub fn merge_split_brain_records(
winner_epoch: ClusterEpoch,
loser_epoch: ClusterEpoch,
mut winner: BTreeMap<String, ReplicatedValueRecord>,
loser: BTreeMap<String, ReplicatedValueRecord>,
policy: &dyn MergePolicy,
) -> ClusterMergeOutcome {
let mut report = SplitBrainReport {
winner_epoch,
loser_epoch,
..SplitBrainReport::default()
};
for (key, loser_record) in loser {
match policy.merge(winner.get(&key), &loser_record) {
Some(record) => {
if winner.get(&key) != Some(&record) {
report.merged_entries = report.merged_entries.saturating_add(1);
}
winner.insert(key, record);
}
None => {
report.discarded_entries = report.discarded_entries.saturating_add(1);
}
}
}
ClusterMergeOutcome {
records: winner,
report,
}
}
pub fn split_brain_winner(left: ClusterEpoch, right: ClusterEpoch) -> (ClusterEpoch, ClusterEpoch) {
if left >= right {
(left, right)
} else {
(right, left)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct PromotionFreezeWindow {
pub observed_ms: u64,
pub bound_ms: u64,
}
impl PromotionFreezeWindow {
pub fn is_bounded(self) -> bool {
self.observed_ms <= self.bound_ms
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct LiveSplitBrainResolution {
pub winner_epoch: ClusterEpoch,
pub loser_epoch: ClusterEpoch,
pub outcome: ClusterMergeOutcome,
}
pub fn resolve_live_split_brain(
left_epoch: ClusterEpoch,
left_records: BTreeMap<String, ReplicatedValueRecord>,
right_epoch: ClusterEpoch,
right_records: BTreeMap<String, ReplicatedValueRecord>,
policy: &dyn MergePolicy,
) -> LiveSplitBrainResolution {
let left_wins = left_epoch >= right_epoch;
let (winner_epoch, loser_epoch) = split_brain_winner(left_epoch, right_epoch);
let (winner_records, loser_records) = if left_wins {
(left_records, right_records)
} else {
(right_records, left_records)
};
let outcome = merge_split_brain_records(
winner_epoch,
loser_epoch,
winner_records,
loser_records,
policy,
);
LiveSplitBrainResolution {
winner_epoch,
loser_epoch,
outcome,
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct LiveReadYourWritesDecision {
pub posture: QuorumPosture,
pub decision: QuorumReadDecision,
}
pub fn live_read_your_writes(
config: ReplicationConfig,
watermark: WriteWatermark,
replicas: impl IntoIterator<Item = ReplicatedValueRecord>,
) -> LiveReadYourWritesDecision {
LiveReadYourWritesDecision {
posture: config.quorum_posture(),
decision: quorum_read_your_writes(watermark, replicas, config.read_quorum),
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct LiveReplicationSend {
pub admitted: bool,
pub max_in_flight: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct LiveReplicationPeer {
node: ClusterNodeId,
window: AdaptiveWindow,
}
impl LiveReplicationPeer {
pub fn new(node: impl Into<ClusterNodeId>, window: AdaptiveWindow) -> Self {
Self {
node: node.into(),
window,
}
}
pub fn node(&self) -> &ClusterNodeId {
&self.node
}
pub fn window(&self) -> AdaptiveWindow {
self.window
}
pub fn send_record<S>(
&mut self,
backup: &mut S,
key: impl Into<String>,
record: ReplicatedValueRecord,
rtt_ok: bool,
) -> Result<LiveReplicationSend, ValueStoreError>
where
S: ReplicatedValueStore,
{
if !self.window.try_acquire() {
return Ok(LiveReplicationSend {
admitted: false,
max_in_flight: self.window.max_in_flight(),
});
}
backup.upsert(key, record)?;
self.window.on_ack(rtt_ok);
Ok(LiveReplicationSend {
admitted: true,
max_in_flight: self.window.max_in_flight(),
})
}
}
pub fn anti_entropy_repair<S>(
backup: &mut S,
records: impl IntoIterator<Item = (String, ReplicatedValueRecord)>,
) -> Result<u64, ValueStoreError>
where
S: ReplicatedValueStore,
{
let mut repaired = 0_u64;
for (key, record) in records {
backup.upsert(key, record)?;
repaired = repaired.saturating_add(1);
}
Ok(repaired)
}