1use std::collections::BTreeMap;
4use std::sync::Mutex;
5use std::time::Duration;
6
7use af_context::{NotificationId, RequestContext};
8use async_trait::async_trait;
9use chrono::Utc;
10
11use crate::{
12 AttemptEvent, DeliveryClass, DeliveryContext, DurableOutbox, NewOutboxItem, Notification,
13 NotificationStatus, NotifyError, NotifyStoreError, OutboxItem, Sender,
14};
15
16pub struct MemoryOutbox {
18 items: Mutex<BTreeMap<NotificationId, OutboxItem>>,
19 idempotency: Mutex<BTreeMap<(String, String, String), NotificationId>>,
20 changes: tokio::sync::broadcast::Sender<NotificationId>,
21}
22
23impl Default for MemoryOutbox {
24 fn default() -> Self {
25 let (changes, _) = tokio::sync::broadcast::channel(32);
26 Self {
27 items: Mutex::new(BTreeMap::new()),
28 idempotency: Mutex::new(BTreeMap::new()),
29 changes,
30 }
31 }
32}
33
34#[async_trait]
35impl DurableOutbox for MemoryOutbox {
36 fn subscribe(&self) -> tokio::sync::broadcast::Receiver<NotificationId> {
37 self.changes.subscribe()
38 }
39 async fn enqueue(&self, item: NewOutboxItem) -> Result<OutboxItem, NotifyStoreError> {
40 let key = (
41 item.tenant_id.to_string(),
42 item.subject_id.to_string(),
43 item.idempotency_key.clone(),
44 );
45 let mut idempotency = self
46 .idempotency
47 .lock()
48 .unwrap_or_else(std::sync::PoisonError::into_inner);
49 let mut items = self
50 .items
51 .lock()
52 .unwrap_or_else(std::sync::PoisonError::into_inner);
53 if let Some(existing) = idempotency.get(&key).and_then(|id| items.get(id)) {
54 return if existing.channel == item.channel
55 && existing.recipient == item.recipient
56 && existing.notification == item.notification
57 && existing.max_attempts == item.max_attempts.max(1)
58 {
59 Ok(existing.clone())
60 } else {
61 Err(NotifyStoreError::IdempotencyConflict)
62 };
63 }
64 let id = NotificationId::parse(format!("notification-{}", items.len() + 1))
65 .map_err(|error| NotifyStoreError::Invalid(error.to_string()))?;
66 let now = Utc::now();
67 let record = OutboxItem {
68 id: id.clone(),
69 tenant_id: item.tenant_id,
70 subject_id: item.subject_id,
71 channel: item.channel,
72 recipient: item.recipient,
73 notification: item.notification,
74 status: NotificationStatus::Pending,
75 attempts: 0,
76 max_attempts: item.max_attempts.max(1),
77 lease_version: 0,
78 last_error: None,
79 created_at: now,
80 updated_at: now,
81 };
82 items.insert(id.clone(), record.clone());
83 idempotency.insert(key, id);
84 let _ = self.changes.send(record.id.clone());
85 Ok(record)
86 }
87 async fn get(
88 &self,
89 context: &RequestContext,
90 id: &NotificationId,
91 ) -> Result<OutboxItem, NotifyStoreError> {
92 self.items
93 .lock()
94 .unwrap_or_else(std::sync::PoisonError::into_inner)
95 .get(id)
96 .filter(|item| {
97 item.tenant_id == context.tenant_id && item.subject_id == context.subject_id
98 })
99 .cloned()
100 .ok_or(NotifyStoreError::NotFound)
101 }
102 async fn list(
103 &self,
104 context: &RequestContext,
105 status: Option<NotificationStatus>,
106 limit: usize,
107 after: Option<&NotificationId>,
108 ) -> Result<Vec<OutboxItem>, NotifyStoreError> {
109 Ok(self
110 .items
111 .lock()
112 .unwrap_or_else(std::sync::PoisonError::into_inner)
113 .values()
114 .filter(|item| {
115 item.tenant_id == context.tenant_id
116 && item.subject_id == context.subject_id
117 && status.is_none_or(|status| item.status == status)
118 && after.is_none_or(|after| item.id > *after)
119 })
120 .take(limit.clamp(1, 101))
121 .cloned()
122 .collect())
123 }
124 async fn attempts(
125 &self,
126 context: &RequestContext,
127 id: &NotificationId,
128 ) -> Result<Vec<AttemptEvent>, NotifyStoreError> {
129 self.get(context, id).await.map(|_| Vec::new())
130 }
131 async fn claim(
132 &self,
133 _worker_id: &str,
134 _lease_secs: i64,
135 batch: usize,
136 ) -> Result<Vec<OutboxItem>, NotifyStoreError> {
137 let mut items = self
138 .items
139 .lock()
140 .unwrap_or_else(std::sync::PoisonError::into_inner);
141 let ids: Vec<_> = items
142 .values()
143 .filter(|item| {
144 matches!(
145 item.status,
146 NotificationStatus::Pending | NotificationStatus::RetryScheduled
147 )
148 })
149 .take(batch.clamp(1, 100))
150 .map(|item| item.id.clone())
151 .collect();
152 let claimed: Vec<_> = ids
153 .into_iter()
154 .filter_map(|id| {
155 let item = items.get_mut(&id)?;
156 item.status = NotificationStatus::Sending;
157 item.attempts += 1;
158 item.lease_version += 1;
159 Some(item.clone())
160 })
161 .collect();
162 for item in &claimed {
163 let _ = self.changes.send(item.id.clone());
164 }
165 Ok(claimed)
166 }
167 async fn mark_sent(
168 &self,
169 id: &NotificationId,
170 lease_version: i64,
171 ) -> Result<(), NotifyStoreError> {
172 let mut items = self
173 .items
174 .lock()
175 .unwrap_or_else(std::sync::PoisonError::into_inner);
176 let item = items.get_mut(id).ok_or(NotifyStoreError::NotFound)?;
177 if item.lease_version != lease_version {
178 return Err(NotifyStoreError::LeaseLost);
179 }
180 item.status = NotificationStatus::Sent;
181 let _ = self.changes.send(id.clone());
182 Ok(())
183 }
184 async fn record_failure(
185 &self,
186 id: &NotificationId,
187 lease_version: i64,
188 class: DeliveryClass,
189 _delay: Duration,
190 ) -> Result<NotificationStatus, NotifyStoreError> {
191 let mut items = self
192 .items
193 .lock()
194 .unwrap_or_else(std::sync::PoisonError::into_inner);
195 let item = items.get_mut(id).ok_or(NotifyStoreError::NotFound)?;
196 if item.lease_version != lease_version {
197 return Err(NotifyStoreError::LeaseLost);
198 }
199 item.status = if class == DeliveryClass::Permanent || item.attempts >= item.max_attempts {
200 NotificationStatus::DeadLetter
201 } else {
202 NotificationStatus::RetryScheduled
203 };
204 let _ = self.changes.send(id.clone());
205 Ok(item.status)
206 }
207 async fn retry_dead_letter(
208 &self,
209 context: &RequestContext,
210 id: &NotificationId,
211 ) -> Result<OutboxItem, NotifyStoreError> {
212 let mut items = self
213 .items
214 .lock()
215 .unwrap_or_else(std::sync::PoisonError::into_inner);
216 let item = items
217 .get_mut(id)
218 .filter(|item| item.tenant_id == context.tenant_id)
219 .ok_or(NotifyStoreError::NotFound)?;
220 if item.status != NotificationStatus::DeadLetter {
221 return Err(NotifyStoreError::Invalid(
222 "notification is not dead-lettered".into(),
223 ));
224 }
225 item.status = NotificationStatus::Pending;
226 item.attempts = 0;
227 let _ = self.changes.send(id.clone());
228 Ok(item.clone())
229 }
230}
231
232#[derive(Default)]
234pub struct CapturingSender {
235 pub deliveries: Mutex<Vec<(String, String)>>,
237}
238
239#[async_trait]
240impl Sender for CapturingSender {
241 fn name(&self) -> &str {
242 "capture"
243 }
244 async fn send(
245 &self,
246 context: &DeliveryContext,
247 recipient: &str,
248 notification: &Notification,
249 ) -> Result<(), NotifyError> {
250 if context.cancellation.is_cancelled() {
251 return Err(NotifyError::Cancelled);
252 }
253 self.deliveries
254 .lock()
255 .unwrap_or_else(std::sync::PoisonError::into_inner)
256 .push((recipient.to_owned(), notification.to_plain_text()));
257 Ok(())
258 }
259}
260
261pub struct FailingSender(pub NotifyError);
263
264#[async_trait]
265impl Sender for FailingSender {
266 fn name(&self) -> &str {
267 "fail"
268 }
269 async fn send(
270 &self,
271 _context: &DeliveryContext,
272 _recipient: &str,
273 _notification: &Notification,
274 ) -> Result<(), NotifyError> {
275 Err(self.0.clone())
276 }
277}