use std::collections::BTreeMap;
use std::fmt;
use serde::{Deserialize, Serialize};
use crate::grid::session_context::{PartitionKey, SessionWatermark, VersionStamp};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CausalSummary {
cap: usize,
deps: BTreeMap<PartitionKey, VersionStamp>,
coarse_floor: Option<VersionStamp>,
coarsened_total: u64,
}
impl CausalSummary {
pub fn new(cap: usize) -> Self {
Self {
cap: cap.max(1),
deps: BTreeMap::new(),
coarse_floor: None,
coarsened_total: 0,
}
}
pub fn from_watermark(watermark: &SessionWatermark) -> Self {
let mut summary = Self::new(watermark.cap());
for (key, stamp) in watermark.entries() {
summary.observe(key.clone(), *stamp);
}
summary
}
pub fn cap(&self) -> usize {
self.cap
}
pub fn len(&self) -> usize {
self.deps.len()
}
pub fn is_empty(&self) -> bool {
self.deps.is_empty() && self.coarse_floor.is_none()
}
pub fn coarsened_total(&self) -> u64 {
self.coarsened_total
}
pub fn coarse_floor(&self) -> Option<VersionStamp> {
self.coarse_floor
}
pub fn dependencies(&self) -> &BTreeMap<PartitionKey, VersionStamp> {
&self.deps
}
pub fn dependency_bytes(&self) -> u64 {
const ENTRY_BYTES: u64 = 56;
const COARSE_FLOOR_BYTES: u64 = 24;
(self.deps.len() as u64 * ENTRY_BYTES)
+ u64::from(self.coarse_floor.is_some()) * COARSE_FLOOR_BYTES
}
pub fn observe(&mut self, key: PartitionKey, stamp: VersionStamp) {
if let Some(existing) = self.deps.get_mut(&key) {
if stamp > *existing {
*existing = stamp;
}
return;
}
if self.deps.len() >= self.cap {
self.coarsen_one();
}
self.deps.insert(key, stamp);
}
pub fn not_yet_applied(&self, applied: &AppliedSet) -> Vec<CausalDependencyMissing> {
let mut missing = Vec::new();
if let Some(required) = self.coarse_floor {
if !applied.covers_coarse_floor(required) {
missing.push(CausalDependencyMissing::CoarseFloor { required });
}
}
for (key, required) in &self.deps {
if !applied.covers(key, *required) {
missing.push(CausalDependencyMissing::Exact {
key: key.clone(),
required: *required,
});
}
}
missing
}
pub fn gc_stable(&mut self, applied: &AppliedSet) -> usize {
let before = self.deps.len() + usize::from(self.coarse_floor.is_some());
self.deps.retain(|key, stamp| !applied.covers(key, *stamp));
if self
.coarse_floor
.is_some_and(|floor| applied.covers_coarse_floor(floor))
{
self.coarse_floor = None;
}
before - (self.deps.len() + usize::from(self.coarse_floor.is_some()))
}
fn coarsen_one(&mut self) {
if let Some((key, stamp)) = self
.deps
.iter()
.min_by_key(|(_, stamp)| **stamp)
.map(|(key, stamp)| (key.clone(), *stamp))
{
self.deps.remove(&key);
self.coarse_floor = Some(self.coarse_floor.map_or(stamp, |floor| floor.max(stamp)));
self.coarsened_total = self.coarsened_total.saturating_add(1);
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct AppliedSet {
applied: BTreeMap<PartitionKey, VersionStamp>,
stable_floor: Option<VersionStamp>,
}
impl AppliedSet {
pub fn new() -> Self {
Self::default()
}
pub fn mark_applied(&mut self, key: PartitionKey, stamp: VersionStamp) {
self.applied
.entry(key)
.and_modify(|existing| *existing = (*existing).max(stamp))
.or_insert(stamp);
}
pub fn mark_stable_floor(&mut self, stamp: VersionStamp) {
self.stable_floor = Some(self.stable_floor.map_or(stamp, |floor| floor.max(stamp)));
}
pub fn stable_floor(&self) -> Option<VersionStamp> {
self.stable_floor
}
pub fn covers(&self, key: &PartitionKey, required: VersionStamp) -> bool {
self.covers_coarse_floor(required)
|| self
.applied
.get(key)
.map(|stamp| stamp.covers(required))
.unwrap_or(false)
}
pub fn covers_coarse_floor(&self, required: VersionStamp) -> bool {
self.stable_floor
.map(|floor| floor.covers(required))
.unwrap_or(false)
}
pub fn entries(&self) -> &BTreeMap<PartitionKey, VersionStamp> {
&self.applied
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CausalDependencyMissing {
Exact {
key: PartitionKey,
required: VersionStamp,
},
CoarseFloor {
required: VersionStamp,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ApplyDecision {
Apply,
Defer {
missing: Vec<CausalDependencyMissing>,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CausalWrite {
pub key: PartitionKey,
pub stamp: VersionStamp,
pub deps: CausalSummary,
}
impl CausalWrite {
pub fn new(key: PartitionKey, stamp: VersionStamp, deps: CausalSummary) -> Self {
Self { key, stamp, deps }
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CausalApplyDeferred {
pub missing: Vec<CausalDependencyMissing>,
}
impl fmt::Display for CausalApplyDeferred {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
formatter,
"causal write deferred with {} missing dependencies",
self.missing.len()
)
}
}
impl std::error::Error for CausalApplyDeferred {}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct CausalConsistencyMetrics {
pub causal_writes_deferred_total: u64,
pub causal_summary_coarsened_total: u64,
pub causal_dependency_bytes: u64,
}
impl From<&CausalSummary> for CausalConsistencyMetrics {
fn from(summary: &CausalSummary) -> Self {
Self {
causal_writes_deferred_total: 0,
causal_summary_coarsened_total: summary.coarsened_total(),
causal_dependency_bytes: summary.dependency_bytes(),
}
}
}
pub fn causal_apply(local_applied: &AppliedSet, write_deps: &CausalSummary) -> ApplyDecision {
let missing = write_deps.not_yet_applied(local_applied);
if missing.is_empty() {
ApplyDecision::Apply
} else {
ApplyDecision::Defer { missing }
}
}
pub fn apply_causal_write(
local_applied: &mut AppliedSet,
write: CausalWrite,
) -> Result<(), CausalApplyDeferred> {
match causal_apply(local_applied, &write.deps) {
ApplyDecision::Apply => {
local_applied.mark_applied(write.key, write.stamp);
Ok(())
}
ApplyDecision::Defer { missing } => Err(CausalApplyDeferred { missing }),
}
}