| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "errors" |
| 6 | "sync" |
| 7 | "testing" |
| 8 | "time" |
| 9 | |
| 10 | "reasonix/internal/event" |
| 11 | "reasonix/internal/provider" |
| 12 | ) |
| 13 | |
| 14 | func TestWaitingStopsOnceBudgetExhausted(t *testing.T) { |
| 15 | p := &transientHeaderProvider{} |
| 16 | oldBudget, oldSleep := recoveryWaitBudget, recoverySleep |
| 17 | defer func() { recoveryWaitBudget, recoverySleep = oldBudget, oldSleep }() |
| 18 | recoveryWaitBudget = 3 * time.Minute |
| 19 | var slept time.Duration |
| 20 | recoverySleep = func(ctx context.Context, d time.Duration) bool { |
| 21 | slept += d |
| 22 | return ctx.Err() == nil |
| 23 | } |
| 24 | sink := &recordSink{} |
| 25 | a := New(p, echoRegistry(), NewSession(""), Options{}, sink) |
| 26 | err := a.Run(withNoClosedLoop(context.Background()), "go") |
| 27 | var exhausted *provider.RecoveryWaitExhaustedError |
| 28 | if !errors.As(err, &exhausted) { |
| 29 | t.Fatalf("calls=%d err=%v", p.calls, err) |
| 30 | } |
| 31 | if exhausted.Phase != "headers" || exhausted.Status != 503 || exhausted.Attempts != 6 || p.calls != 6 { |
| 32 | t.Fatalf("calls=%d exhausted=%+v", p.calls, exhausted) |
| 33 | } |
| 34 | if exhausted.Waited > recoveryWaitBudget || exhausted.Waited+time.Minute <= recoveryWaitBudget || slept > exhausted.Waited { |
| 35 | t.Fatalf("waited=%s slept=%s budget=%s", exhausted.Waited, slept, recoveryWaitBudget) |
| 36 | } |
| 37 | if provider.ClassifyRecovery(err).Retryable { |
| 38 | t.Fatal("exhausted wait classified as retryable") |
| 39 | } |
| 40 | waiting := 0 |
| 41 | for _, e := range sink.kinds(event.Retrying) { |
| 42 | if e.Recovery == nil { |
| 43 | continue |
| 44 | } |
| 45 | if !e.Recovery.Waiting { |
| 46 | if e.Recovery.WaitBudgetMs != 0 { |
| 47 | t.Fatalf("short retry advertised a wait budget: %+v", e.Recovery) |
| 48 | } |
| 49 | continue |
| 50 | } |
| 51 | waiting++ |
| 52 | if e.Recovery.WaitBudgetMs != recoveryWaitBudget.Milliseconds() || e.Recovery.NextAttemptAt == 0 || e.Recovery.WaitedMs < 0 { |
| 53 | t.Fatalf("recovery=%+v", e.Recovery) |
| 54 | } |
| 55 | } |
| 56 | if waiting != 2 { |
| 57 | t.Fatalf("waiting retries=%d", waiting) |
| 58 | } |
| 59 | } |
| 60 | |
| 61 | func TestOversizedRetryAfterNeverStartsAnUnaffordableWait(t *testing.T) { |
| 62 | p := &retryAfterProvider{after: time.Hour} |
| 63 | old := recoverySleep |
| 64 | defer func() { recoverySleep = old }() |
| 65 | recoverySleep = func(context.Context, time.Duration) bool { return true } |
| 66 | a := New(p, echoRegistry(), NewSession(""), Options{}, event.Discard) |
| 67 | err := a.Run(withNoClosedLoop(context.Background()), "go") |
| 68 | var exhausted *provider.RecoveryWaitExhaustedError |
| 69 | if !errors.As(err, &exhausted) || exhausted.Attempts != maxSamplingAttempts || p.calls != maxSamplingAttempts { |
| 70 | t.Fatalf("calls=%d err=%v", p.calls, err) |
| 71 | } |
| 72 | } |
| 73 | |
| 74 | type retryAfterProvider struct { |
| 75 | calls int |
| 76 | after time.Duration |
| 77 | } |
| 78 | |
| 79 | func (*retryAfterProvider) Name() string { return "retry-after" } |
| 80 | func (p *retryAfterProvider) Stream(context.Context, provider.Request) (<-chan provider.Chunk, error) { |
| 81 | p.calls++ |
| 82 | return nil, &provider.APIError{Status: 429, RetryAfter: p.after} |
| 83 | } |
| 84 | |
| 85 | type cancelOnWaitSink struct { |
| 86 | recordSink |
| 87 | once sync.Once |
| 88 | cancel context.CancelFunc |
| 89 | } |
| 90 | |
| 91 | func (s *cancelOnWaitSink) Emit(e event.Event) { |
| 92 | s.recordSink.Emit(e) |
| 93 | if e.Kind == event.Retrying && e.Recovery != nil && e.Recovery.Waiting { |
| 94 | s.once.Do(s.cancel) |
| 95 | } |
| 96 | } |
| 97 | |
| 98 | func TestWaitingCancelsPromptlyWithRealTimer(t *testing.T) { |
| 99 | p := &transientHeaderProvider{} |
| 100 | old := recoverySleep |
| 101 | defer func() { recoverySleep = old }() |
| 102 | recoverySleep = sleepRecovery |
| 103 | ctx, cancel := context.WithCancel(context.Background()) |
| 104 | defer cancel() |
| 105 | sink := &cancelOnWaitSink{cancel: cancel} |
| 106 | a := New(p, echoRegistry(), NewSession(""), Options{}, sink) |
| 107 | started := time.Now() |
| 108 | err := a.Run(withNoClosedLoop(ctx), "go") |
| 109 | if !errors.Is(err, context.Canceled) || p.calls != maxSamplingAttempts { |
| 110 | t.Fatalf("calls=%d err=%v", p.calls, err) |
| 111 | } |
| 112 | if elapsed := time.Since(started); elapsed > 5*time.Second { |
| 113 | t.Fatalf("cancel during the minute-long wait took %s", elapsed) |
| 114 | } |
| 115 | } |
| 116 |