use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::fmt;
use std::time::Duration;
use serde::{Deserialize, Serialize};
use crate::cluster::{ClusterEpoch, ClusterNodeId, PartitionId};
use crate::grid::hardening::{ReplicatedValueRecord, SealedBytes, ValueVersion};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Hint {
pub target: ClusterNodeId,
pub key: String,
pub partition: PartitionId,
pub version: ValueVersion,
pub epoch: ClusterEpoch,
pub sealed: SealedBytes,
pub created_at_millis: u64,
}
impl Hint {
pub fn new(
target: impl Into<ClusterNodeId>,
key: impl Into<String>,
partition: PartitionId,
version: ValueVersion,
epoch: ClusterEpoch,
sealed: impl Into<SealedBytes>,
created_at_millis: u64,
) -> Self {
Self {
target: target.into(),
key: key.into(),
partition,
version,
epoch,
sealed: sealed.into(),
created_at_millis,
}
}
pub fn approx_bytes(&self) -> u64 {
self.key.len().saturating_add(self.sealed.len()).max(1) as u64
}
fn is_expired(&self, now_millis: u64, budget: HintBudget) -> bool {
let max_age_millis = budget.max_age.as_millis().min(u128::from(u64::MAX)) as u64;
now_millis.saturating_sub(self.created_at_millis) > max_age_millis
}
fn to_record(&self) -> ReplicatedValueRecord {
ReplicatedValueRecord::value(
self.partition,
self.version,
self.epoch,
self.sealed.clone(),
)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct HintBudget {
pub max_hints: usize,
pub max_bytes: u64,
pub max_age: Duration,
}
impl HintBudget {
pub fn new(max_hints: usize, max_bytes: u64, max_age: Duration) -> Self {
Self {
max_hints: max_hints.max(1),
max_bytes: max_bytes.max(1),
max_age: if max_age.is_zero() {
Duration::from_millis(1)
} else {
max_age
},
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum HintOutcome {
Stored,
Replayed,
DroppedOverBudget,
DroppedExpired,
RequiredReplicaMiss,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum HintReplayDecision {
Replayed { record: ReplicatedValueRecord },
SuppressedByNewerValue { current: ReplicatedValueRecord },
SuppressedByTombstone { current: ReplicatedValueRecord },
}
impl HintReplayDecision {
pub fn outcome(&self) -> HintOutcome {
match self {
Self::Replayed { .. } => HintOutcome::Replayed,
Self::SuppressedByNewerValue { .. } | Self::SuppressedByTombstone { .. } => {
HintOutcome::DroppedExpired
}
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct HintMetrics {
pub hints_stored_total: u64,
pub hints_replayed_total: u64,
pub hints_dropped_total: u64,
pub hint_store_bytes: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct HintError {
message: String,
}
impl fmt::Display for HintError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.message)
}
}
impl std::error::Error for HintError {}
pub trait HintStore {
fn store(
&mut self,
hint: Hint,
required_by_consistency: bool,
now_millis: u64,
) -> Result<HintOutcome, HintError>;
fn drain_for(&mut self, target: &ClusterNodeId) -> Result<Vec<Hint>, HintError>;
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct InMemoryHintStore {
budget: HintBudget,
hints: BTreeMap<ClusterNodeId, VecDeque<Hint>>,
repair_marks: BTreeSet<String>,
metrics: HintMetrics,
}
impl InMemoryHintStore {
pub fn new(budget: HintBudget) -> Self {
Self {
budget,
hints: BTreeMap::new(),
repair_marks: BTreeSet::new(),
metrics: HintMetrics::default(),
}
}
pub fn metrics(&self) -> HintMetrics {
let mut metrics = self.metrics;
metrics.hint_store_bytes = self.total_bytes();
metrics
}
pub fn is_marked_for_repair(&self, key: &str) -> bool {
self.repair_marks.contains(key)
}
pub fn len(&self) -> usize {
self.hints.values().map(VecDeque::len).sum()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
fn total_bytes(&self) -> u64 {
self.hints
.values()
.flat_map(|hints| hints.iter())
.map(Hint::approx_bytes)
.sum()
}
fn mark_for_repair(&mut self, key: &str) {
self.repair_marks.insert(key.to_owned());
}
}
impl HintStore for InMemoryHintStore {
fn store(
&mut self,
hint: Hint,
required_by_consistency: bool,
now_millis: u64,
) -> Result<HintOutcome, HintError> {
if required_by_consistency {
return Ok(HintOutcome::RequiredReplicaMiss);
}
if hint.is_expired(now_millis, self.budget) {
self.metrics.hints_dropped_total = self.metrics.hints_dropped_total.saturating_add(1);
self.mark_for_repair(&hint.key);
return Ok(HintOutcome::DroppedExpired);
}
let next_count = self.len().saturating_add(1);
let next_bytes = self.total_bytes().saturating_add(hint.approx_bytes());
if next_count > self.budget.max_hints || next_bytes > self.budget.max_bytes {
self.metrics.hints_dropped_total = self.metrics.hints_dropped_total.saturating_add(1);
self.mark_for_repair(&hint.key);
return Ok(HintOutcome::DroppedOverBudget);
}
self.hints
.entry(hint.target.clone())
.or_default()
.push_back(hint);
self.metrics.hints_stored_total = self.metrics.hints_stored_total.saturating_add(1);
Ok(HintOutcome::Stored)
}
fn drain_for(&mut self, target: &ClusterNodeId) -> Result<Vec<Hint>, HintError> {
let Some(hints) = self.hints.remove(target) else {
return Ok(Vec::new());
};
Ok(hints.into_iter().collect())
}
}
pub fn apply_hint(
current: Option<&ReplicatedValueRecord>,
hint: &Hint,
) -> Result<HintReplayDecision, HintError> {
let candidate = hint.to_record();
let Some(current) = current else {
return Ok(HintReplayDecision::Replayed { record: candidate });
};
if current.is_tombstone() && (current.version, current.epoch) >= (hint.version, hint.epoch) {
return Ok(HintReplayDecision::SuppressedByTombstone {
current: current.clone(),
});
}
let merged = current.clone().merge(candidate.clone());
if merged == candidate {
Ok(HintReplayDecision::Replayed { record: candidate })
} else {
Ok(HintReplayDecision::SuppressedByNewerValue {
current: current.clone(),
})
}
}
pub fn replay_hints<I>(
hints: I,
current_by_key: &BTreeMap<String, ReplicatedValueRecord>,
) -> Result<Vec<HintReplayDecision>, HintError>
where
I: IntoIterator<Item = Hint>,
{
hints
.into_iter()
.map(|hint| apply_hint(current_by_key.get(&hint.key), &hint))
.collect()
}