ruststream_gcp_pubsub/
subscriber.rs

1//! [`PubSubSubscriber`]: a stream of deliveries backed by a pump task.
2//!
3//! The client's `MessageStream` exposes an inherent `next()` that is not documented
4//! cancel-safe, so the crate owns a pump task per subscription: it drives `next`, converts
5//! deliveries, and forwards them into a bounded channel the subscriber polls. Settlement needs
6//! no round trip through the pump - the client's `Handler` travels inside each message and is
7//! consumed by `ack`/`nack` directly.
8
9use futures::Stream;
10
11use google_cloud_pubsub::subscriber::{MessageStream, ShutdownToken};
12use ruststream::Subscriber;
13use tokio::sync::mpsc;
14
15use crate::broker::Core;
16use crate::error::{PubSubError, box_err};
17use crate::message::PubSubMessage;
18use crate::subscription::PubSubSubscription;
19
20/// How many converted deliveries may sit between the pump and the consumer. Real prefetch is
21/// the client's own flow control (`max_outstanding`); this only decouples the two loops.
22const CHANNEL_CAPACITY: usize = 16;
23
24/// A subscription to one Pub/Sub subscription; yields [`PubSubMessage`]s.
25///
26/// Dropping the subscriber signals the client's shutdown token, which drains the stream and
27/// stops the pump task.
28pub struct PubSubSubscriber {
29    subscription: String,
30    rx: mpsc::Receiver<Result<PubSubMessage, PubSubError>>,
31    shutdown: ShutdownToken,
32}
33
34impl std::fmt::Debug for PubSubSubscriber {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        f.debug_struct("PubSubSubscriber")
37            .field("subscription", &self.subscription)
38            .finish_non_exhaustive()
39    }
40}
41
42impl PubSubSubscriber {
43    /// The full resource name of the subscription this stream consumes from.
44    #[must_use]
45    pub fn subscription(&self) -> &str {
46        &self.subscription
47    }
48
49    /// Opens the stream synchronously (the client connects lazily) and spawns the pump.
50    pub(crate) fn open(core: &Core, descriptor: &PubSubSubscription) -> Self {
51        let name = core.subscription_name(descriptor.subscription());
52        let mut builder = core.subscriber.subscribe(name.clone());
53        if let Some(messages) = descriptor.max_outstanding_value() {
54            builder = builder.set_max_outstanding_messages(messages);
55        }
56        if let Some(extension) = descriptor.ack_extension_value() {
57            builder = builder.set_max_lease_extension(extension);
58        }
59        let stream = builder.build();
60        let shutdown = stream.shutdown_token();
61
62        let (tx, rx) = mpsc::channel(CHANNEL_CAPACITY);
63        tokio::spawn(pump(stream, tx, name.clone()));
64
65        Self {
66            subscription: name,
67            rx,
68            shutdown,
69        }
70    }
71}
72
73impl Drop for PubSubSubscriber {
74    fn drop(&mut self) {
75        // `shutdown` is async and destructors are sync; the spawned signal drains the stream,
76        // which ends the pump task. Best effort by design: if the runtime is already gone, the
77        // pump dies with it.
78        let token = self.shutdown.clone();
79        if let Ok(handle) = tokio::runtime::Handle::try_current() {
80            handle.spawn(async move {
81                token.shutdown().await;
82            });
83        }
84    }
85}
86
87impl Subscriber for PubSubSubscriber {
88    type Message = PubSubMessage;
89    type Error = PubSubError;
90
91    fn stream(&mut self) -> impl Stream<Item = Result<PubSubMessage, PubSubError>> + Send + '_ {
92        // Poll the channel in place rather than wrapping it in an owning stream, so `stream`
93        // can be called again after the returned stream is dropped (the runtime and the
94        // conformance helpers re-enter it per call).
95        futures::stream::poll_fn(move |cx| self.rx.poll_recv(cx))
96    }
97}
98
99async fn pump(
100    mut stream: MessageStream,
101    out: mpsc::Sender<Result<PubSubMessage, PubSubError>>,
102    subscription: String,
103) {
104    while let Some(item) = stream.next().await {
105        match item {
106            Ok((message, handler)) => {
107                if out
108                    .send(Ok(PubSubMessage::new(message, handler)))
109                    .await
110                    .is_err()
111                {
112                    // Subscriber dropped; its Drop has already signalled shutdown.
113                    break;
114                }
115            }
116            Err(err) => {
117                // The client absorbs transient failures internally; an error here is
118                // permanent and closes the stream.
119                let _ = out
120                    .send(Err(PubSubError::Receive {
121                        subscription: subscription.clone(),
122                        source: box_err(err),
123                    }))
124                    .await;
125                break;
126            }
127        }
128    }
129}
ruststream_gcp_pubsub/subscriber.rs

ruststream_gcp_pubsub/
subscriber.rs
