| 1 | package bot |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "log/slog" |
| 6 | "time" |
| 7 | ) |
| 8 | |
| 9 | // This file holds the shared connection-lifecycle primitives for the platform |
| 10 | // adapters (qq / feishu / weixin). Each adapter previously hand-rolled the same |
| 11 | // "loop forever, attempt a connection, sleep a fixed 5s on failure" shape; two of |
| 12 | // them used time.Sleep, which ignores ctx and so left a Stop() blocked for the |
| 13 | // remaining delay. SleepCtx fixes that; RunWithRetry folds the persistent- |
| 14 | // connection reconnect loop (qq / feishu) into one cancellation-aware, |
| 15 | // exponentially-backing-off helper. |
| 16 | |
| 17 | const ( |
| 18 | defaultInitialDelay = 1 * time.Second |
| 19 | defaultMaxDelay = 30 * time.Second |
| 20 | defaultResetAfter = 60 * time.Second |
| 21 | ) |
| 22 | |
| 23 | // RetryConfig controls the reconnect/error backoff for RunWithRetry. The zero |
| 24 | // value uses sane defaults (1s → 30s exponential, reset after 60s up). |
| 25 | type RetryConfig struct { |
| 26 | // InitialDelay is the wait after the first failed/closed attempt. |
| 27 | InitialDelay time.Duration |
| 28 | // MaxDelay caps the exponential backoff. |
| 29 | MaxDelay time.Duration |
| 30 | // ResetAfter is the minimum time an attempt must stay connected for the |
| 31 | // backoff to reset to InitialDelay on its next failure — so a flaky reconnect |
| 32 | // escalates while a long-healthy connection that finally drops retries fast. |
| 33 | ResetAfter time.Duration |
| 34 | } |
| 35 | |
| 36 | func (c RetryConfig) withDefaults() RetryConfig { |
| 37 | if c.InitialDelay <= 0 { |
| 38 | c.InitialDelay = defaultInitialDelay |
| 39 | } |
| 40 | if c.MaxDelay <= 0 { |
| 41 | c.MaxDelay = defaultMaxDelay |
| 42 | } |
| 43 | if c.MaxDelay < c.InitialDelay { |
| 44 | c.MaxDelay = c.InitialDelay |
| 45 | } |
| 46 | if c.ResetAfter <= 0 { |
| 47 | c.ResetAfter = defaultResetAfter |
| 48 | } |
| 49 | return c |
| 50 | } |
| 51 | |
| 52 | // SleepCtx waits for d or until ctx is cancelled, whichever comes first. It |
| 53 | // returns true if the full duration elapsed and false if ctx was cancelled (or |
| 54 | // was already cancelled on entry). Adapter loops use it instead of time.Sleep so |
| 55 | // a Stop() takes effect promptly rather than blocking out the remaining delay. |
| 56 | func SleepCtx(ctx context.Context, d time.Duration) bool { |
| 57 | if ctx.Err() != nil { |
| 58 | return false |
| 59 | } |
| 60 | if d <= 0 { |
| 61 | return true |
| 62 | } |
| 63 | t := time.NewTimer(d) |
| 64 | defer t.Stop() |
| 65 | select { |
| 66 | case <-ctx.Done(): |
| 67 | return false |
| 68 | case <-t.C: |
| 69 | return true |
| 70 | } |
| 71 | } |
| 72 | |
| 73 | // nextDelay doubles cur, capped at max. The <= 0 guard catches overflow on very |
| 74 | // large durations. |
| 75 | func nextDelay(cur, max time.Duration) time.Duration { |
| 76 | next := cur * 2 |
| 77 | if next <= 0 || next > max { |
| 78 | next = max |
| 79 | } |
| 80 | return next |
| 81 | } |
| 82 | |
| 83 | // RunWithRetry drives a persistent-connection adapter. It calls attempt(ctx) — |
| 84 | // one full connection lifetime that blocks until the connection drops or errors — |
| 85 | // then waits a cancellation-aware exponential backoff and reconnects, repeating |
| 86 | // until ctx is cancelled. attempt MUST honor ctx and clean up its own connection |
| 87 | // before returning; it returns nil for a clean close and an error otherwise (both |
| 88 | // trigger a reconnect — only ctx cancellation stops the loop). |
| 89 | func RunWithRetry(ctx context.Context, log *slog.Logger, name string, cfg RetryConfig, attempt func(context.Context) error) { |
| 90 | if log == nil { |
| 91 | log = slog.Default() |
| 92 | } |
| 93 | cfg = cfg.withDefaults() |
| 94 | delay := cfg.InitialDelay |
| 95 | for { |
| 96 | if ctx.Err() != nil { |
| 97 | return |
| 98 | } |
| 99 | start := time.Now() |
| 100 | err := attempt(ctx) |
| 101 | if ctx.Err() != nil { |
| 102 | return |
| 103 | } |
| 104 | if err != nil { |
| 105 | log.Error(name+" connection failed", "err", err) |
| 106 | } else { |
| 107 | log.Warn(name + " connection closed") |
| 108 | } |
| 109 | if time.Since(start) >= cfg.ResetAfter { |
| 110 | delay = cfg.InitialDelay |
| 111 | } |
| 112 | if !SleepCtx(ctx, delay) { |
| 113 | return |
| 114 | } |
| 115 | delay = nextDelay(delay, cfg.MaxDelay) |
| 116 | } |
| 117 | } |
| 118 |