use std::collections::BTreeMap;
use std::sync::Mutex;
use std::time::Duration;
use af_context::{NotificationId, RequestContext};
use async_trait::async_trait;
use chrono::Utc;
use crate::{
AttemptEvent, DeliveryClass, DeliveryContext, DurableOutbox, NewOutboxItem, Notification,
NotificationStatus, NotifyError, NotifyStoreError, OutboxItem, Sender,
};
pub struct MemoryOutbox {
items: Mutex<BTreeMap<NotificationId, OutboxItem>>,
idempotency: Mutex<BTreeMap<(String, String, String), NotificationId>>,
changes: tokio::sync::broadcast::Sender<NotificationId>,
}
impl Default for MemoryOutbox {
fn default() -> Self {
let (changes, _) = tokio::sync::broadcast::channel(32);
Self {
items: Mutex::new(BTreeMap::new()),
idempotency: Mutex::new(BTreeMap::new()),
changes,
}
}
}
#[async_trait]
impl DurableOutbox for MemoryOutbox {
fn subscribe(&self) -> tokio::sync::broadcast::Receiver<NotificationId> {
self.changes.subscribe()
}
async fn enqueue(&self, item: NewOutboxItem) -> Result<OutboxItem, NotifyStoreError> {
let key = (
item.tenant_id.to_string(),
item.subject_id.to_string(),
item.idempotency_key.clone(),
);
let mut idempotency = self
.idempotency
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let mut items = self
.items
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(existing) = idempotency.get(&key).and_then(|id| items.get(id)) {
return if existing.channel == item.channel
&& existing.recipient == item.recipient
&& existing.notification == item.notification
&& existing.max_attempts == item.max_attempts.max(1)
{
Ok(existing.clone())
} else {
Err(NotifyStoreError::IdempotencyConflict)
};
}
let id = NotificationId::parse(format!("notification-{}", items.len() + 1))
.map_err(|error| NotifyStoreError::Invalid(error.to_string()))?;
let now = Utc::now();
let record = OutboxItem {
id: id.clone(),
tenant_id: item.tenant_id,
subject_id: item.subject_id,
channel: item.channel,
recipient: item.recipient,
notification: item.notification,
status: NotificationStatus::Pending,
attempts: 0,
max_attempts: item.max_attempts.max(1),
lease_version: 0,
last_error: None,
created_at: now,
updated_at: now,
};
items.insert(id.clone(), record.clone());
idempotency.insert(key, id);
let _ = self.changes.send(record.id.clone());
Ok(record)
}
async fn get(
&self,
context: &RequestContext,
id: &NotificationId,
) -> Result<OutboxItem, NotifyStoreError> {
self.items
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.get(id)
.filter(|item| {
item.tenant_id == context.tenant_id && item.subject_id == context.subject_id
})
.cloned()
.ok_or(NotifyStoreError::NotFound)
}
async fn list(
&self,
context: &RequestContext,
status: Option<NotificationStatus>,
limit: usize,
after: Option<&NotificationId>,
) -> Result<Vec<OutboxItem>, NotifyStoreError> {
Ok(self
.items
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.values()
.filter(|item| {
item.tenant_id == context.tenant_id
&& item.subject_id == context.subject_id
&& status.is_none_or(|status| item.status == status)
&& after.is_none_or(|after| item.id > *after)
})
.take(limit.clamp(1, 101))
.cloned()
.collect())
}
async fn attempts(
&self,
context: &RequestContext,
id: &NotificationId,
) -> Result<Vec<AttemptEvent>, NotifyStoreError> {
self.get(context, id).await.map(|_| Vec::new())
}
async fn claim(
&self,
_worker_id: &str,
_lease_secs: i64,
batch: usize,
) -> Result<Vec<OutboxItem>, NotifyStoreError> {
let mut items = self
.items
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let ids: Vec<_> = items
.values()
.filter(|item| {
matches!(
item.status,
NotificationStatus::Pending | NotificationStatus::RetryScheduled
)
})
.take(batch.clamp(1, 100))
.map(|item| item.id.clone())
.collect();
let claimed: Vec<_> = ids
.into_iter()
.filter_map(|id| {
let item = items.get_mut(&id)?;
item.status = NotificationStatus::Sending;
item.attempts += 1;
item.lease_version += 1;
Some(item.clone())
})
.collect();
for item in &claimed {
let _ = self.changes.send(item.id.clone());
}
Ok(claimed)
}
async fn mark_sent(
&self,
id: &NotificationId,
lease_version: i64,
) -> Result<(), NotifyStoreError> {
let mut items = self
.items
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let item = items.get_mut(id).ok_or(NotifyStoreError::NotFound)?;
if item.lease_version != lease_version {
return Err(NotifyStoreError::LeaseLost);
}
item.status = NotificationStatus::Sent;
let _ = self.changes.send(id.clone());
Ok(())
}
async fn record_failure(
&self,
id: &NotificationId,
lease_version: i64,
class: DeliveryClass,
_delay: Duration,
) -> Result<NotificationStatus, NotifyStoreError> {
let mut items = self
.items
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let item = items.get_mut(id).ok_or(NotifyStoreError::NotFound)?;
if item.lease_version != lease_version {
return Err(NotifyStoreError::LeaseLost);
}
item.status = if class == DeliveryClass::Permanent || item.attempts >= item.max_attempts {
NotificationStatus::DeadLetter
} else {
NotificationStatus::RetryScheduled
};
let _ = self.changes.send(id.clone());
Ok(item.status)
}
async fn retry_dead_letter(
&self,
context: &RequestContext,
id: &NotificationId,
) -> Result<OutboxItem, NotifyStoreError> {
let mut items = self
.items
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let item = items
.get_mut(id)
.filter(|item| item.tenant_id == context.tenant_id)
.ok_or(NotifyStoreError::NotFound)?;
if item.status != NotificationStatus::DeadLetter {
return Err(NotifyStoreError::Invalid(
"notification is not dead-lettered".into(),
));
}
item.status = NotificationStatus::Pending;
item.attempts = 0;
let _ = self.changes.send(id.clone());
Ok(item.clone())
}
}
#[derive(Default)]
pub struct CapturingSender {
pub deliveries: Mutex<Vec<(String, String)>>,
}
#[async_trait]
impl Sender for CapturingSender {
fn name(&self) -> &str {
"capture"
}
async fn send(
&self,
context: &DeliveryContext,
recipient: &str,
notification: &Notification,
) -> Result<(), NotifyError> {
if context.cancellation.is_cancelled() {
return Err(NotifyError::Cancelled);
}
self.deliveries
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.push((recipient.to_owned(), notification.to_plain_text()));
Ok(())
}
}
pub struct FailingSender(pub NotifyError);
#[async_trait]
impl Sender for FailingSender {
fn name(&self) -> &str {
"fail"
}
async fn send(
&self,
_context: &DeliveryContext,
_recipient: &str,
_notification: &Notification,
) -> Result<(), NotifyError> {
Err(self.0.clone())
}
}