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

xmtp_mls/subscriptions/local_delivery/
acknowledgement.rs
