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 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
36pub struct EventConnectionObservation {
37 pub state: EventConnectionState,
39 pub reconnect_attempt: u32,
41 pub backoff: Duration,
43}
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub struct EventReconnectPolicy {
48 pub initial_delay: Duration,
50 pub max_delay: Duration,
52 pub stable_reset_after: Duration,
54 pub jitter_min_percent: u16,
56 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 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#[derive(Debug)]
91pub enum EventSubscriptionUpdate<T> {
92 Connection(EventConnectionObservation),
94 Item(T),
96}
97
98#[async_trait]
103pub trait EventSubscriptionSource: Send + Sync {
104 type Item: Send;
106 type Subscription: Send;
108 type Error: Display + Send + Sync;
110
111 async fn subscribe(&self) -> Result<Self::Subscription, Self::Error>;
113
114 async fn receive(
116 &self,
117 subscription: &mut Self::Subscription,
118 ) -> Result<Self::Item, Self::Error>;
119}
120
121pub 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}