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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum EventConnectionState {
12 Connected,
14 Disconnected,
16 Reconnecting,
18 Recovered,
20}
21
22impl EventConnectionState {
23 #[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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub struct EventConnectionObservation {
38 pub state: EventConnectionState,
40 pub reconnect_attempt: u32,
42 pub backoff: Duration,
44}
45
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub struct EventReconnectPolicy {
49 pub initial_delay: Duration,
51 pub max_delay: Duration,
53 pub stable_reset_after: Duration,
55 pub jitter_min_percent: u16,
57 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 #[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#[derive(Debug)]
93pub enum EventSubscriptionUpdate<T> {
94 Connection(EventConnectionObservation),
96 Item(T),
98}
99
100#[async_trait]
105pub trait EventSubscriptionSource: Send + Sync {
106 type Item: Send;
108 type Subscription: Send;
110 type Error: Display + Send + Sync;
112
113 async fn subscribe(&self) -> Result<Self::Subscription, Self::Error>;
115
116 async fn receive(
118 &self,
119 subscription: &mut Self::Subscription,
120 ) -> Result<Self::Item, Self::Error>;
121}
122
123#[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}