1use std::{collections::BTreeMap, pin::pin, sync::Arc, time::Duration};
12
13use async_trait::async_trait;
14use futures::{
15 Stream, StreamExt, TryFutureExt,
16 future::{Either, select},
17 stream::unfold,
18};
19use http::StatusCode;
20use ruma::{
21 EventId, MilliSecondsSinceUnixEpoch, OwnedEventId, OwnedRoomId, OwnedUserId, UserId,
22 api::{
23 client::delayed_events::get_delayed_event::v1::Finalized,
24 error::{ErrorKind, LimitExceededErrorData, RetryAfter, StandardErrorBody},
25 },
26 events::TimelineEventType,
27};
28use serde::{Deserialize, Serialize};
29use tokio::{sync::Notify, time::timeout};
30use tuwunel_core::{
31 Err, Error, Result, debug_warn, err, error, implement,
32 result::NotFound,
33 smallstr::SmallString,
34 utils::{
35 MutexMap, MutexMapGuard,
36 rand::string_array,
37 stream::{BroadbandExt, ReadyExt, TryIgnore, WidebandExt, automatic_width},
38 time::now_millis,
39 },
40 warn,
41};
42use tuwunel_database::{Deserialized, Ignore, Interfix, Json, Map};
43use tuwunel_matrix::{
44 StateKey,
45 pdu::{Content, DELAY_ID, PduBuilder},
46};
47
48use crate::{
49 service::make_name,
50 transaction_ids::{RoomScope, RoomTxn},
51};
52
53#[cfg(test)]
54mod tests;
55
56pub type DelayId = SmallString<[u8; DELAY_ID_LEN]>;
61
62pub struct Service {
67 services: Arc<crate::services::OnceServices>,
68 db: Data,
69
70 locks: MutexMap<DelayId, ()>,
75
76 user_locks: MutexMap<OwnedUserId, ()>,
78
79 wake: Notify,
81}
82
83struct Data {
84 delayid_event: Arc<Map>,
85 delayid_finalised: Arc<Map>,
86 dueat_delayid: Arc<Map>,
87 userid_delayid: Arc<Map>,
88}
89
90#[derive(Debug, Deserialize, Serialize)]
95pub struct Event {
96 pub user_id: OwnedUserId,
98
99 pub room_id: OwnedRoomId,
101
102 #[serde(rename = "type")]
104 pub event_type: TimelineEventType,
105
106 #[serde(default, skip_serializing_if = "Option::is_none")]
108 pub state_key: Option<StateKey>,
109
110 pub content: Content,
112
113 #[serde(default, skip_serializing_if = "Option::is_none")]
115 pub redacts: Option<OwnedEventId>,
116
117 #[serde(default, skip_serializing_if = "Option::is_none")]
119 pub timestamp: Option<MilliSecondsSinceUnixEpoch>,
120
121 pub delay_ms: u64,
123
124 pub due_ms: u64,
126
127 #[serde(default, skip_serializing_if = "Option::is_none")]
129 pub finalised: Option<Finalized>,
130}
131
132#[derive(Clone, Copy, Debug)]
137pub struct Limits {
138 pub max_delay_ms: u64,
140
141 pub max_scheduled: usize,
143}
144
145const DELAY_ID_LEN: usize = 32;
147
148const SEND_WIDTH: usize = 8;
150
151const MAX_WAIT: Duration = Duration::from_mins(1);
152
153#[async_trait]
154impl crate::Service for Service {
155 fn build(args: &crate::Args<'_>) -> Result<Arc<Self>> {
156 let db = Data {
157 delayid_event: args.db["delayid_event"].clone(),
158 delayid_finalised: args.db["delayid_finalised"].clone(),
159 dueat_delayid: args.db["dueat_delayid"].clone(),
160 userid_delayid: args.db["userid_delayid"].clone(),
161 };
162
163 Ok(Arc::new(Self {
164 services: args.services.clone(),
165 db,
166 locks: MutexMap::new(),
167 user_locks: MutexMap::new(),
168 wake: Notify::new(),
169 }))
170 }
171
172 async fn worker(self: Arc<Self>) -> Result {
173 if self.services.db.is_read_only() {
174 return Ok(()); }
176
177 self.due_events()
178 .for_each_concurrent(SEND_WIDTH, async |(due, delay_id)| {
179 self.send_scheduled(&delay_id, due).await;
180 })
181 .await;
182
183 Ok(())
184 }
185
186 fn name(&self) -> &str { make_name(module_path!()) }
187}
188
189#[implement(Service)]
193#[must_use]
194pub fn limits(&self) -> Option<Limits> {
195 let config = &self.services.config.rooms.delayed_events;
196 let max_delay_ms = config.max_delay_s.saturating_mul(1000);
197 let enabled = max_delay_ms > 0 && config.max_scheduled > 0;
198
199 enabled.then_some(Limits {
200 max_delay_ms,
201 max_scheduled: config.max_scheduled,
202 })
203}
204
205#[implement(Service)]
212pub async fn check_capacity(&self, user_id: &UserId, max_scheduled: usize) -> Result {
213 let scheduled = self
214 .db
215 .userid_delayid
216 .count_prefix(&(user_id, Interfix))
217 .await;
218
219 if scheduled < max_scheduled || self.reclaim_orphans(user_id).await < max_scheduled {
220 return Ok(());
221 }
222
223 let retry_after = self
224 .next_due_of(user_id)
225 .await
226 .map(|due| due.saturating_sub(now_millis()).div_ceil(1000))
227 .map(Duration::from_secs);
228
229 Err(limit_exceeded(retry_after))
230}
231
232#[implement(Service)]
237pub async fn existing_txn(&self, txn: RoomTxn<'_>, event_type: &str) -> Result<Option<DelayId>> {
238 self.services
239 .transaction_ids
240 .existing_room_txnid(txn, RoomScope::Delayed(event_type))
241 .await
242 .optional()?
243 .map(|response| {
244 str::from_utf8(&response)
245 .map(DelayId::from)
246 .map_err(|e| err!(Database("Unreadable delayed event transaction response: {e}")))
247 })
248 .transpose()
249}
250
251#[implement(Service)]
257pub fn schedule(&self, event: &Event, txn: Option<RoomTxn<'_>>) -> DelayId {
258 let delay_id = DelayId::from(string_array::<DELAY_ID_LEN>().as_str());
259 let event_type = event.event_type.to_cow_str();
260 let response = txn.map(|txn| (txn, RoomScope::Delayed(&event_type)));
261
262 self.insert(&delay_id, event, response);
263 self.wake.notify_one();
264
265 delay_id
266}
267
268#[implement(Service)]
275pub async fn get(&self, delay_id: &str, owner: Option<&UserId>) -> Result<Event> {
276 self.get_event(delay_id)
277 .or_else(|_| read_event(&self.db.delayid_finalised, delay_id))
278 .await
279 .ok()
280 .filter(|event| owner.is_none_or(|owner| owner == event.user_id))
281 .ok_or_else(|| err!(Request(NotFound("No delayed event with this ID."))))
282}
283
284#[implement(Service)]
288pub fn scheduled<'a>(
289 &'a self,
290 user_id: &'a UserId,
291) -> impl Stream<Item = (DelayId, Event)> + Send + 'a {
292 self.owned_ids(user_id)
293 .wide_filter_map(async |delay_id| {
294 let event = self.get_event(&delay_id).await.ok()?;
295
296 Some((delay_id, event))
297 })
298}
299
300#[implement(Service)]
305#[tracing::instrument(level = "debug", skip(self))]
306pub async fn restart(&self, delay_id: &str, owner: Option<&UserId>) -> Result {
307 let _lock = self.locks.lock(&DelayId::from(delay_id)).await;
308 let event = self.get(delay_id, owner).await?;
309
310 if let Some(finalised) = &event.finalised {
311 return Err(conflict(finalised));
312 }
313
314 let old_due = event.due_ms;
315 let event = Event {
316 due_ms: now_millis().saturating_add(event.delay_ms),
317 ..event
318 };
319
320 self.reschedule(delay_id, &event, old_due);
321
322 Ok(())
323}
324
325#[implement(Service)]
330#[tracing::instrument(level = "debug", skip(self))]
331pub async fn cancel(&self, delay_id: &str, owner: Option<&UserId>) -> Result {
332 let _lock = self.locks.lock(&DelayId::from(delay_id)).await;
333 let event = self.get(delay_id, owner).await?;
334
335 match event.finalised {
336 | Some(Finalized::Error { .. } | Finalized::Canceled { .. }) => Ok(()),
337 | Some(ref finalised @ Finalized::Sent { .. }) => Err(conflict(finalised)),
338 | None => {
339 let finalized_ts = MilliSecondsSinceUnixEpoch::now();
340 let finalised = Some(Finalized::Canceled { finalized_ts });
341
342 self.finalise(delay_id, &Event { finalised, ..event });
343
344 Ok(())
345 },
346 }
347}
348
349#[implement(Service)]
356#[tracing::instrument(level = "debug", skip(self))]
357pub async fn send(&self, delay_id: &str, owner: Option<&UserId>) -> Result<OwnedEventId> {
358 let _lock = self.locks.lock(&DelayId::from(delay_id)).await;
359 let event = self.get(delay_id, owner).await?;
360
361 match event.finalised {
362 | None => self.send_and_finalise(delay_id, event).await,
363 | Some(Finalized::Sent { event_id, .. }) => Ok(event_id),
364 | Some(ref finalised @ (Finalized::Error { .. } | Finalized::Canceled { .. })) =>
365 Err(conflict(finalised)),
366 }
367}
368
369#[implement(Service)]
376#[tracing::instrument(level = "debug", skip(self))]
377pub async fn remove_user(&self, user_id: &UserId) {
378 let _user_lock = self.lock_user(user_id).await;
379
380 self.owned_ids(user_id)
381 .for_each_concurrent(automatic_width(), async |delay_id| {
382 self.remove_owned(user_id, &delay_id).await;
383 })
384 .await;
385}
386
387#[implement(Service)]
392pub async fn lock_user(&self, user_id: &UserId) -> MutexMapGuard<OwnedUserId, ()> {
393 self.user_locks.lock(user_id).await
394}
395
396#[implement(Service)]
403fn due_events(&self) -> impl Stream<Item = (u64, DelayId)> + Send + '_ {
404 unfold(None, async move |last: Option<u64>| {
405 if let Some(last) = last {
406 self.wait_for_due(last).await?;
407 }
408
409 let now = now_millis();
410
411 Some((self.due_at(now), Some(now)))
412 })
413 .flatten()
414}
415
416#[implement(Service)]
421async fn wait_for_due(&self, last: u64) -> Option<()> {
422 let wait = self
423 .next_due_after(last)
424 .await
425 .map_or(MAX_WAIT, |due| Duration::from_millis(due.saturating_sub(now_millis())))
426 .min(MAX_WAIT);
427
428 let wake = timeout(wait, self.wake.notified());
429 let shutdown = self.services.server.until_shutdown();
430
431 match select(pin!(wake), pin!(shutdown)).await {
432 | Either::Left(_) => Some(()),
433 | Either::Right(_) => None,
434 }
435}
436
437#[implement(Service)]
442fn due_at(&self, now: u64) -> impl Stream<Item = (u64, DelayId)> + Send + '_ {
443 self.db
444 .dueat_delayid
445 .keys()
446 .ready_take_while(move |key: &Result<(u64, &str)>| {
447 key.as_ref().is_ok_and(|&(due, _)| due <= now) || key.is_err()
448 })
449 .ready_filter_map(|key| {
450 key.inspect_err(|e| error!(%e, "Unreadable delayed event index entry"))
451 .ok()
452 })
453 .map(|(due, delay_id)| (due, DelayId::from(delay_id)))
454 .ready_filter(|(_, delay_id)| !self.locks.contains(delay_id))
455}
456
457#[implement(Service)]
464#[tracing::instrument(level = "debug", skip(self))]
465async fn send_scheduled(&self, delay_id: &DelayId, due: u64) {
466 let _lock = self.locks.lock(delay_id).await;
467 let event = match self.get_event(delay_id).await {
468 | Ok(event) if event.due_ms == due => event,
469 | Ok(_) => return,
470 | Err(e) if e.is_not_found() => return self.drop_due(delay_id, due),
471 | Err(e) => {
472 error!(%delay_id, %e, "Dropping unreadable delayed event");
473 return self.drop_due(delay_id, due);
474 },
475 };
476
477 let user_id = event.user_id.clone();
478
479 self.send_and_finalise(delay_id, event)
480 .await
481 .inspect_err(|e| log_send_error(delay_id, &user_id, e, None))
482 .ok();
483}
484
485#[implement(Service)]
490async fn send_and_finalise(&self, delay_id: &str, event: Event) -> Result<OwnedEventId> {
491 let result = self.append(delay_id, &event).await;
492 let finalised = Some(outcome(&result));
493
494 self.finalise(delay_id, &Event { finalised, ..event });
495
496 result
497}
498
499#[implement(Service)]
506#[tracing::instrument(level = "debug", skip_all)]
507async fn append(&self, delay_id: &str, event: &Event) -> Result<OwnedEventId> {
508 let Event { user_id, room_id, .. } = event;
509
510 if self.services.users.is_suspended(user_id).await {
511 return Err!(Request(UserSuspended("Cannot send delayed events while suspended.")));
512 }
513
514 let pdu = PduBuilder {
515 event_type: event.event_type.clone(),
516 content: event.content.clone(),
517 unsigned: Some(BTreeMap::from([(DELAY_ID.to_owned(), delay_id.into())])),
518 state_key: event.state_key.clone(),
519 redacts: event.redacts.clone(),
520 timestamp: event.timestamp,
521 };
522
523 let state_lock = self.services.state.mutex.lock(room_id).await;
524 let (event_id, appended) = self
525 .services
526 .timeline
527 .build_and_store_pdu(pdu, user_id, room_id, &state_lock)
528 .await?;
529
530 appended
531 .inspect_err(|e| log_send_error(delay_id, user_id, e, Some(&event_id)))
532 .ok();
533
534 Ok(event_id)
535}
536
537fn log_send_error(delay_id: &str, user_id: &UserId, error: &Error, event_id: Option<&EventId>) {
541 match error.status_code().is_server_error() {
542 | true => warn!(%delay_id, %user_id, %error, ?event_id, "Delayed event send failed"),
543 | false =>
544 debug_warn!(%delay_id, %user_id, %error, ?event_id, "Delayed event send failed"),
545 }
546}
547
548#[implement(Service)]
553#[tracing::instrument(level = "debug", skip(self))]
554async fn remove_owned(&self, user_id: &UserId, delay_id: &DelayId) {
555 let _lock = self.locks.lock(delay_id).await;
556
557 match self.get_event(delay_id).await {
558 | Ok(event) => self.remove(delay_id, &event.user_id, event.due_ms),
559 | Err(_) => self.drop_owned(user_id, delay_id),
560 }
561}
562
563#[implement(Service)]
569#[tracing::instrument(level = "debug", skip(self))]
570async fn reclaim_orphans(&self, user_id: &UserId) -> usize {
571 self.owned_ids(user_id)
572 .broad_filter_map(async |delay_id| {
573 let _lock = self.locks.lock(&delay_id).await;
574
575 self.get_event(&delay_id)
576 .await
577 .inspect_err(|_| self.drop_owned(user_id, &delay_id))
578 .ok()
579 })
580 .count()
581 .await
582}
583
584#[implement(Service)]
589#[tracing::instrument(level = "trace", skip(self))]
590async fn next_due_after(&self, now: u64) -> Option<u64> {
591 pin!(
592 self.db
593 .dueat_delayid
594 .keys_from(&(now.saturating_add(1),))
595 )
596 .next()
597 .await
598 .and_then(|key| {
599 key.inspect_err(|e| error!(%e, "Unreadable delayed event index entry"))
600 .ok()
601 })
602 .map(|(due, _): (u64, Ignore)| due)
603}
604
605#[implement(Service)]
610#[tracing::instrument(level = "trace", skip(self))]
611async fn next_due_of(&self, user_id: &UserId) -> Option<u64> {
612 self.owned_ids(user_id)
613 .broad_filter_map(async |delay_id| self.get_event(&delay_id).await.ok())
614 .ready_fold(None, |next: Option<u64>, event| {
615 Some(next.map_or(event.due_ms, |next| next.min(event.due_ms)))
616 })
617 .await
618}
619
620#[implement(Service)]
621fn owned_ids<'a>(&'a self, user_id: &'a UserId) -> impl Stream<Item = DelayId> + Send + 'a {
622 self.db
623 .userid_delayid
624 .keys_prefix(&(user_id, Interfix))
625 .ignore_err()
626 .map(|(_, delay_id): (Ignore, &str)| DelayId::from(delay_id))
627}
628
629#[implement(Service)]
630#[tracing::instrument(level = "trace", skip(self))]
631async fn get_event(&self, delay_id: &str) -> Result<Event> {
632 read_event(&self.db.delayid_event, delay_id).await
633}
634
635async fn read_event(map: &Arc<Map>, delay_id: &str) -> Result<Event> {
636 map.get(delay_id)
637 .await
638 .deserialized()
639 .map(|Json(event): Json<Event>| event)
640}
641
642#[implement(Service)]
647fn insert(&self, delay_id: &str, event: &Event, response: Option<(RoomTxn<'_>, RoomScope<'_>)>) {
648 let mut batch = self.services.db.txn();
649
650 batch.raw_put(&self.db.delayid_event, delay_id, Json(event));
651 batch.put_raw(&self.db.dueat_delayid, (event.due_ms, delay_id), []);
652 batch.put_raw(&self.db.userid_delayid, (&event.user_id, delay_id), []);
653
654 if let Some((txn, scope)) = response {
655 self.services
656 .transaction_ids
657 .add_room_txnid_batch(&mut batch, txn, scope, delay_id.as_bytes());
658 }
659
660 batch.execute();
661}
662
663#[implement(Service)]
664fn reschedule(&self, delay_id: &str, event: &Event, old_due: u64) {
665 let mut batch = self.services.db.txn();
666
667 batch.raw_put(&self.db.delayid_event, delay_id, Json(event));
668 batch.del(&self.db.dueat_delayid, (old_due, delay_id));
669 batch.put_raw(&self.db.dueat_delayid, (event.due_ms, delay_id), []);
670 batch.execute();
671}
672
673#[implement(Service)]
678fn finalise(&self, delay_id: &str, event: &Event) {
679 let mut batch = self.services.db.txn();
680
681 batch.raw_put(&self.db.delayid_finalised, delay_id, Json(event));
682 batch.del_raw(&self.db.delayid_event, delay_id);
683 batch.del(&self.db.dueat_delayid, (event.due_ms, delay_id));
684 batch.del(&self.db.userid_delayid, (&event.user_id, delay_id));
685 batch.execute();
686}
687
688#[implement(Service)]
689fn remove(&self, delay_id: &str, user_id: &UserId, due: u64) {
690 let mut batch = self.services.db.txn();
691
692 batch.del_raw(&self.db.delayid_event, delay_id);
693 batch.del(&self.db.dueat_delayid, (due, delay_id));
694 batch.del(&self.db.userid_delayid, (user_id, delay_id));
695 batch.execute();
696}
697
698#[implement(Service)]
699fn drop_due(&self, delay_id: &str, due: u64) {
700 let mut batch = self.services.db.txn();
701
702 batch.del_raw(&self.db.delayid_event, delay_id);
703 batch.del(&self.db.dueat_delayid, (due, delay_id));
704 batch.execute();
705}
706
707#[implement(Service)]
708fn drop_owned(&self, user_id: &UserId, delay_id: &str) {
709 let mut batch = self.services.db.txn();
710
711 batch.del_raw(&self.db.delayid_event, delay_id);
712 batch.del(&self.db.userid_delayid, (user_id, delay_id));
713 batch.execute();
714}
715
716fn limit_exceeded(retry_after: Option<Duration>) -> Error {
717 let retry_after = retry_after.map(RetryAfter::Delay);
718 let kind = ErrorKind::LimitExceeded(LimitExceededErrorData { retry_after });
719 let message = "The maximum number of delayed events has been reached.";
720
721 Error::Request(kind, message.into(), StatusCode::TOO_MANY_REQUESTS)
722}
723
724fn conflict(finalised: &Finalized) -> Error {
725 match finalised {
726 | Finalized::Sent { .. } => err!(Conflict("The delayed event has already been sent.")),
727 | Finalized::Error { .. } | Finalized::Canceled { .. } =>
728 err!(Conflict("The delayed event has already been cancelled.")),
729 }
730}
731
732fn outcome(result: &Result<OwnedEventId>) -> Finalized {
736 let finalized_ts = MilliSecondsSinceUnixEpoch::now();
737
738 match result {
739 | Ok(event_id) => Finalized::Sent { finalized_ts, event_id: event_id.clone() },
740 | Err(e) => Finalized::Error {
741 finalized_ts,
742 error: StandardErrorBody::new(e.kind(), e.sanitized_message()),
743 },
744 }
745}