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, acknowledgement::DeliveryAckRequest},
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    database_id: [u8; 16],
22    pub(super) cancel: CancellationToken,
23    closed: AtomicBool,
24    failure: Mutex<Option<Arc<LocalDeliveryError>>>,
25}
26
27impl<Context: XmtpSharedContext> DeliverySession<Context> {
28    pub(super) fn new(
29        context: Context,
30        owner: Option<DeliveryOwner>,
31        database_id: [u8; 16],
32    ) -> Self {
33        let cancel = context.cancellation_token().child_token();
34        Self {
35            context,
36            owner,
37            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(self
51                .background_error()
52                .unwrap_or(LocalDeliveryError::Closed));
53        }
54        if self.context.db().stream_database_id()? != self.database_id {
55            return Err(StorageError::from(
56                xmtp_db::stream_storage::StreamStorageError::ForeignCursor,
57            )
58            .into());
59        }
60        if let Some(owner) = self.owner {
61            self.context
62                .db()
63                .check_delivery_owner_with_clock(owner, now_ns)?;
64        }
65        Ok(())
66    }
67
68    pub(super) fn owner(&self) -> Option<DeliveryOwner> {
69        self.owner
70    }
71
72    pub(super) fn renew(&self, lease_ns: i64) -> std::result::Result<(), StorageError> {
73        if let Some(owner) = self.owner {
74            self.context
75                .db()
76                .renew_delivery_owner_with_clock(owner, lease_ns, now_ns)?;
77        }
78        Ok(())
79    }
80
81    pub(super) fn close(&self) {
82        if self.closed.swap(true, Ordering::AcqRel) {
83            return;
84        }
85        self.cancel.cancel();
86        if let Some(owner) = self.owner {
87            let mut registered = self.context.delivery_owner().lock();
88            match self.context.db().release_delivery_owner(owner) {
89                Ok(()) => {
90                    if *registered == Some(owner) {
91                        *registered = None;
92                    }
93                }
94                Err(error) => {
95                    if *registered == Some(owner) {
96                        *self.context.retired_delivery_owner().lock() = Some(owner);
97                    }
98                    tracing::warn!(%error, "Failed to release the message delivery owner")
99                }
100            }
101        }
102    }
103
104    pub(super) fn fail(&self, error: LocalDeliveryError) -> LocalDeliveryError {
105        let error = Arc::clone(self.failure.lock().get_or_insert_with(|| Arc::new(error)));
106        self.close();
107        LocalDeliveryError::SessionFailure(error)
108    }
109
110    pub(super) fn background_error(&self) -> Option<LocalDeliveryError> {
111        self.failure
112            .lock()
113            .as_ref()
114            .map(|error| LocalDeliveryError::SessionFailure(Arc::clone(error)))
115    }
116
117    pub(super) fn end_result<T>(&self) -> Result<Option<T>> {
118        match self.background_error() {
119            Some(error) => Err(error),
120            None => Ok(None),
121        }
122    }
123}
124
125#[derive(Clone, Copy)]
126pub(super) enum AcknowledgementState {
127    Waiting,
128    Dispatched,
129    Reselect,
130    Acknowledged,
131    Rejected,
132    Failed,
133}
134
135pub(super) struct PendingAcknowledgement {
136    pub(super) cursor: DeliveryCursor,
137    pub(super) revision: u64,
138    pub(super) state: Mutex<AcknowledgementState>,
139    pub(super) changed: Notify,
140}
141
142impl PendingAcknowledgement {
143    pub(super) fn new(cursor: DeliveryCursor, revision: u64) -> Self {
144        Self {
145            cursor,
146            revision,
147            state: Mutex::new(AcknowledgementState::Waiting),
148            changed: Notify::new(),
149        }
150    }
151}
152
153/// An opaque, single-item token. Dropping it never acknowledges delivery.
154pub struct DeliveryAcknowledgement<Context: XmtpSharedContext> {
155    session: Arc<DeliverySession<Context>>,
156    pending: Arc<PendingAcknowledgement>,
157    group_id: GroupId,
158    control: LocalDeliveryControl,
159    message_id: Vec<u8>,
160}
161
162impl<Context: XmtpSharedContext> DeliveryAcknowledgement<Context> {
163    pub(super) fn new(
164        session: Arc<DeliverySession<Context>>,
165        pending: Arc<PendingAcknowledgement>,
166        group_id: GroupId,
167        control: LocalDeliveryControl,
168        message_id: Vec<u8>,
169    ) -> Self {
170        Self {
171            session,
172            pending,
173            group_id,
174            control,
175            message_id,
176        }
177    }
178
179    /// Call on the host thread immediately before a callback that was queued earlier.
180    /// SelectionChanged discards the queued item. Continue reading without acknowledging it.
181    pub fn check_owner(&self) -> Result<()> {
182        self.check_dispatch(true)
183    }
184
185    /// Keep a cancelled iterator request from admitting a prepared message.
186    pub fn check_owner_cancellable(&self, request: &DeliveryAckRequest) -> Result<bool> {
187        self.check_dispatch_with_request(true, Some(request))
188    }
189
190    fn check_dispatch(&self, dispatch: bool) -> Result<()> {
191        self.check_dispatch_with_request(dispatch, None).map(|_| ())
192    }
193
194    fn check_dispatch_with_request(
195        &self,
196        dispatch: bool,
197        request: Option<&DeliveryAckRequest>,
198    ) -> Result<bool> {
199        if let Some(error) = self.session.background_error() {
200            return Err(error);
201        }
202        let result = self.begin_dispatch(&mut self.pending.state.lock(), dispatch, request);
203        match result {
204            Err(LocalDeliveryError::SelectionChanged) => Err(LocalDeliveryError::SelectionChanged),
205            Err(error) => Err(self.session.fail(error)),
206            result => result,
207        }
208    }
209
210    fn begin_dispatch(
211        &self,
212        state: &mut AcknowledgementState,
213        dispatch: bool,
214        request: Option<&DeliveryAckRequest>,
215    ) -> Result<bool> {
216        match *state {
217            AcknowledgementState::Reselect => return Err(LocalDeliveryError::SelectionChanged),
218            AcknowledgementState::Rejected => {
219                return Err(LocalDeliveryError::AcknowledgementRejected);
220            }
221            AcknowledgementState::Failed => return Err(LocalDeliveryError::AcknowledgementFailed),
222            AcknowledgementState::Dispatched | AcknowledgementState::Acknowledged => {
223                self.session.check_owner()?;
224                return Ok(request.is_none_or(DeliveryAckRequest::admit_handoff));
225            }
226            AcknowledgementState::Waiting => {}
227        }
228        // This lock makes dispatch and a scope change have one order.
229        let selection = self.control.selection.lock();
230        self.session.check_owner()?;
231        let retained = self
232            .session
233            .context
234            .db()
235            .retained_delivery_message(&self.message_id, self.pending.cursor, now_ns())?
236            .is_some();
237        if selection.revision != self.pending.revision
238            || !retained
239            || !super::matches_filter(
240                &self.session.context,
241                self.group_id,
242                &selection.scope,
243                &selection.filter,
244            )?
245        {
246            *state = AcknowledgementState::Reselect;
247            self.pending.changed.notify_one();
248            return Err(LocalDeliveryError::SelectionChanged);
249        }
250        if dispatch {
251            if !request.is_none_or(DeliveryAckRequest::admit_handoff) {
252                return Ok(false);
253            }
254            *state = AcknowledgementState::Dispatched;
255        }
256        Ok(true)
257    }
258
259    /// Persist after the callback returns successfully, or at the next iterator request.
260    // implements: PROC-052, PROC-031, PROC-040
261    pub fn acknowledge(&self) -> Result<()> {
262        self.acknowledge_with_request(None).map(|_| ())
263    }
264
265    /// Keep the prior iterator item eligible if cancellation wins commit admission.
266    pub fn acknowledge_cancellable(&self, request: &DeliveryAckRequest) -> Result<bool> {
267        if request.is_cancelled() {
268            return Ok(false);
269        }
270        self.acknowledge_with_request(Some(request))
271    }
272
273    fn acknowledge_with_request(&self, request: Option<&DeliveryAckRequest>) -> Result<bool> {
274        let mut state = self.pending.state.lock();
275        if matches!(*state, AcknowledgementState::Acknowledged) {
276            return Ok(true);
277        }
278        if let Some(error) = self.session.background_error() {
279            return Err(error);
280        }
281        match *state {
282            AcknowledgementState::Acknowledged => return Ok(true),
283            AcknowledgementState::Rejected => {
284                return Err(LocalDeliveryError::AcknowledgementRejected);
285            }
286            AcknowledgementState::Failed => return Err(LocalDeliveryError::AcknowledgementFailed),
287            AcknowledgementState::Reselect => return Err(LocalDeliveryError::SelectionChanged),
288            AcknowledgementState::Waiting | AcknowledgementState::Dispatched => {}
289        }
290        let result = self.begin_dispatch(&mut state, true, None).and_then(|_| {
291            self.session.check_owner()?;
292            if let Some(owner) = self.session.owner {
293                let db = self.session.context.db();
294                if let Some(request) = request {
295                    return Ok(db.acknowledge_delivery_cancellable_with_clock(
296                        owner,
297                        self.group_id,
298                        self.pending.cursor,
299                        now_ns,
300                        request,
301                    )?);
302                }
303                db.acknowledge_delivery_with_clock(
304                    owner,
305                    self.group_id,
306                    self.pending.cursor,
307                    now_ns,
308                )?;
309            }
310            Ok(request.is_none_or(DeliveryAckRequest::admit_replay_acknowledgement))
311        });
312        match result {
313            Err(LocalDeliveryError::SelectionChanged) => Err(LocalDeliveryError::SelectionChanged),
314            Ok(false) => Ok(false),
315            Ok(true) => {
316                *state = AcknowledgementState::Acknowledged;
317                self.pending.changed.notify_one();
318                Ok(true)
319            }
320            Err(error) => {
321                *state = AcknowledgementState::Failed;
322                self.pending.changed.notify_one();
323                Err(self.session.fail(error))
324            }
325        }
326    }
327
328    /// Read and enrich the queued item without starting dispatch or hiding storage errors.
329    /// The host must still call `check_owner` immediately before its handoff.
330    pub fn enriched_message(&self) -> Result<crate::messages::decoded_message::DecodedMessage> {
331        use crate::messages::enrichment::{EnrichMessageError, enrich_messages};
332        use xmtp_db::group_message::QueryGroupMessage;
333        let result = (|| {
334            self.check_dispatch(false)?;
335            let db = self.session.context.db();
336            let message = db
337                .get_app_visible_group_message(&self.message_id, xmtp_common::time::now_ns())
338                .map_err(StorageError::from)?;
339            let message = match message {
340                Some(message) => message,
341                None => {
342                    self.check_dispatch(false)?;
343                    return Err(LocalDeliveryError::EnrichedMessageUnavailable);
344                }
345            };
346            // Content that fails to decode is delivered as an undecodable
347            // body; only a database failure stops the read.
348            // implements: CTYPE-008
349            let mut messages =
350                enrich_messages(db, &self.group_id, vec![message]).map_err(|error| {
351                    LocalDeliveryError::Storage(match error {
352                        EnrichMessageError::DbConnection(error) => StorageError::Connection(error),
353                        EnrichMessageError::Storage(error) => error,
354                    })
355                })?;
356            messages
357                .pop()
358                .ok_or(LocalDeliveryError::EnrichedMessageUnavailable)
359        })();
360        match result {
361            Err(LocalDeliveryError::SelectionChanged) => Err(LocalDeliveryError::SelectionChanged),
362            Err(error) => Err(self.session.fail(error)),
363            Ok(message) => Ok(message),
364        }
365    }
366
367    /// A failed callback stops this reader and leaves its default delivery position unchanged.
368    pub fn reject(&self) {
369        let mut state = self.pending.state.lock();
370        if matches!(*state, AcknowledgementState::Waiting)
371            && self.control.selection.lock().revision != self.pending.revision
372        {
373            *state = AcknowledgementState::Reselect;
374            self.pending.changed.notify_one();
375        } else if matches!(
376            *state,
377            AcknowledgementState::Waiting | AcknowledgementState::Dispatched
378        ) {
379            *state = AcknowledgementState::Rejected;
380            self.pending.changed.notify_one();
381            self.session.close();
382        }
383    }
384}
385
386impl<Context: XmtpSharedContext> Drop for DeliveryAcknowledgement<Context> {
387    fn drop(&mut self) {
388        self.reject();
389    }
390}
xmtp_mls/subscriptions/local_delivery/acknowledgement.rs

xmtp_mls/subscriptions/local_delivery/
acknowledgement.rs
