aster_forge_events/
supervisor.rs

1use std::fmt::Display;
2use std::sync::Arc;
3use std::time::{Duration, Instant};
4
5use async_trait::async_trait;
6use tokio::sync::mpsc;
7use tokio_util::sync::CancellationToken;
8
9/// Connection lifecycle state emitted by a reconnecting event subscriber.
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum EventConnectionState {
12    /// A first subscription became ready.
13    Connected,
14    /// An established subscription stopped or could not be opened.
15    Disconnected,
16    /// The subscriber is waiting before another connection attempt.
17    Reconnecting,
18    /// A subscription recovered after a previous disconnect.
19    Recovered,
20}
21
22impl EventConnectionState {
23    /// Returns the stable low-cardinality label for this state.
24    #[must_use]
25    pub const fn as_label(self) -> &'static str {
26        match self {
27            Self::Connected => "connected",
28            Self::Disconnected => "disconnected",
29            Self::Reconnecting => "reconnecting",
30            Self::Recovered => "recovered",
31        }
32    }
33}
34
35/// Connection lifecycle observation emitted by a subscription supervisor.
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub struct EventConnectionObservation {
38    /// Current connection state.
39    pub state: EventConnectionState,
40    /// One-based reconnect attempt number, or zero for the first connection.
41    pub reconnect_attempt: u32,
42    /// Backoff selected for the next reconnect attempt.
43    pub backoff: Duration,
44}
45
46/// Reconnect policy for a transient event subscription.
47#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub struct EventReconnectPolicy {
49    /// Initial reconnect delay.
50    pub initial_delay: Duration,
51    /// Maximum reconnect delay.
52    pub max_delay: Duration,
53    /// Stable connection duration after which the attempt counter resets.
54    pub stable_reset_after: Duration,
55    /// Minimum jitter percentage applied to the selected delay.
56    pub jitter_min_percent: u16,
57    /// Maximum jitter percentage applied to the selected delay.
58    pub jitter_max_percent: u16,
59}
60
61impl Default for EventReconnectPolicy {
62    fn default() -> Self {
63        Self {
64            initial_delay: Duration::from_millis(250),
65            max_delay: Duration::from_secs(30),
66            stable_reset_after: Duration::from_secs(30),
67            jitter_min_percent: 80,
68            jitter_max_percent: 120,
69        }
70    }
71}
72
73impl EventReconnectPolicy {
74    /// Calculates the bounded jittered delay for a one-based reconnect attempt.
75    #[must_use]
76    pub fn reconnect_delay(self, reconnect_attempt: u32) -> Duration {
77        let raw = aster_forge_utils::backoff::exponential_delay(
78            self.initial_delay,
79            reconnect_attempt.saturating_sub(1),
80        );
81        let capped = aster_forge_utils::backoff::cap_delay(raw, self.max_delay);
82        let jittered = aster_forge_utils::backoff::randomized_jitter(
83            capped,
84            self.jitter_min_percent,
85            self.jitter_max_percent,
86        );
87        aster_forge_utils::backoff::cap_delay(jittered, self.max_delay)
88    }
89}
90
91/// One update emitted by a reconnecting subscription supervisor.
92#[derive(Debug)]
93pub enum EventSubscriptionUpdate<T> {
94    /// The transport connection changed state.
95    Connection(EventConnectionObservation),
96    /// The transport delivered one product-owned item.
97    Item(T),
98}
99
100/// A transport-specific source that can open and receive from one subscription.
101///
102/// The shared supervisor owns retries and lifecycle observations. Implementations own only one
103/// connection attempt and one-item receive semantics.
104#[async_trait]
105pub trait EventSubscriptionSource: Send + Sync {
106    /// Item emitted by the transport.
107    type Item: Send;
108    /// One active transport subscription.
109    type Subscription: Send;
110    /// Transport error returned by subscribe or receive.
111    type Error: Display + Send + Sync;
112
113    /// Opens one subscription attempt.
114    async fn subscribe(&self) -> Result<Self::Subscription, Self::Error>;
115
116    /// Receives one item from an active subscription.
117    async fn receive(
118        &self,
119        subscription: &mut Self::Subscription,
120    ) -> Result<Self::Item, Self::Error>;
121}
122
123/// Supervises one transient subscription until shutdown or receiver closure.
124///
125/// Consumers handle updates sequentially. `Connected`/`Recovered` is therefore observed before
126/// the first item from that subscription, so products can reconcile authoritative state first.
127#[expect(
128    clippy::too_many_lines,
129    reason = "The supervisor keeps one explicit subscription and reconnect state machine in a single loop."
130)]
131pub async fn supervise_event_subscription<S>(
132    source: Arc<S>,
133    reconnect_policy: EventReconnectPolicy,
134    shutdown: CancellationToken,
135    updates: mpsc::Sender<EventSubscriptionUpdate<S::Item>>,
136) where
137    S: EventSubscriptionSource + ?Sized,
138{
139    let mut reconnect_attempt = 0_u32;
140
141    loop {
142        let subscription = tokio::select! {
143            () = shutdown.cancelled() => return,
144            result = source.subscribe() => result,
145        };
146        let mut subscription = match subscription {
147            Ok(subscription) => subscription,
148            Err(error) => {
149                reconnect_attempt = reconnect_attempt.saturating_add(1);
150                let delay = reconnect_policy.reconnect_delay(reconnect_attempt);
151                if !send_connection_update(
152                    &updates,
153                    &shutdown,
154                    EventConnectionState::Disconnected,
155                    reconnect_attempt,
156                    Duration::ZERO,
157                )
158                .await
159                    || !send_connection_update(
160                        &updates,
161                        &shutdown,
162                        EventConnectionState::Reconnecting,
163                        reconnect_attempt,
164                        delay,
165                    )
166                    .await
167                {
168                    return;
169                }
170                tracing::warn!(
171                    reconnect_attempt,
172                    backoff_ms = delay.as_millis(),
173                    error = %error,
174                    "event subscription attempt failed"
175                );
176                if sleep_or_shutdown(&shutdown, delay).await {
177                    return;
178                }
179                continue;
180            }
181        };
182
183        let connected_state = if reconnect_attempt == 0 {
184            EventConnectionState::Connected
185        } else {
186            EventConnectionState::Recovered
187        };
188        if !send_connection_update(
189            &updates,
190            &shutdown,
191            connected_state,
192            reconnect_attempt,
193            Duration::ZERO,
194        )
195        .await
196        {
197            return;
198        }
199
200        let connected_at = Instant::now();
201        loop {
202            let item = tokio::select! {
203                () = shutdown.cancelled() => return,
204                result = source.receive(&mut subscription) => result,
205            };
206            match item {
207                Ok(item) => {
208                    reconnect_attempt = 0;
209                    if !send_update(&updates, &shutdown, EventSubscriptionUpdate::Item(item)).await
210                    {
211                        return;
212                    }
213                }
214                Err(error) => {
215                    if connected_at.elapsed() >= reconnect_policy.stable_reset_after {
216                        reconnect_attempt = 0;
217                    }
218                    reconnect_attempt = reconnect_attempt.saturating_add(1);
219                    let delay = reconnect_policy.reconnect_delay(reconnect_attempt);
220                    if !send_connection_update(
221                        &updates,
222                        &shutdown,
223                        EventConnectionState::Disconnected,
224                        reconnect_attempt,
225                        Duration::ZERO,
226                    )
227                    .await
228                        || !send_connection_update(
229                            &updates,
230                            &shutdown,
231                            EventConnectionState::Reconnecting,
232                            reconnect_attempt,
233                            delay,
234                        )
235                        .await
236                    {
237                        return;
238                    }
239                    tracing::warn!(
240                        reconnect_attempt,
241                        backoff_ms = delay.as_millis(),
242                        error = %error,
243                        "event subscription disconnected"
244                    );
245                    if sleep_or_shutdown(&shutdown, delay).await {
246                        return;
247                    }
248                    break;
249                }
250            }
251        }
252    }
253}
254
255async fn send_connection_update<T>(
256    updates: &mpsc::Sender<EventSubscriptionUpdate<T>>,
257    shutdown: &CancellationToken,
258    state: EventConnectionState,
259    reconnect_attempt: u32,
260    backoff: Duration,
261) -> bool {
262    send_update(
263        updates,
264        shutdown,
265        EventSubscriptionUpdate::Connection(EventConnectionObservation {
266            state,
267            reconnect_attempt,
268            backoff,
269        }),
270    )
271    .await
272}
273
274async fn send_update<T>(
275    updates: &mpsc::Sender<EventSubscriptionUpdate<T>>,
276    shutdown: &CancellationToken,
277    update: EventSubscriptionUpdate<T>,
278) -> bool {
279    tokio::select! {
280        () = shutdown.cancelled() => false,
281        result = updates.send(update) => result.is_ok(),
282    }
283}
284
285async fn sleep_or_shutdown(shutdown: &CancellationToken, delay: Duration) -> bool {
286    tokio::select! {
287        () = shutdown.cancelled() => true,
288        () = tokio::time::sleep(delay) => false,
289    }
290}
291
292#[cfg(test)]
293mod tests {
294    use super::EventReconnectPolicy;
295    use std::time::Duration;
296
297    #[test]
298    fn reconnect_delay_doubles_and_caps_with_fixed_jitter() {
299        let policy = EventReconnectPolicy {
300            initial_delay: Duration::from_millis(100),
301            max_delay: Duration::from_millis(250),
302            stable_reset_after: Duration::from_secs(1),
303            jitter_min_percent: 100,
304            jitter_max_percent: 100,
305        };
306
307        assert_eq!(policy.reconnect_delay(1), Duration::from_millis(100));
308        assert_eq!(policy.reconnect_delay(2), Duration::from_millis(200));
309        assert_eq!(policy.reconnect_delay(3), Duration::from_millis(250));
310    }
311}