Skip to main content

wowlab_centrifuge/
subscription.rs

1use std::time::Duration;
2
3use tokio::sync::mpsc;
4use wowlab_types::sensitive::Sensitive;
5
6use crate::{
7    SubscriptionFilter, proto,
8    types::{StreamPosition, SubscribeResult, SubscriptionEvent},
9};
10
11const DEFAULT_MIN_RESUBSCRIBE_DELAY: Duration = Duration::from_millis(500);
12const DEFAULT_MAX_RESUBSCRIBE_DELAY: Duration = Duration::from_secs(20);
13
14pub(crate) const FLAG_CHANNEL_COMPACTION: i64 = 1;
15
16/// Lifecycle state of a channel subscription.
17#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
18#[non_exhaustive]
19pub enum SubscriptionState {
20    #[default]
21    Unsubscribed,
22    Subscribing,
23    Subscribed,
24}
25
26/// Channel subscription options and recovery settings.
27#[derive(Clone)]
28pub struct SubscriptionConfig {
29    pub channel: String,
30    pub token: Option<Sensitive<String>>,
31    pub data: Option<Vec<u8>>,
32    pub positioned: bool,
33    pub recoverable: bool,
34    pub join_leave: bool,
35    pub delta: Option<String>,
36    pub tags_filter: Option<SubscriptionFilter>,
37    pub min_resubscribe_delay: Duration,
38    pub max_resubscribe_delay: Duration,
39}
40
41impl SubscriptionConfig {
42    /// Creates default subscription settings for `channel`.
43    #[must_use]
44    pub fn new(channel: impl Into<String>) -> Self {
45        Self {
46            channel: channel.into(),
47            token: None,
48            data: None,
49            positioned: false,
50            recoverable: false,
51            join_leave: false,
52            delta: None,
53            tags_filter: None,
54            min_resubscribe_delay: DEFAULT_MIN_RESUBSCRIBE_DELAY,
55            max_resubscribe_delay: DEFAULT_MAX_RESUBSCRIBE_DELAY,
56        }
57    }
58
59    /// Sets the private-channel token.
60    #[must_use]
61    pub fn token(mut self, token: impl Into<String>) -> Self {
62        self.token = Some(Sensitive::new(token.into()));
63
64        self
65    }
66
67    /// Sets application data sent with the subscribe command.
68    #[must_use]
69    pub fn data(mut self, data: Vec<u8>) -> Self {
70        self.data = Some(data);
71
72        self
73    }
74
75    /// Enables or disables stream-position information.
76    #[must_use]
77    pub fn positioned(mut self, enabled: bool) -> Self {
78        self.positioned = enabled;
79
80        self
81    }
82
83    /// Enables or disables stream recovery after reconnecting.
84    #[must_use]
85    pub fn recoverable(mut self, enabled: bool) -> Self {
86        self.recoverable = enabled;
87
88        self
89    }
90
91    /// Enables or disables join and leave publications.
92    #[must_use]
93    pub fn join_leave(mut self, enabled: bool) -> Self {
94        self.join_leave = enabled;
95
96        self
97    }
98
99    /// Requests the named delta-compression format.
100    #[must_use]
101    pub fn delta(mut self, format: impl Into<String>) -> Self {
102        self.delta = Some(format.into());
103
104        self
105    }
106
107    /// Sets the server-side publication tag filter.
108    #[must_use]
109    pub fn tags_filter(mut self, filter: SubscriptionFilter) -> Self {
110        self.tags_filter = Some(filter);
111
112        self
113    }
114}
115
116impl std::fmt::Debug for SubscriptionConfig {
117    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
118        f.debug_struct("SubscriptionConfig")
119            .field("channel", &self.channel)
120            .field("token", &self.token.as_ref().map(|_| "[REDACTED]"))
121            .field("data", &self.data)
122            .field("positioned", &self.positioned)
123            .field("recoverable", &self.recoverable)
124            .field("join_leave", &self.join_leave)
125            .field("delta", &self.delta)
126            .field("tags_filter", &self.tags_filter)
127            .field("min_resubscribe_delay", &self.min_resubscribe_delay)
128            .field("max_resubscribe_delay", &self.max_resubscribe_delay)
129            .finish()
130    }
131}
132
133pub(crate) struct SubscriptionInner {
134    pub config: SubscriptionConfig,
135    pub state: SubscriptionState,
136    pub stream_position: Option<StreamPosition>,
137    pub event_tx: mpsc::Sender<SubscriptionEvent>,
138}
139
140impl SubscriptionInner {
141    pub(crate) fn new(
142        config: SubscriptionConfig,
143        event_tx: mpsc::Sender<SubscriptionEvent>,
144    ) -> Self {
145        Self {
146            config,
147            state: SubscriptionState::Unsubscribed,
148            stream_position: None,
149            event_tx,
150        }
151    }
152
153    pub(crate) fn on_subscribed(&mut self, result: &SubscribeResult) {
154        self.state = SubscriptionState::Subscribed;
155
156        if result.recoverable {
157            self.stream_position = Some(StreamPosition {
158                offset: result.offset,
159                epoch: result.epoch.clone(),
160            });
161        }
162    }
163
164    pub(crate) fn update_position(&mut self, offset: u64) {
165        if let Some(ref mut pos) = self.stream_position {
166            pos.offset = offset;
167        }
168    }
169
170    pub(crate) fn build_subscribe_request(&self, recover: bool) -> proto::SubscribeRequest {
171        let mut req = proto::SubscribeRequest {
172            channel: self.config.channel.clone(),
173            token: self
174                .config
175                .token
176                .as_ref()
177                .map(|t| t.expose().clone())
178                .unwrap_or_default(),
179            data: self.config.data.clone().unwrap_or_default(),
180            positioned: self.config.positioned,
181            recoverable: self.config.recoverable,
182            join_leave: self.config.join_leave,
183            delta: self.config.delta.clone().unwrap_or_default(),
184            tf: self
185                .config
186                .tags_filter
187                .as_ref()
188                .map(proto::FilterNode::from),
189            flag: FLAG_CHANNEL_COMPACTION,
190            ..Default::default()
191        };
192
193        if recover {
194            if let Some(ref pos) = self.stream_position {
195                req.recover = true;
196                req.offset = pos.offset;
197                req.epoch.clone_from(&pos.epoch);
198            }
199        }
200
201        req
202    }
203}
204
205/// Handle for receiving events from a subscribed channel; drop to stop receiving.
206#[derive(Debug)]
207pub struct Subscription {
208    pub channel: String,
209    event_rx: mpsc::Receiver<SubscriptionEvent>,
210}
211
212impl Subscription {
213    pub(crate) fn new(channel: String, event_rx: mpsc::Receiver<SubscriptionEvent>) -> Self {
214        Self { channel, event_rx }
215    }
216
217    /// Waits for the next subscription event, or returns `None` when the client drops the sender.
218    pub async fn recv(&mut self) -> Option<SubscriptionEvent> {
219        self.event_rx.recv().await
220    }
221
222    /// Returns the next queued event without waiting.
223    pub fn try_recv(&mut self) -> Option<SubscriptionEvent> {
224        self.event_rx.try_recv().ok()
225    }
226}
227
228#[cfg(test)]
229mod tests {
230    use googletest::prelude::*;
231
232    use super::SubscriptionConfig;
233
234    #[gtest]
235    fn debug_redacts_subscription_token() -> Result<()> {
236        let config = SubscriptionConfig::new("jobs").token("subscription-secret");
237        let debug = format!("{config:?}");
238
239        verify_that!(debug.as_str(), contains_substring("[REDACTED]"))?;
240
241        verify_that!(
242            debug.as_str(),
243            not(contains_substring("subscription-secret"))
244        )
245    }
246}