use std::collections::{BTreeMap, HashMap};
use std::sync::{Mutex, MutexGuard, PoisonError};
use super::codec::{self, DedupRecord};
use super::{DedupCache, DedupEntry, InFlightClaim};
use crate::durability::DurabilityError;
use crate::durability::head_checkpoint::READ_BATCH_SIZE;
pub(super) const IN_FLIGHT_CHECKPOINT_EVERY: u64 = 256;
const CLAIM_PAGE: usize = 2;
const INDEX_APPEND_ATTEMPTS: usize = 64;
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
struct OpenClaim {
row: u64,
claim: u64,
}
#[derive(Debug)]
struct IndexState {
start: u64,
tail: u64,
open: HashMap<String, OpenClaim>,
reserved: BTreeMap<u64, usize>,
carrying: bool,
recorded: bool,
}
impl IndexState {
fn empty() -> Self {
Self {
start: 0,
tail: 0,
open: HashMap::new(),
reserved: BTreeMap::new(),
carrying: false,
recorded: false,
}
}
fn floor(&self) -> u64 {
let open = self.open.values().map(|claim| claim.row).min();
let reserved = self.reserved.keys().next().copied();
[open, reserved]
.into_iter()
.flatten()
.fold(self.tail, u64::min)
}
fn reserve(&mut self) -> u64 {
let floor = self.tail;
*self.reserved.entry(floor).or_insert(0) += 1;
floor
}
fn unreserve(&mut self, floor: u64) {
if let Some(count) = self.reserved.get_mut(&floor) {
*count = count.saturating_sub(1);
if *count == 0 {
self.reserved.remove(&floor);
}
}
}
fn appended(&mut self, after: u64) {
self.tail = self.tail.max(after);
}
fn leave(&mut self, key: &str, settled: u64) -> bool {
if self.open.get(key).is_some_and(|open| open.claim <= settled) {
self.open.remove(key);
return true;
}
false
}
fn plan(&mut self) -> Option<Advance> {
if self.carrying {
return None;
}
let floor = self.floor();
if floor >= self.start.saturating_add(IN_FLIGHT_CHECKPOINT_EVERY) {
return Some(Advance::Record(floor));
}
if self.tail < self.start.saturating_add(2 * IN_FLIGHT_CHECKPOINT_EVERY) {
return None;
}
let horizon = self.tail.saturating_sub(IN_FLIGHT_CHECKPOINT_EVERY);
let mut keys: Vec<(String, OpenClaim)> = self
.open
.iter()
.filter(|(_, open)| open.row < horizon)
.map(|(key, open)| (key.clone(), *open))
.collect();
if keys.is_empty() {
return None;
}
keys.sort_unstable();
self.carrying = true;
let at = self.reserve();
Some(Advance::Carry { keys, at })
}
fn carried(
&mut self,
keys: &[(String, OpenClaim)],
first: Option<u64>,
at: u64,
) -> Option<u64> {
if let Some(first) = first {
let mut row = first;
for (key, carried) in keys {
if let Some(open) = self.open.get_mut(key) {
if open.row == carried.row {
open.row = row;
}
}
row = row.saturating_add(1);
}
self.appended(row);
}
self.unreserve(at);
self.carrying = false;
let floor = self.floor();
(floor > self.start).then_some(floor)
}
}
enum Advance {
Record(u64),
Carry {
keys: Vec<(String, OpenClaim)>,
at: u64,
},
}
#[derive(Debug, Default)]
pub(super) struct InFlightIndex(Mutex<Option<IndexState>>);
impl InFlightIndex {
pub(super) fn empty() -> Self {
Self(Mutex::new(Some(IndexState::empty())))
}
fn state(&self) -> MutexGuard<'_, Option<IndexState>> {
self.0.lock().unwrap_or_else(PoisonError::into_inner)
}
fn is_loaded(&self) -> bool {
self.state().is_some()
}
fn install(&self, loaded: IndexState) -> bool {
let mut state = self.state();
if state.is_some() {
return false;
}
*state = Some(loaded);
true
}
fn reserve(&self) -> Option<u64> {
self.state().as_mut().map(IndexState::reserve)
}
fn unreserve(&self, floor: u64) {
if let Some(state) = self.state().as_mut() {
state.unreserve(floor);
}
}
fn joined(&self, key: &str, row: u64, claim: u64) {
if let Some(state) = self.state().as_mut() {
state.appended(row.saturating_add(1));
state.open.insert(key.to_owned(), OpenClaim { row, claim });
}
}
fn unjoin(&self, key: &str, row: u64) {
if let Some(state) = self.state().as_mut() {
if state.open.get(key).is_some_and(|open| open.row == row) {
state.open.remove(key);
}
}
}
fn leave(&self, key: &str, settled: u64) -> bool {
self.with_state(|state| state.leave(key, settled))
.unwrap_or(false)
}
fn open_claims(&self) -> Vec<(String, u64)> {
self.state()
.as_ref()
.map(|state| {
state
.open
.iter()
.map(|(key, open)| (key.clone(), open.claim))
.collect()
})
.unwrap_or_default()
}
fn awaits_record(&self) -> bool {
self.state().as_ref().is_some_and(|state| !state.recorded)
}
fn start_recorded(&self) {
if let Some(state) = self.state().as_mut() {
state.recorded = true;
}
}
fn start(&self) -> Option<u64> {
self.state().as_ref().map(|state| state.start)
}
fn recorded(&self, start: u64) {
if let Some(state) = self.state().as_mut() {
state.start = state.start.max(start);
}
}
fn plan(&self) -> Option<Advance> {
self.with_state(IndexState::plan).flatten()
}
fn carried(&self, keys: &[(String, OpenClaim)], first: Option<u64>, at: u64) -> Option<u64> {
self.with_state(|state| state.carried(keys, first, at))
.flatten()
}
fn with_state<T>(&self, apply: impl FnOnce(&mut IndexState) -> T) -> Option<T> {
self.state().as_mut().map(apply)
}
}
impl DedupCache {
fn index_stream(&self) -> String {
format!("{}/in-flight", self.namespace)
}
fn index_start_key(&self) -> String {
format!("{}/in-flight-start", self.namespace)
}
pub async fn in_flight_claims(&self) -> Result<Vec<InFlightClaim>, DurabilityError> {
if !self.in_flight.is_loaded() {
let (state, claims) = self.load_in_flight_index().await?;
if self.in_flight.install(state) {
self.advance_in_flight().await;
return Ok(claims);
}
}
let mut held = self.in_flight.open_claims();
held.sort_unstable();
let mut claims = Vec::new();
for (key, claim) in held {
if let Some(entry) = self.claim_at(&key, claim).await? {
claims.push(InFlightClaim::of(key, &entry));
}
}
Ok(claims)
}
pub(super) async fn ensure_in_flight_index(&self) -> Result<(), DurabilityError> {
if self.in_flight.is_loaded() {
return Ok(());
}
let (state, _) = self.load_in_flight_index().await?;
if self.in_flight.install(state) {
self.advance_in_flight().await;
}
Ok(())
}
pub(super) async fn join_in_flight(
&self,
key: &str,
claim: u64,
) -> Result<u64, DurabilityError> {
self.ensure_in_flight_index().await?;
self.record_empty_start().await?;
let floor = self.in_flight.reserve().ok_or_else(|| {
DurabilityError::ConfigError("in-flight index is not loaded".to_owned())
})?;
let appended = self
.append_index_rows(vec![codec::encode_join(key, claim)?], floor)
.await;
if let Ok(row) = appended {
self.in_flight.joined(key, row, claim);
}
self.in_flight.unreserve(floor);
appended
}
pub(super) async fn record_empty_start(&self) -> Result<(), DurabilityError> {
if !self.in_flight.awaits_record() {
return Ok(());
}
match self.store.cas(&self.index_start_key(), 0, 1).await {
Ok(()) | Err(DurabilityError::CursorRegression { .. }) => {
self.in_flight.start_recorded();
Ok(())
}
Err(error) => Err(error),
}
}
pub(super) fn unjoin_in_flight(&self, key: &str, row: u64) {
self.in_flight.unjoin(key, row);
}
pub(super) async fn leave_in_flight(&self, key: &str, settled: u64) {
if self.in_flight.leave(key, settled) {
self.advance_in_flight().await;
}
}
async fn advance_in_flight(&self) {
if let Err(error) = self.try_advance_in_flight().await {
tracing::warn!(
namespace = %self.namespace,
%error,
"the in-flight index start could not be recorded; the next start reads the \
rows since the previous record"
);
}
}
async fn try_advance_in_flight(&self) -> Result<(), DurabilityError> {
let target = match self.in_flight.plan() {
None => return Ok(()),
Some(Advance::Record(target)) => target,
Some(Advance::Carry { keys, at }) => {
let rows = keys
.iter()
.map(|(key, open)| codec::encode_join(key, open.claim))
.collect::<Result<Vec<_>, _>>();
let appended = match rows {
Ok(rows) => self.append_index_rows(rows, at).await,
Err(error) => Err(error),
};
let target = self
.in_flight
.carried(&keys, appended.as_ref().ok().copied(), at);
appended?;
let Some(target) = target else {
return Ok(());
};
target
}
};
let Some(start) = self.in_flight.start() else {
return Ok(());
};
match self
.store
.cas(
&self.index_start_key(),
start.saturating_add(1),
target.saturating_add(1),
)
.await
{
Ok(()) => self.in_flight.recorded(target),
Err(DurabilityError::CursorRegression { stored, .. }) => {
self.in_flight.recorded(stored.saturating_sub(1));
}
Err(error) => return Err(error),
}
Ok(())
}
async fn append_index_rows(
&self,
rows: Vec<Vec<u8>>,
expected: u64,
) -> Result<u64, DurabilityError> {
let stream = self.index_stream();
let mut expected = expected;
let mut last_conflict = None;
for _ in 0..INDEX_APPEND_ATTEMPTS {
match self.store.append_all(&stream, rows.clone(), expected).await {
Ok(_) => return Ok(expected),
Err(DurabilityError::SequenceConflict { actual, .. }) if actual > expected => {
last_conflict = Some(actual);
expected = actual;
}
Err(error) => return Err(error),
}
}
Err(DurabilityError::SequenceConflict {
expected,
actual: last_conflict.unwrap_or(expected),
})
}
async fn load_in_flight_index(
&self,
) -> Result<(IndexState, Vec<InFlightClaim>), DurabilityError> {
let start = match self.store.read_value(&self.index_start_key()).await? {
Some(recorded) if recorded > 0 => recorded - 1,
_ => self.migrate_in_flight_index().await?,
};
let (named, tail) = self.read_index_rows(start).await?;
let mut open = HashMap::new();
let mut claims = Vec::new();
for (key, (row, claim)) in named {
let held = match claim {
Some(claim) => self
.claim_at(&key, claim)
.await?
.map(|entry| (claim, entry)),
None => self.claim_without_sequence(&key).await?,
};
let Some((claim, entry)) = held else {
continue;
};
open.insert(key.clone(), OpenClaim { row, claim });
claims.push(InFlightClaim::of(key, &entry));
}
let state = IndexState {
start,
tail,
open,
reserved: BTreeMap::new(),
carrying: false,
recorded: true,
};
Ok((state, claims))
}
async fn read_index_rows(
&self,
start: u64,
) -> Result<(BTreeMap<String, (u64, Option<u64>)>, u64), DurabilityError> {
let stream = self.index_stream();
let mut named = BTreeMap::new();
let mut tail = start;
loop {
let page = self.store.read_from(&stream, tail, READ_BATCH_SIZE).await?;
for row in &page {
if row.sequence < tail {
return Err(DurabilityError::ConfigError(format!(
"stream {stream} answered sequence {} to a read from {tail}",
row.sequence
)));
}
let (key, claim) = codec::decode_join(&row.payload)?;
named.insert(key, (row.sequence, claim));
tail = row.sequence.checked_add(1).ok_or_else(|| {
DurabilityError::ConfigError(format!("stream {stream} sequence overflow"))
})?;
}
if page.len() < READ_BATCH_SIZE {
break;
}
}
Ok((named, tail))
}
async fn claim_at(&self, key: &str, claim: u64) -> Result<Option<DedupEntry>, DurabilityError> {
let stream = self.stream_key_for(key);
let rows = self.store.read_from(&stream, claim, CLAIM_PAGE).await?;
let mut expected = claim;
for row in &rows {
if row.sequence != expected {
return Err(DurabilityError::ConfigError(format!(
"stream {stream} answered sequence {} where {expected} was read",
row.sequence
)));
}
expected = expected.saturating_add(1);
}
let Some(claimed) = rows.first() else {
return Ok(None);
};
let record = DedupRecord::deserialize(&claimed.payload)?;
if record.idempotency_key() != key {
return Err(DurabilityError::DedupCollision {
key: key.to_owned(),
});
}
if rows.len() > 1 {
return Ok(None);
}
Ok(record
.into_active()
.filter(|entry| entry.receipt().is_none()))
}
async fn claim_without_sequence(
&self,
key: &str,
) -> Result<Option<(u64, DedupEntry)>, DurabilityError> {
let snapshot = self.load_snapshot(&self.stream_key_for(key), key).await?;
let claim = match snapshot.current.as_ref() {
Some(entry) if entry.receipt().is_none() => snapshot.claim_sequence()?,
_ => return Ok(None),
};
Ok(snapshot.current.map(|entry| (claim, entry)))
}
async fn migrate_in_flight_index(&self) -> Result<u64, DurabilityError> {
let mut keys = Vec::new();
let mut walk = self.latest_records();
while let Some(batch) = walk.next_batch().await? {
for candidate in batch {
if let DedupRecord::Active(entry) = candidate.record {
if entry.receipt().is_none() {
keys.push((entry.idempotency_key, candidate.sequence));
}
}
}
}
let start = if keys.is_empty() {
0
} else {
let rows = keys
.iter()
.map(|(key, claim)| codec::encode_join(key, *claim))
.collect::<Result<Vec<_>, _>>()?;
self.append_index_rows(rows, 0).await?
};
let recorded = match self
.store
.cas(&self.index_start_key(), 0, start.saturating_add(1))
.await
{
Ok(()) => start,
Err(DurabilityError::CursorRegression { stored, .. }) if stored > 0 => stored - 1,
Err(error) => return Err(error),
};
tracing::info!(
namespace = %self.namespace,
claims = keys.len(),
scanned = walk.scanned(),
"the dedup namespace was walked once to build its in-flight index"
);
Ok(recorded)
}
}