use std::collections::{HashMap, VecDeque};
use std::sync::Mutex;
use crate::error::{Error, Result};
use crate::tokens::StoreFuture;
pub const DEFAULT_WINDOW_SECS: i64 = 24 * 60 * 60;
pub const DEFAULT_CAPACITY: usize = 10_000;
pub const DEFAULT_LEASE_SECS: i64 = 5 * 60;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum Claim {
Fresh,
InFlight,
Done,
}
pub trait DeliveryStore: Send + Sync + 'static {
fn claim<'a>(&'a self, webhook_id: &'a str, now: i64) -> StoreFuture<'a, Claim>;
fn complete<'a>(&'a self, webhook_id: &'a str, claimed_at: i64) -> StoreFuture<'a, ()>;
fn release<'a>(&'a self, webhook_id: &'a str, claimed_at: i64) -> StoreFuture<'a, ()>;
}
#[derive(Debug)]
pub struct InMemoryDeliveryStore {
window_secs: i64,
lease_secs: i64,
capacity: usize,
inner: Mutex<Claims>,
}
#[derive(Debug, Default)]
struct Claims {
by_id: HashMap<String, Entry>,
order: VecDeque<(String, i64)>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Entry {
at: i64,
done: bool,
}
impl Default for InMemoryDeliveryStore {
fn default() -> Self {
Self::new(DEFAULT_WINDOW_SECS, DEFAULT_CAPACITY)
}
}
impl InMemoryDeliveryStore {
#[must_use]
pub fn new(window_secs: i64, capacity: usize) -> Self {
Self {
window_secs,
lease_secs: DEFAULT_LEASE_SECS,
capacity: capacity.max(1),
inner: Mutex::new(Claims::default()),
}
}
#[must_use]
pub const fn with_lease(mut self, lease_secs: i64) -> Self {
self.lease_secs = lease_secs;
self
}
pub fn len(&self) -> Result<usize> {
self.lock().map(|claims| claims.by_id.len())
}
pub fn is_empty(&self) -> Result<bool> {
self.len().map(|len| len == 0)
}
fn lock(&self) -> Result<std::sync::MutexGuard<'_, Claims>> {
self.inner
.lock()
.map_err(|_| Error::Store("delivery store lock is poisoned".to_owned()))
}
#[allow(
clippy::significant_drop_tightening,
reason = "one lock covers the check and the insert"
)]
fn claim_now(&self, webhook_id: &str, now: i64) -> Result<Claim> {
let mut claims = self.lock()?;
claims.prune(now.saturating_sub(self.window_secs));
if let Some(entry) = claims.by_id.get(webhook_id).copied() {
if entry.done {
return Ok(Claim::Done);
}
if now.saturating_sub(entry.at) < self.lease_secs {
return Ok(Claim::InFlight);
}
}
while claims.by_id.len() >= self.capacity && !claims.by_id.contains_key(webhook_id) {
if !claims.pop_oldest() {
break;
}
}
claims.by_id.insert(
webhook_id.to_owned(),
Entry {
at: now,
done: false,
},
);
claims.order.push_back((webhook_id.to_owned(), now));
claims.compact(self.capacity);
Ok(Claim::Fresh)
}
}
impl DeliveryStore for InMemoryDeliveryStore {
fn claim<'a>(&'a self, webhook_id: &'a str, now: i64) -> StoreFuture<'a, Claim> {
let result = self.claim_now(webhook_id, now);
Box::pin(async move { result })
}
fn complete<'a>(&'a self, webhook_id: &'a str, claimed_at: i64) -> StoreFuture<'a, ()> {
let result = self.lock().map(|mut claims| {
if let Some(entry) = claims.by_id.get_mut(webhook_id)
&& entry.at == claimed_at
{
entry.done = true;
}
});
Box::pin(async move { result })
}
fn release<'a>(&'a self, webhook_id: &'a str, claimed_at: i64) -> StoreFuture<'a, ()> {
let result = self.lock().map(|mut claims| {
if claims
.by_id
.get(webhook_id)
.is_some_and(|entry| entry.at == claimed_at)
{
claims.by_id.remove(webhook_id);
}
});
Box::pin(async move { result })
}
}
impl Claims {
fn prune(&mut self, cutoff: i64) {
while self.order.front().is_some_and(|(_, at)| *at <= cutoff) {
self.pop_oldest();
}
}
fn compact(&mut self, capacity: usize) {
if self.order.len() > capacity.saturating_mul(2) {
let by_id = &self.by_id;
self.order
.retain(|(id, at)| by_id.get(id).is_some_and(|entry| entry.at == *at));
}
}
fn pop_oldest(&mut self) -> bool {
let Some((id, at)) = self.order.pop_front() else {
return false;
};
if self.by_id.get(&id).is_some_and(|entry| entry.at == at) {
self.by_id.remove(&id);
}
true
}
}
#[cfg(test)]
mod tests;