use std::{
collections::BTreeSet,
fmt,
sync::{Arc, Mutex},
};
use jacquard_macros::public_model;
use serde::{Deserialize, Serialize};
#[public_model]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ClaimRejected;
impl fmt::Display for ClaimRejected {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("claim is already held")
}
}
impl std::error::Error for ClaimRejected {}
#[derive(Clone)]
pub struct PendingClaims<Key: Ord> {
claimed: Arc<Mutex<BTreeSet<Key>>>,
}
impl<Key: Ord> Default for PendingClaims<Key> {
fn default() -> Self {
Self {
claimed: Arc::new(Mutex::new(BTreeSet::new())),
}
}
}
pub struct ClaimGuard<Key: Ord> {
claimed: Arc<Mutex<BTreeSet<Key>>>,
key: Option<Key>,
}
impl<Key> PendingClaims<Key>
where
Key: Clone + Ord,
{
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn try_claim(&self, key: Key) -> Result<ClaimGuard<Key>, ClaimRejected> {
let mut guard = self.claimed.lock().expect("pending claims lock");
if !guard.insert(key.clone()) {
return Err(ClaimRejected);
}
Ok(ClaimGuard {
claimed: Arc::clone(&self.claimed),
key: Some(key),
})
}
#[must_use]
pub fn contains(&self, key: &Key) -> bool {
self.claimed
.lock()
.expect("pending claims lock")
.contains(key)
}
}
impl<Key: Ord> ClaimGuard<Key> {
#[must_use]
pub fn key(&self) -> &Key {
self.key.as_ref().expect("claim guard key")
}
}
impl<Key: Ord> Drop for ClaimGuard<Key> {
fn drop(&mut self) {
if let Some(key) = self.key.take() {
self.claimed
.lock()
.expect("pending claims lock")
.remove(&key);
}
}
}
#[cfg(test)]
mod tests {
use super::{ClaimRejected, PendingClaims};
#[test]
fn duplicate_claims_are_rejected() {
let claims = PendingClaims::new();
let _guard = claims.try_claim("peer-a").expect("first claim");
assert_eq!(claims.try_claim("peer-a").err(), Some(ClaimRejected));
}
#[test]
fn dropped_guards_release_claims() {
let claims = PendingClaims::new();
let guard = claims.try_claim("peer-a").expect("claim");
assert!(claims.contains(&"peer-a"));
assert_eq!(guard.key(), &"peer-a");
drop(guard);
assert!(!claims.contains(&"peer-a"));
}
}