Skip to main content

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