๐Ÿ“ฆ RightNow-AI / openfang

๐Ÿ“„ retry.rs ยท 514 lines
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514//! Generic retry with exponential backoff and jitter.
//!
//! Provides a configurable, async-aware retry utility that can be used for
//! LLM API calls, network operations, channel message delivery, and any
//! other fallible async operation across the OpenFang codebase.
//!
//! Jitter uses `std::time::SystemTime` UNIX nanos as a seed to avoid
//! requiring the `rand` crate as a dependency.

use tracing::{debug, warn};

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

/// Configuration for retry behavior.
#[derive(Debug, Clone)]
pub struct RetryConfig {
    /// Maximum number of attempts (including the first try).
    pub max_attempts: u32,
    /// Minimum delay between retries in milliseconds.
    pub min_delay_ms: u64,
    /// Maximum delay between retries in milliseconds.
    pub max_delay_ms: u64,
    /// Jitter factor (0.0 = no jitter, 1.0 = full jitter).
    ///
    /// The actual sleep is `delay * (1 + random_fraction * jitter)`, where
    /// `random_fraction` is in `[0, 1)`.
    pub jitter: f64,
}

impl Default for RetryConfig {
    fn default() -> Self {
        Self {
            max_attempts: 3,
            min_delay_ms: 300,
            max_delay_ms: 30_000,
            jitter: 0.2,
        }
    }
}

/// Result of a retry operation.
#[derive(Debug)]
pub enum RetryOutcome<T, E> {
    /// The operation succeeded.
    Success {
        /// The successful result.
        result: T,
        /// Total number of attempts made (1 = first try succeeded).
        attempts: u32,
    },
    /// All retries exhausted without success.
    Exhausted {
        /// The error from the last attempt.
        last_error: E,
        /// Total number of attempts made.
        attempts: u32,
    },
}

// ---------------------------------------------------------------------------
// Backoff computation
// ---------------------------------------------------------------------------

/// Compute the delay for a given attempt (0-indexed).
///
/// Formula: `min(min_delay * 2^attempt, max_delay) * (1 + random * jitter)`
///
/// Uses `std::time::SystemTime` nanos as a lightweight pseudo-random source
/// instead of requiring the `rand` crate.
pub fn compute_backoff(config: &RetryConfig, attempt: u32) -> u64 {
    // Exponential base: min_delay * 2^attempt, capped at max_delay.
    let base = config
        .min_delay_ms
        .saturating_mul(1u64.checked_shl(attempt).unwrap_or(u64::MAX));
    let capped = base.min(config.max_delay_ms);

    // Jitter: multiply by (1 + random_fraction * jitter).
    if config.jitter <= 0.0 {
        return capped;
    }

    let frac = pseudo_random_fraction();
    let jitter_offset = (capped as f64) * frac * config.jitter;
    let with_jitter = (capped as f64) + jitter_offset;

    // Clamp to max_delay (jitter can push slightly above).
    (with_jitter as u64).min(config.max_delay_ms)
}

/// Return a pseudo-random fraction in `[0, 1)` using the current system time
/// nanos. This is NOT cryptographically secure, but good enough for jitter.
fn pseudo_random_fraction() -> f64 {
    let nanos = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .subsec_nanos();
    // Mix the bits a bit to reduce predictability.
    let mixed = nanos.wrapping_mul(2654435761); // Knuth multiplicative hash
    (mixed as f64) / (u32::MAX as f64)
}

// ---------------------------------------------------------------------------
// Core retry function
// ---------------------------------------------------------------------------

/// Execute an async operation with retry.
///
/// # Parameters
///
/// - `config` โ€” retry configuration (attempts, delays, jitter).
/// - `operation` โ€” the async closure to execute. Called once per attempt.
/// - `should_retry` โ€” predicate that inspects the error and returns `true`
///   if the operation should be retried.
/// - `retry_after_hint` โ€” optional hint extractor. If it returns `Some(ms)`,
///   that delay is used instead of the computed backoff (but still capped at
///   `max_delay_ms`).
///
/// # Returns
///
/// A `RetryOutcome` indicating success or exhaustion.
pub async fn retry_async<F, Fut, T, E, P, H>(
    config: &RetryConfig,
    mut operation: F,
    should_retry: P,
    retry_after_hint: H,
) -> RetryOutcome<T, E>
where
    F: FnMut() -> Fut,
    Fut: std::future::Future<Output = Result<T, E>>,
    P: Fn(&E) -> bool,
    H: Fn(&E) -> Option<u64>,
    E: std::fmt::Debug,
{
    let max = config.max_attempts.max(1);
    let mut last_error: Option<E> = None;

    for attempt in 0..max {
        match operation().await {
            Ok(result) => {
                if attempt > 0 {
                    debug!(
                        attempt = attempt + 1,
                        "retry succeeded after {} previous failures", attempt
                    );
                }
                return RetryOutcome::Success {
                    result,
                    attempts: attempt + 1,
                };
            }
            Err(err) => {
                let is_last = attempt + 1 >= max;

                if is_last || !should_retry(&err) {
                    if !should_retry(&err) {
                        debug!(
                            attempt = attempt + 1,
                            "error is not retryable, giving up: {:?}", err
                        );
                    } else {
                        warn!(
                            attempt = attempt + 1,
                            max_attempts = max,
                            "all retry attempts exhausted: {:?}",
                            err
                        );
                    }
                    return RetryOutcome::Exhausted {
                        last_error: err,
                        attempts: attempt + 1,
                    };
                }

                // Determine delay.
                let hint = retry_after_hint(&err);
                let delay_ms = if let Some(hinted) = hint {
                    // Respect the hint, but cap it.
                    hinted.min(config.max_delay_ms)
                } else {
                    compute_backoff(config, attempt)
                };

                debug!(
                    attempt = attempt + 1,
                    delay_ms, "retrying after error: {:?}", err
                );

                tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await;

                last_error = Some(err);
            }
        }
    }

    // Should not be reachable, but handle gracefully.
    RetryOutcome::Exhausted {
        last_error: last_error.expect("at least one attempt should have been made"),
        attempts: max,
    }
}

