futu_mcp/state/push_subscribers/

modern_delivery.rs

use super::*;

impl ModernResourceDelivery {
    pub(super) fn new(handle: String, permit: Arc<OwnedSemaphorePermit>) -> Self {
        Self {
            handle,
            inner: StdMutex::new(ModernResourceInner {
                queue: VecDeque::with_capacity(MODERN_PUSH_QUEUE_CAPACITY),
                dropped_events: 0,
                enqueued_generation: 0,
                next_listener_token: 0,
                active_listener: None,
                pending_notification: None,
                #[cfg(test)]
                notification_send_pause: None,
                #[cfg(test)]
                notification_send_attempts: 0,
                #[cfg(test)]
                notification_send_failures: 0,
                closed: false,
            }),
            _permit: permit,
        }
    }

    /// Queue one already-authorized event. At most one notification future may
    /// be pending per handle; additional pushes remain visible through read and
    /// are coalesced behind that notification.
    pub fn enqueue(&self, event: serde_json::Value) -> Option<PendingModernNotification> {
        let mut inner = self.lock_inner();
        if inner.closed {
            return None;
        }
        if inner.queue.len() == MODERN_PUSH_QUEUE_CAPACITY {
            inner.queue.pop_front();
            inner.dropped_events = inner.dropped_events.saturating_add(1);
        }
        inner.queue.push_back(event);
        inner.enqueued_generation = inner.enqueued_generation.wrapping_add(1);
        if inner.pending_notification.is_some() {
            return None;
        }
        let listener = inner.active_listener.as_ref()?;
        let work = PendingModernNotification {
            listener_token: listener.token,
            sink: listener.sink.clone(),
            cancel: listener.cancel.clone(),
            uri: push_resource_uri(&self.handle),
            generation: inner.enqueued_generation,
        };
        inner.pending_notification = Some(listener.token);
        Some(work)
    }

    pub async fn send_pending_notification(&self, mut work: PendingModernNotification) {
        loop {
            #[cfg(test)]
            {
                let mut inner = self.lock_inner();
                inner.notification_send_attempts =
                    inner.notification_send_attempts.saturating_add(1);
            }
            let result = tokio::select! {
                biased;
                () = work.cancel.cancelled() => Ok(()),
                result = work.sink.notify_resource_updated(work.uri.clone()) => result,
            };
            #[cfg(test)]
            if result.is_err() {
                let mut inner = self.lock_inner();
                inner.notification_send_failures =
                    inner.notification_send_failures.saturating_add(1);
            }
            #[cfg(test)]
            let notification_send_pause = self.lock_inner().notification_send_pause.take();
            #[cfg(test)]
            if let Some((sent, resume)) = notification_send_pause {
                sent.wait().await;
                resume.wait().await;
            }
            let (listener_to_cancel, next_work) = {
                let mut inner = self.lock_inner();
                if inner.pending_notification != Some(work.listener_token) {
                    return;
                }
                inner.pending_notification = None;
                let active_listener_matches = inner
                    .active_listener
                    .as_ref()
                    .is_some_and(|listener| listener.token == work.listener_token);
                if result.is_err() && active_listener_matches {
                    (inner.active_listener.take(), None)
                } else if result.is_ok()
                    && active_listener_matches
                    && !inner.queue.is_empty()
                    && inner.enqueued_generation != work.generation
                {
                    let Some(listener) = inner.active_listener.as_ref() else {
                        return;
                    };
                    let next = PendingModernNotification {
                        listener_token: listener.token,
                        sink: listener.sink.clone(),
                        cancel: listener.cancel.clone(),
                        uri: push_resource_uri(&self.handle),
                        generation: inner.enqueued_generation,
                    };
                    inner.pending_notification = Some(listener.token);
                    (None, Some(next))
                } else {
                    (None, None)
                }
            };
            if let Some(listener) = listener_to_cancel {
                listener.cancel.cancel();
            }
            let Some(next_work) = next_work else {
                return;
            };
            work = next_work;
        }
    }

    pub(super) fn attach_listener(
        &self,
        sink: SubscriptionSink,
    ) -> Result<
        (
            u64,
            tokio_util::sync::CancellationToken,
            Option<PendingModernNotification>,
        ),
        String,
    > {
        let cancel = tokio_util::sync::CancellationToken::new();
        let listener_sink = sink.clone();
        let (listener_token, previous, pending) = {
            let mut inner = self.lock_inner();
            if inner.closed {
                return Err("push resource not found for current caller".to_string());
            }
            inner.next_listener_token = inner.next_listener_token.wrapping_add(1);
            let listener_token = inner.next_listener_token;
            inner.pending_notification = None;
            let previous = inner
                .active_listener
                .replace(ActiveModernListener {
                    token: listener_token,
                    sink,
                    cancel: cancel.clone(),
                })
                .map(|listener| listener.cancel);
            let pending = if inner.queue.is_empty() {
                None
            } else {
                inner.pending_notification = Some(listener_token);
                Some(PendingModernNotification {
                    listener_token,
                    sink: listener_sink,
                    cancel: cancel.clone(),
                    uri: push_resource_uri(&self.handle),
                    generation: inner.enqueued_generation,
                })
            };
            (listener_token, previous, pending)
        };
        if let Some(previous) = previous {
            previous.cancel();
        }
        Ok((listener_token, cancel, pending))
    }

    pub(super) fn detach_listener_if(&self, listener_token: u64) {
        let removed = {
            let mut inner = self.lock_inner();
            if inner
                .active_listener
                .as_ref()
                .is_some_and(|listener| listener.token == listener_token)
            {
                inner.pending_notification = None;
                inner.active_listener.take()
            } else {
                None
            }
        };
        if let Some(listener) = removed {
            listener.cancel.cancel();
        }
    }

    pub(super) fn close(&self) {
        let listener = {
            let mut inner = self.lock_inner();
            inner.closed = true;
            inner.queue.clear();
            inner.dropped_events = 0;
            inner.pending_notification = None;
            inner.active_listener.take()
        };
        if let Some(listener) = listener {
            listener.cancel.cancel();
        }
    }

    pub(super) fn drain(&self) -> Result<(Vec<serde_json::Value>, u64), String> {
        let mut inner = self.lock_inner();
        if inner.closed {
            return Err("push resource not found for current caller".to_string());
        }
        let events = inner.queue.drain(..).collect();
        let dropped = std::mem::take(&mut inner.dropped_events);
        Ok((events, dropped))
    }

    pub(super) fn lock_inner(&self) -> std::sync::MutexGuard<'_, ModernResourceInner> {
        match self.inner.lock() {
            Ok(inner) => inner,
            Err(poisoned) => {
                tracing::error!(
                    handle = %self.handle,
                    "recovering poisoned modern push queue lock"
                );
                self.inner.clear_poison();
                poisoned.into_inner()
            }
        }
    }
}