use core::fmt;
use core::num::NonZeroUsize;
use core::str::FromStr;
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use thiserror::Error;
use crate::domain::RequestId;
use crate::domain::session::SessionId;
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct IdempotencyKey(Arc<str>);
impl IdempotencyKey {
#[must_use]
pub fn new(value: impl AsRef<str>) -> Self {
Self(Arc::from(value.as_ref()))
}
#[must_use]
pub fn from_request_id(id: &RequestId) -> Self {
Self::new(id.to_string())
}
#[must_use]
pub fn for_request(session: &SessionId, id: &RequestId) -> Self {
Self(Arc::from(format!("{session}/{id}").as_str()))
}
#[must_use]
pub fn for_method_call(
session: &SessionId,
id: &RequestId,
method: &str,
target: &str,
) -> Self {
Self(Arc::from(
format!("{session}/{id}/{method}/{target}").as_str(),
))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for IdempotencyKey {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.0)
}
}
impl From<&RequestId> for IdempotencyKey {
fn from(id: &RequestId) -> Self {
Self::from_request_id(id)
}
}
impl FromStr for IdempotencyKey {
type Err = core::convert::Infallible;
fn from_str(text: &str) -> Result<Self, Self::Err> {
Ok(Self::new(text))
}
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ClaimToken {
key: IdempotencyKey,
generation: u64,
}
impl ClaimToken {
#[must_use]
pub const fn key(&self) -> &IdempotencyKey {
&self.key
}
#[must_use]
pub const fn generation(&self) -> u64 {
self.generation
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Claim<T> {
Claimed {
token: ClaimToken,
},
Completed(T),
InFlight,
}
impl<T> Claim<T> {
#[must_use]
pub const fn is_execution_mine(&self) -> bool {
matches!(self, Self::Claimed { .. })
}
#[must_use]
pub const fn recorded(&self) -> Option<&T> {
match self {
Self::Completed(result) => Some(result),
Self::Claimed { .. } | Self::InFlight => None,
}
}
#[must_use]
pub const fn token(&self) -> Option<&ClaimToken> {
match self {
Self::Claimed { token } => Some(token),
Self::Completed(_) | Self::InFlight => None,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum Slot<T> {
Held {
generation: u64,
},
Done(T),
}
#[derive(Debug)]
pub struct CompletedResults<T> {
slots: HashMap<IdempotencyKey, Slot<T>>,
completed_order: VecDeque<IdempotencyKey>,
next_generation: u64,
capacity: NonZeroUsize,
}
impl<T> CompletedResults<T> {
#[must_use]
pub fn with_capacity(capacity: NonZeroUsize) -> Self {
Self {
slots: HashMap::with_capacity(capacity.get()),
completed_order: VecDeque::new(),
next_generation: 0,
capacity,
}
}
#[must_use]
pub fn len(&self) -> usize {
self.slots.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.slots.is_empty()
}
#[must_use]
pub const fn capacity(&self) -> usize {
self.capacity.get()
}
pub fn claim(&mut self, key: &IdempotencyKey) -> Result<Claim<T>, IdempotencyError>
where
T: Clone,
{
if let Some(slot) = self.slots.get(key) {
return Ok(match slot {
Slot::Done(result) => Claim::Completed(result.clone()),
Slot::Held { .. } => Claim::InFlight,
});
}
if self.slots.len() >= self.capacity.get() {
let Some(oldest) = self.completed_order.pop_front() else {
return Err(IdempotencyError::StoreFull {
capacity: self.capacity.get(),
});
};
self.slots.remove(&oldest);
}
self.next_generation += 1;
let generation = self.next_generation;
self.slots.insert(key.clone(), Slot::Held { generation });
Ok(Claim::Claimed {
token: ClaimToken {
key: key.clone(),
generation,
},
})
}
pub fn complete(&mut self, token: &ClaimToken, result: T) -> Result<(), StaleClaim> {
match self.slots.get(&token.key) {
Some(Slot::Held { generation }) if *generation == token.generation => {}
_ => return Err(StaleClaim::new(token)),
}
self.slots.insert(token.key.clone(), Slot::Done(result));
self.completed_order.push_back(token.key.clone());
Ok(())
}
pub fn abandon(&mut self, token: &ClaimToken) -> Result<(), StaleClaim> {
match self.slots.get(&token.key) {
Some(Slot::Held { generation }) if *generation == token.generation => {}
_ => return Err(StaleClaim::new(token)),
}
self.slots.remove(&token.key);
Ok(())
}
#[must_use]
pub fn completed(&self, key: &IdempotencyKey) -> Option<&T> {
match self.slots.get(key) {
Some(Slot::Done(result)) => Some(result),
Some(Slot::Held { .. }) | None => None,
}
}
#[must_use]
pub fn is_in_flight(&self, key: &IdempotencyKey) -> bool {
matches!(self.slots.get(key), Some(Slot::Held { .. }))
}
}
#[derive(Debug, Clone, PartialEq, Eq, Error)]
#[error(
"the claim on key {key} from attempt {stale_generation} is stale: the key has \
since been released and re-claimed, and the later attempt owns the result"
)]
pub struct StaleClaim {
key: IdempotencyKey,
stale_generation: u64,
}
impl StaleClaim {
#[must_use]
pub fn new(token: &ClaimToken) -> Self {
Self {
key: token.key.clone(),
stale_generation: token.generation,
}
}
#[must_use]
pub const fn key(&self) -> &IdempotencyKey {
&self.key
}
#[must_use]
pub const fn stale_generation(&self) -> u64 {
self.stale_generation
}
}
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum IdempotencyError {
#[error(
"the idempotency store holds {capacity} keys and cannot claim another without risking a lost result"
)]
StoreFull {
capacity: usize,
},
#[error(transparent)]
Stale(#[from] StaleClaim),
#[error("the idempotency store could not answer: {reason}")]
Unavailable {
reason: &'static str,
},
}
#[cfg(test)]
mod tests {
use super::{
Claim, ClaimToken, CompletedResults, IdempotencyError, IdempotencyKey, StaleClaim,
};
use crate::domain::RequestId;
use core::num::NonZeroUsize;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::thread;
fn store(capacity: usize) -> Result<CompletedResults<String>, Failure> {
let capacity = NonZeroUsize::new(capacity).ok_or(Failure::ZeroCapacity)?;
Ok(CompletedResults::with_capacity(capacity))
}
#[derive(Debug, PartialEq, Eq)]
enum Failure {
ZeroCapacity,
WrongClaim(&'static str),
WrongResult(&'static str),
WrongError(&'static str),
NotStale,
Stale,
Unexpected(String),
}
impl From<StaleClaim> for Failure {
fn from(_: StaleClaim) -> Self {
Self::Stale
}
}
impl From<IdempotencyError> for Failure {
fn from(error: IdempotencyError) -> Self {
Self::Unexpected(error.to_string())
}
}
fn token_of<T>(claim: &Claim<T>) -> Result<ClaimToken, Failure> {
claim
.token()
.cloned()
.ok_or(Failure::WrongClaim("not claimed"))
}
#[test]
fn a_token_reads_back_the_key_and_the_attempt() -> Result<(), Failure> {
let mut results = store(4)?;
let key = IdempotencyKey::new("req-1");
let token = token_of(&results.claim(&key)?)?;
if token.key() != &key {
return Err(Failure::WrongResult("a different key"));
}
let first = token.generation();
if first == 0 {
return Err(Failure::WrongResult("generation zero"));
}
results.abandon(&token)?;
let retaken = token_of(&results.claim(&key)?)?;
if retaken.generation() <= first {
return Err(Failure::WrongResult(
"a re-claim did not advance the generation",
));
}
Ok(())
}
#[test]
fn a_fresh_key_is_claimed_and_then_held() -> Result<(), Failure> {
let mut results = store(4)?;
let key = IdempotencyKey::new("req-1");
let claim = results.claim(&key)?;
if !claim.is_execution_mine() {
return Err(Failure::WrongClaim("a fresh key was not claimed"));
}
if !results.is_in_flight(&key) {
return Err(Failure::WrongResult("a claimed key did not read as held"));
}
if results.completed(&key).is_some() {
return Err(Failure::WrongResult("a held key read as completed"));
}
Ok(())
}
#[test]
fn a_held_key_refuses_a_second_claim() -> Result<(), Failure> {
let mut results = store(4)?;
let key = IdempotencyKey::new("req-1");
let _first = results.claim(&key)?;
match results.claim(&key)? {
Claim::InFlight => Ok(()),
Claim::Claimed { .. } => Err(Failure::WrongClaim("a held key was re-claimed")),
Claim::Completed(_) => Err(Failure::WrongClaim("a held key reported a result")),
}
}
#[test]
fn a_completed_key_hands_back_the_result_without_a_token() -> Result<(), Failure> {
let mut results = store(4)?;
let key = IdempotencyKey::new("req-1");
let token = token_of(&results.claim(&key)?)?;
results.complete(&token, "original".to_owned())?;
match results.claim(&key)? {
Claim::Completed(result) => {
if result != "original" {
return Err(Failure::WrongResult(
"the recorded result was not handed back",
));
}
if results.is_in_flight(&key) {
return Err(Failure::WrongResult("a completed key read as held"));
}
Ok(())
}
Claim::Claimed { .. } => Err(Failure::WrongClaim("a completed key was re-claimed")),
Claim::InFlight => Err(Failure::WrongClaim("a completed key reported in flight")),
}
}
#[test]
fn abandoning_releases_the_key_for_a_retry() -> Result<(), Failure> {
let mut results = store(4)?;
let key = IdempotencyKey::new("req-1");
let token = token_of(&results.claim(&key)?)?;
results.abandon(&token)?;
if results.is_in_flight(&key) {
return Err(Failure::WrongResult("an abandoned key still reads as held"));
}
if results.completed(&key).is_some() {
return Err(Failure::WrongResult("an abandoned key read as completed"));
}
if !results.claim(&key)?.is_execution_mine() {
return Err(Failure::WrongClaim("an abandoned key was not re-claimable"));
}
Ok(())
}
#[test]
fn a_late_completion_from_an_abandoned_claim_is_refused() -> Result<(), Failure> {
let mut results = store(4)?;
let key = IdempotencyKey::new("req-1");
let slow = token_of(&results.claim(&key)?)?;
results.abandon(&slow)?;
let retry = token_of(&results.claim(&key)?)?;
results.complete(&retry, "second".to_owned())?;
if results.complete(&slow, "first".to_owned()).is_ok() {
return Err(Failure::NotStale);
}
match results.completed(&key) {
Some(result) if result == "second" => Ok(()),
Some(_) => Err(Failure::WrongResult(
"the stale result overwrote the current one",
)),
None => Err(Failure::WrongResult("the completed key read as absent")),
}
}
#[test]
fn a_late_abandon_from_an_abandoned_claim_is_refused() -> Result<(), Failure> {
let mut results = store(4)?;
let key = IdempotencyKey::new("req-1");
let slow = token_of(&results.claim(&key)?)?;
results.abandon(&slow)?;
let retry = token_of(&results.claim(&key)?)?;
if results.abandon(&slow).is_ok() {
return Err(Failure::NotStale);
}
if !results.is_in_flight(&key) {
return Err(Failure::WrongResult(
"a late abandon released the live claim",
));
}
if results.claim(&key)?.is_execution_mine() {
return Err(Failure::WrongClaim("the live claim was stealable"));
}
results.complete(&retry, "second".to_owned())?;
Ok(())
}
#[test]
fn completing_an_unclaimed_key_is_refused() -> Result<(), Failure> {
let mut results = store(4)?;
let key = IdempotencyKey::new("never-seen");
let forged = ClaimToken {
key: key.clone(),
generation: 999,
};
if results.complete(&forged, "invented".to_owned()).is_ok() {
return Err(Failure::NotStale);
}
if results.completed(&key).is_some() {
return Err(Failure::WrongResult("a forged token installed a result"));
}
Ok(())
}
#[test]
fn exactly_one_of_many_concurrent_callers_executes() -> Result<(), Failure> {
const THREADS: usize = 64;
let results = Arc::new(Mutex::new(store(8)?));
let executions = Arc::new(AtomicUsize::new(0));
let mut workers = Vec::with_capacity(THREADS);
for _ in 0..THREADS {
let results = Arc::clone(&results);
let executions = Arc::clone(&executions);
workers.push(thread::spawn(move || {
let key = IdempotencyKey::new("req-1");
let mut held = results
.lock()
.map_err(|_| Failure::Unexpected("poisoned".to_owned()))?;
match held.claim(&key) {
Err(error) => Err(Failure::Unexpected(error.to_string())),
Ok(Claim::Claimed { token }) => {
executions.fetch_add(1, Ordering::SeqCst);
held.complete(&token, "ran".to_owned())
.map_err(Failure::from)
}
Ok(Claim::InFlight | Claim::Completed(_)) => Ok(()),
}
}));
}
for worker in workers {
worker
.join()
.map_err(|_| Failure::Unexpected("panicked".to_owned()))??;
}
let count = executions.load(Ordering::SeqCst);
if count != 1 {
return Err(Failure::Unexpected(format!("{count} executions")));
}
let held = results
.lock()
.map_err(|_| Failure::Unexpected("poisoned".to_owned()))?;
if held.len() != 1 {
return Err(Failure::Unexpected(format!("{} keys", held.len())));
}
drop(held);
Ok(())
}
#[test]
fn distinct_keys_do_not_interfere() -> Result<(), Failure> {
let mut results = store(4)?;
let first = IdempotencyKey::new("req-1");
let second = IdempotencyKey::new("req-2");
let token = token_of(&results.claim(&first)?)?;
if !results.claim(&second)?.is_execution_mine() {
return Err(Failure::WrongClaim("a second key was blocked by the first"));
}
results.complete(&token, "one".to_owned())?;
if results.completed(&second).is_some() {
return Err(Failure::WrongResult("completing one key completed another"));
}
Ok(())
}
#[test]
fn a_full_store_drops_its_oldest_result_for_a_new_key() -> Result<(), Failure> {
let mut results = store(2)?;
for index in 0..2 {
let key = IdempotencyKey::new(format!("req-{index}"));
let token = token_of(&results.claim(&key)?)?;
results.complete(&token, format!("result-{index}"))?;
}
let overflow = IdempotencyKey::new("req-overflow");
if !results.claim(&overflow)?.is_execution_mine() {
return Err(Failure::WrongClaim("a full store refused a new key"));
}
if results.len() != 2 {
return Err(Failure::WrongResult("the ceiling was exceeded"));
}
let kept = IdempotencyKey::new("req-1");
match results.claim(&kept)? {
Claim::Completed(result) if result == "result-1" => {}
_ => return Err(Failure::WrongResult("a full store lost the newer result")),
}
let dropped = IdempotencyKey::new("req-0");
if results.completed(&dropped).is_some() {
return Err(Failure::WrongResult("the oldest result was kept"));
}
Ok(())
}
#[test]
fn a_store_full_of_running_work_refuses_a_new_key() -> Result<(), Failure> {
let mut results = store(2)?;
for index in 0..2 {
let key = IdempotencyKey::new(format!("req-{index}"));
let _ = results.claim(&key)?;
}
match results.claim(&IdempotencyKey::new("req-overflow")) {
Err(IdempotencyError::StoreFull { capacity }) => {
if capacity != 2 {
return Err(Failure::WrongError("the ceiling was misreported"));
}
Ok(())
}
Err(other) => Err(Failure::Unexpected(other.to_string())),
Ok(_) => Err(Failure::WrongClaim("a full store handed out a claim")),
}
}
#[test]
fn the_key_is_the_request_id_a_retry_repeats() -> Result<(), Failure> {
let numeric = IdempotencyKey::from(&RequestId::number(7));
let text = IdempotencyKey::new("7");
if numeric != text {
return Err(Failure::Unexpected(format!("{numeric} != {text}")));
}
if numeric.as_str() != "7" {
return Err(Failure::Unexpected(numeric.to_string()));
}
let mut results = store(4)?;
let token = token_of(&results.claim(&numeric)?)?;
results.complete(&token, "once".to_owned())?;
match results.claim(&text)? {
Claim::Completed(result) if result == "once" => Ok(()),
_ => Err(Failure::WrongClaim("a re-spelled key was treated as new")),
}
}
#[test]
fn an_empty_store_claims_cleanly() -> Result<(), Failure> {
let results = store(4)?;
if !results.is_empty() {
return Err(Failure::WrongResult("a fresh store was not empty"));
}
let key = IdempotencyKey::new("req-1");
if results.completed(&key).is_some() || results.is_in_flight(&key) {
return Err(Failure::WrongResult(
"a fresh store reported an existing key",
));
}
Ok(())
}
}