// ---------------------------------------------------------------------------
// Pre-built configs
// ---------------------------------------------------------------------------

/// Retry config for LLM API calls.
///
/// 3 attempts, 1s initial delay, up to 60s, 20% jitter.
pub fn llm_retry_config() -> RetryConfig {
    RetryConfig {
        max_attempts: 3,
        min_delay_ms: 1_000,
        max_delay_ms: 60_000,
        jitter: 0.2,
    }
}

/// Retry config for network operations (webhooks, fetches).
///
/// 3 attempts, 500ms initial delay, up to 30s, 10% jitter.
pub fn network_retry_config() -> RetryConfig {
    RetryConfig {
        max_attempts: 3,
        min_delay_ms: 500,
        max_delay_ms: 30_000,
        jitter: 0.1,
    }
}

/// Retry config for channel message delivery.
///
/// 3 attempts, 400ms initial delay, up to 15s, 10% jitter.
pub fn channel_retry_config() -> RetryConfig {
    RetryConfig {
        max_attempts: 3,
        min_delay_ms: 400,
        max_delay_ms: 15_000,
        jitter: 0.1,
    }
}

// ===========================================================================
// Tests
// ===========================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicU32, Ordering};
    use std::sync::Arc;

    #[test]
    fn test_retry_config_defaults() {
        let config = RetryConfig::default();
        assert_eq!(config.max_attempts, 3);
        assert_eq!(config.min_delay_ms, 300);
        assert_eq!(config.max_delay_ms, 30_000);
        assert!((config.jitter - 0.2).abs() < f64::EPSILON);
    }

    #[test]
    fn test_compute_backoff_exponential() {
        let config = RetryConfig {
            max_attempts: 5,
            min_delay_ms: 100,
            max_delay_ms: 100_000,
            jitter: 0.0, // no jitter for deterministic test
        };

        // 100 * 2^0 = 100
        assert_eq!(compute_backoff(&config, 0), 100);
        // 100 * 2^1 = 200
        assert_eq!(compute_backoff(&config, 1), 200);
        // 100 * 2^2 = 400
        assert_eq!(compute_backoff(&config, 2), 400);
        // 100 * 2^3 = 800
        assert_eq!(compute_backoff(&config, 3), 800);
    }

    #[test]
    fn test_compute_backoff_capped() {
        let config = RetryConfig {
            max_attempts: 10,
            min_delay_ms: 1_000,
            max_delay_ms: 5_000,
            jitter: 0.0,
        };

        // 1000 * 2^0 = 1000
        assert_eq!(compute_backoff(&config, 0), 1_000);
        // 1000 * 2^1 = 2000
        assert_eq!(compute_backoff(&config, 1), 2_000);
        // 1000 * 2^2 = 4000
        assert_eq!(compute_backoff(&config, 2), 4_000);
        // 1000 * 2^3 = 8000, capped at 5000
        assert_eq!(compute_backoff(&config, 3), 5_000);
        // Further attempts stay capped
        assert_eq!(compute_backoff(&config, 10), 5_000);
    }

    #[tokio::test]
    async fn test_retry_success_first_try() {
        let config = RetryConfig {
            max_attempts: 3,
            min_delay_ms: 10,
            max_delay_ms: 100,
            jitter: 0.0,
        };

        let outcome = retry_async(
            &config,
            || async { Ok::<&str, &str>("hello") },
            |_| true,
            |_: &&str| None,
        )
        .await;

        match outcome {
            RetryOutcome::Success { result, attempts } => {
                assert_eq!(result, "hello");
                assert_eq!(attempts, 1);
            }
            _ => panic!("expected success"),
        }
    }

    #[tokio::test]
    async fn test_retry_success_after_failures() {
        let config = RetryConfig {
            max_attempts: 5,
            min_delay_ms: 1, // tiny delays for test speed
            max_delay_ms: 10,
            jitter: 0.0,
        };

        let counter = Arc::new(AtomicU32::new(0));
        let counter_clone = counter.clone();

        let outcome = retry_async(
            &config,
            move || {
                let c = counter_clone.clone();
                async move {
                    let n = c.fetch_add(1, Ordering::SeqCst);
                    if n < 2 {
                        Err("not yet")
                    } else {
                        Ok("finally")
                    }
                }
            },
            |_| true,
            |_: &&str| None,
        )
        .await;

        match outcome {
            RetryOutcome::Success { result, attempts } => {
                assert_eq!(result, "finally");
                assert_eq!(attempts, 3); // failed twice, succeeded on 3rd
            }
            _ => panic!("expected success"),
        }
    }

    #[tokio::test]
    async fn test_retry_exhausted() {
        let config = RetryConfig {
            max_attempts: 3,
            min_delay_ms: 1,
            max_delay_ms: 10,
            jitter: 0.0,
        };

        let outcome = retry_async(
            &config,
            || async { Err::<(), &str>("always fails") },
            |_| true,
            |_: &&str| None,
        )
        .await;

        match outcome {
            RetryOutcome::Exhausted {
                last_error,
                attempts,
            } => {
                assert_eq!(last_error, "always fails");
                assert_eq!(attempts, 3);
            }
            _ => panic!("expected exhausted"),
        }
    }

    #[tokio::test]
    async fn test_retry_non_retryable_error() {
        let config = RetryConfig {
            max_attempts: 5,
            min_delay_ms: 1,
            max_delay_ms: 10,
            jitter: 0.0,
        };

        let counter = Arc::new(AtomicU32::new(0));
        let counter_clone = counter.clone();

        let outcome = retry_async(
            &config,
            move || {
                let c = counter_clone.clone();
                async move {
                    c.fetch_add(1, Ordering::SeqCst);
                    Err::<(), &str>("fatal error")
                }
            },
            |_| false, // never retry
            |_: &&str| None,
        )
        .await;

        match outcome {
            RetryOutcome::Exhausted {
                last_error,
                attempts,
            } => {
                assert_eq!(last_error, "fatal error");
                assert_eq!(attempts, 1); // gave up immediately
            }
            _ => panic!("expected exhausted"),
        }

        assert_eq!(counter.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn test_retry_with_hint_delay() {
        let config = RetryConfig {
            max_attempts: 3,
            min_delay_ms: 10_000, // large base delay
            max_delay_ms: 60_000,
            jitter: 0.0,
        };

        let counter = Arc::new(AtomicU32::new(0));
        let counter_clone = counter.clone();

        let start = std::time::Instant::now();

        let outcome = retry_async(
            &config,
            move || {
                let c = counter_clone.clone();
                async move {
                    let n = c.fetch_add(1, Ordering::SeqCst);
                    if n < 1 {
                        Err("transient")
                    } else {
                        Ok("ok")
                    }
                }
            },
            |_| true,
            |_: &&str| Some(1), // hint: 1ms delay (overrides 10s base)
        )
        .await;

        let elapsed = start.elapsed();

        match outcome {
            RetryOutcome::Success { result, attempts } => {
                assert_eq!(result, "ok");
                assert_eq!(attempts, 2);
                // Should complete in well under 1 second (hint was 1ms,
                // not the 10s base delay).
                assert!(
                    elapsed.as_millis() < 5_000,
                    "retry took too long: {:?} โ€” hint should have overridden base delay",
                    elapsed
                );
            }
            _ => panic!("expected success"),
        }
    }

    #[test]
    fn test_llm_retry_config() {
        let config = llm_retry_config();
        assert_eq!(config.max_attempts, 3);
        assert_eq!(config.min_delay_ms, 1_000);
        assert_eq!(config.max_delay_ms, 60_000);
        assert!((config.jitter - 0.2).abs() < f64::EPSILON);
    }

    #[test]
    fn test_channel_retry_config() {
        let config = channel_retry_config();
        assert_eq!(config.max_attempts, 3);
        assert_eq!(config.min_delay_ms, 400);
        assert_eq!(config.max_delay_ms, 15_000);
        assert!((config.jitter - 0.1).abs() < f64::EPSILON);
    }

    #[test]
    fn test_network_retry_config() {
        let config = network_retry_config();
        assert_eq!(config.max_attempts, 3);
        assert_eq!(config.min_delay_ms, 500);
        assert_eq!(config.max_delay_ms, 30_000);
        assert!((config.jitter - 0.1).abs() < f64::EPSILON);
    }
}