xmtp_mls/subscriptions/local_delivery/
acknowledgement.rs

1use parking_lot::Mutex;
2use std::sync::{
3    Arc,
4    atomic::{AtomicBool, Ordering},
5};
6use tokio::sync::Notify;
7use tokio_util::sync::CancellationToken;
8use xmtp_common::time::now_ns;
9use xmtp_db::{
10    StorageError,
11    delivery::{DeliveryCursor, DeliveryOwner, QueryDelivery},
12};
13use xmtp_proto::types::GroupId;
14
15use super::{LocalDeliveryControl, LocalDeliveryError, Result};
16use crate::context::XmtpSharedContext;
17
18pub(super) struct DeliverySession<Context: XmtpSharedContext> {
19    pub(super) context: Context,
20    pub(super) owner: Option<DeliveryOwner>,
21    replay_database_id: Option<[u8; 16]>,
22    pub(super) cancel: CancellationToken,
23    closed: AtomicBool,
24    failure: Mutex<Option<LocalDeliveryError>>,
25}
26
27impl<Context: XmtpSharedContext> DeliverySession<Context> {
28    pub(super) fn new(
29        context: Context,
30        owner: Option<DeliveryOwner>,
31        replay_database_id: Option<[u8; 16]>,
32    ) -> Self {
33        let cancel = context.cancellation_token().child_token();
34        Self {
35            context,
36            owner,
37            replay_database_id,
38            cancel,
39            closed: AtomicBool::new(false),
40            failure: Mutex::new(None),
41        }
42    }
43
44    pub(super) fn is_closed(&self) -> bool {
45        self.closed.load(Ordering::Acquire) || self.cancel.is_cancelled()
46    }
47
48    pub(super) fn check_owner(&self) -> Result<()> {
49        if self.is_closed() {
50            return Err(LocalDeliveryError::Closed);
51        }
52        if let Some(owner) = self.owner {
53            self.context
54                .db()
55                .check_delivery_owner_with_clock(owner, now_ns)?;
56        }
57        if let Some(identity) = self.replay_database_id
58            && self.context.db().stream_database_id()? != identity
59        {
60            return Err(StorageError::from(
61                xmtp_db::stream_storage::StreamStorageError::ForeignCursor,
62            )
63            .into());
64        }
65        Ok(())
66    }
67
68    pub(super) fn renew(&self, lease_ns: i64) -> std::result::Result<(), StorageError> {
69        if let Some(owner) = self.owner {
70            self.context
71                .db()
72                .renew_delivery_owner_with_clock(owner, lease_ns, now_ns)?;
73        }
74        Ok(())
75    }
76
77    pub(super) fn close(&self) {
78        if self.closed.swap(true, Ordering::AcqRel) {
79            return;
80        }
81        self.cancel.cancel();
82        if let Some(owner) = self.owner {
83            let mut registered = self.context.delivery_owner().lock();
84            match self.context.db().release_delivery_owner(owner) {
85                Ok(()) => {
86                    if *registered == Some(owner) {
87                        *registered = None;
88                    }
89                }
90                Err(error) => {
91                    tracing::warn!(%error, "Failed to release the message delivery owner")
92                }
93            }
94        }
95    }
96
97    pub(super) fn fail(&self, error: LocalDeliveryError) {
98        *self.failure.lock() = Some(error);
99        self.close();
100    }
101
102    pub(super) fn end_result<T>(&self) -> Result<Option<T>> {
103        match self.failure.lock().take() {
104            Some(error) => Err(error),
105            None => Ok(None),
106        }
107    }
108}
109
110#[derive(Clone, Copy)]
111pub(super) enum AcknowledgementState {
112    Waiting,
113    Dispatched,
114    Reselect,
115    Acknowledged,
116    Rejected,
117    Failed,
118}
119
120pub(super) struct PendingAcknowledgement {
121    pub(super) cursor: DeliveryCursor,
122    pub(super) revision: u64,
123    pub(super) state: Mutex<AcknowledgementState>,
124    pub(super) changed: Notify,
125}
126
127impl PendingAcknowledgement {
128    pub(super) fn new(cursor: DeliveryCursor, revision: u64) -> Self {
129        Self {
130            cursor,
131            revision,
132            state: Mutex::new(AcknowledgementState::Waiting),
133            changed: Notify::new(),
134        }
135    }
136}
137
138/// An opaque, single-item token. Dropping it never acknowledges delivery.
139pub struct DeliveryAcknowledgement<Context: XmtpSharedContext> {
140    session: Arc<DeliverySession<Context>>,
141    pending: Arc<PendingAcknowledgement>,
142    group_id: GroupId,
143    control: LocalDeliveryControl,
144    message_id: Vec<u8>,
145}
146
147impl<Context: XmtpSharedContext> DeliveryAcknowledgement<Context> {
148    pub(super) fn new(
149        session: Arc<DeliverySession<Context>>,
150        pending: Arc<PendingAcknowledgement>,
151        group_id: GroupId,
152        control: LocalDeliveryControl,
153        message_id: Vec<u8>,
154    ) -> Self {
155        Self {
156            session,
157            pending,
158            group_id,
159            control,
160            message_id,
161        }
162    }
163
164    /// Call on the host thread immediately before a callback that was queued earlier.
165    /// SelectionChanged discards the queued item. Continue reading without acknowledging it.
166    pub fn check_owner(&self) -> Result<()> {
167        self.begin_dispatch(&mut self.pending.state.lock())
168    }
169
170    fn begin_dispatch(&self, state: &mut AcknowledgementState) -> Result<()> {
171        match *state {
172            AcknowledgementState::Reselect => return Err(LocalDeliveryError::SelectionChanged),
173            AcknowledgementState::Rejected => {
174                return Err(LocalDeliveryError::AcknowledgementRejected);
175            }
176            AcknowledgementState::Failed => return Err(LocalDeliveryError::AcknowledgementFailed),
177            AcknowledgementState::Dispatched | AcknowledgementState::Acknowledged => {
178                return self.session.check_owner();
179            }
180            AcknowledgementState::Waiting => {}
181        }
182        // This lock makes dispatch and a scope change have one order.
183        let selection = self.control.selection.lock();
184        self.session.check_owner()?;
185        let retained = self.session.context.db().delivery_message_is_retained(
186            &self.message_id,
187            self.pending.cursor,
188            now_ns(),
189        )?;
190        if selection.revision != self.pending.revision
191            || !retained
192            || !super::matches_filter(&self.session.context, self.group_id, &selection.filter)?
193        {
194            *state = AcknowledgementState::Reselect;
195            self.pending.changed.notify_one();
196            return Err(LocalDeliveryError::SelectionChanged);
197        }
198        *state = AcknowledgementState::Dispatched;
199        Ok(())
200    }
201
202    /// Persist after the callback returns successfully, or at the next iterator request.
203    pub fn acknowledge(&self) -> Result<()> {
204        let mut state = self.pending.state.lock();
205        match *state {
206            AcknowledgementState::Acknowledged => return Ok(()),
207            AcknowledgementState::Rejected => {
208                return Err(LocalDeliveryError::AcknowledgementRejected);
209            }
210            AcknowledgementState::Failed => return Err(LocalDeliveryError::AcknowledgementFailed),
211            AcknowledgementState::Reselect => return Err(LocalDeliveryError::SelectionChanged),
212            AcknowledgementState::Waiting | AcknowledgementState::Dispatched => {}
213        }
214        let result = self.begin_dispatch(&mut state).and_then(|()| {
215            self.session.check_owner()?;
216            if let Some(owner) = self.session.owner {
217                self.session.context.db().acknowledge_delivery_with_clock(
218                    owner,
219                    self.group_id,
220                    self.pending.cursor,
221                    now_ns,
222                )?;
223            }
224            Ok(())
225        });
226        match result {
227            Err(LocalDeliveryError::SelectionChanged) => Err(LocalDeliveryError::SelectionChanged),
228            Ok(()) => {
229                *state = AcknowledgementState::Acknowledged;
230                self.pending.changed.notify_one();
231                Ok(())
232            }
233            Err(error) => {
234                *state = AcknowledgementState::Failed;
235                self.pending.changed.notify_one();
236                self.session.close();
237                Err(error)
238            }
239        }
240    }
241
242    /// A failed callback stops this reader and leaves its default delivery position unchanged.
243    pub fn reject(&self) {
244        let mut state = self.pending.state.lock();
245        if matches!(*state, AcknowledgementState::Waiting)
246            && self.control.selection.lock().revision != self.pending.revision
247        {
248            *state = AcknowledgementState::Reselect;
249            self.pending.changed.notify_one();
250        } else if matches!(
251            *state,
252            AcknowledgementState::Waiting | AcknowledgementState::Dispatched
253        ) {
254            *state = AcknowledgementState::Rejected;
255            self.pending.changed.notify_one();
256            self.session.close();
257        }
258    }
259}
260
261impl<Context: XmtpSharedContext> Drop for DeliveryAcknowledgement<Context> {
262    fn drop(&mut self) {
263        self.reject();
264    }
265}
xmtp_mls/subscriptions/local_delivery/acknowledgement.rs

xmtp_mls/subscriptions/local_delivery/
acknowledgement.rs
