diff --git a/pkg/cli/logs_orchestrator.go b/pkg/cli/logs_orchestrator.go index 1d3742eb61d..f8a651bf7df 100644 --- a/pkg/cli/logs_orchestrator.go +++ b/pkg/cli/logs_orchestrator.go @@ -367,7 +367,13 @@ outerLoop: // always sleeps at least APICallCooldown but will also block until the // reset window when the remaining budget is nearly exhausted. if iteration > 0 { - if rlErr := checkAndWaitForRateLimit(verbose); rlErr != nil { + if rlErr := checkAndWaitForRateLimit(activeCtx, verbose); rlErr != nil { + if errors.Is(rlErr, context.Canceled) || errors.Is(rlErr, context.DeadlineExceeded) { + // Context was cancelled or timed out during the rate-limit wait. + // Use continue (not break) so the top-of-loop activeCtx.Done() path + // preserves Canceled vs DeadlineExceeded behavior. + continue + } logsOrchestratorLog.Printf("Rate limit check failed (using static cooldown): %v", rlErr) } } diff --git a/pkg/cli/logs_rate_limit.go b/pkg/cli/logs_rate_limit.go index 592732d545a..81e1deae5cf 100644 --- a/pkg/cli/logs_rate_limit.go +++ b/pkg/cli/logs_rate_limit.go @@ -11,7 +11,9 @@ package cli import ( + "context" "encoding/json" + "errors" "fmt" "os" "time" @@ -22,6 +24,7 @@ import ( ) var logsRateLimitLog = logger.New("cli:logs_rate_limit") +var fetchRateLimitFunc = fetchRateLimit // rateLimitResponse models the JSON returned by `gh api rate_limit`. // Only the "core" resource bucket is used because log downloads and @@ -70,21 +73,49 @@ func fetchRateLimit() (rateLimitResource, error) { return resp.Resources.Core, nil } +// sleepWithContext pauses for duration d and returns nil when the timer fires. +// If ctx is cancelled before the timer expires, it stops the timer and returns +// ctx.Err() so callers can propagate cancellation immediately. +func sleepWithContext(ctx context.Context, d time.Duration) error { + if ctx == nil { + ctx = context.Background() + } + select { + case <-ctx.Done(): + return ctx.Err() + default: + } + + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-timer.C: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + // checkAndWaitForRateLimit queries the GitHub API rate limit and sleeps until // the reset window when the remaining core request budget falls at or below // RateLimitThreshold. It always waits at least APICallCooldown between // successive calls so that even when requests are plentiful the orchestrator // does not hammer the API. // +// ctx is checked on every sleep so that a user cancellation (Ctrl-C) or +// deadline expiry wakes the function early and propagates the context error. +// // If the rate limit cannot be fetched (e.g. network error) the function falls // back to the static APICallCooldown sleep and returns the error so callers // can decide whether to surface it. -func checkAndWaitForRateLimit(verbose bool) error { - rl, err := fetchRateLimit() +func checkAndWaitForRateLimit(ctx context.Context, verbose bool) error { + rl, err := fetchRateLimitFunc() if err != nil { // Best-effort: fall back to static cooldown so the caller can continue. logsRateLimitLog.Printf("Could not fetch rate limit, using static cooldown: %v", err) - time.Sleep(APICallCooldown) + if sleepErr := sleepWithContext(ctx, APICallCooldown); sleepErr != nil { + return fmt.Errorf("rate-limit fetch failed and context was canceled or timed out during fallback cooldown: %w", errors.Join(err, sleepErr)) + } return err } @@ -94,8 +125,7 @@ func checkAndWaitForRateLimit(verbose bool) error { if waitDur <= 0 { // Reset has already passed; apply minimal cooldown and carry on. logsRateLimitLog.Print("Rate limit reset has already passed, applying minimal cooldown") - time.Sleep(APICallCooldown) - return nil + return sleepWithContext(ctx, APICallCooldown) } // Add a small buffer so we don't resume right on the boundary. @@ -107,8 +137,7 @@ func checkAndWaitForRateLimit(verbose bool) error { ) fmt.Fprintln(os.Stderr, console.FormatWarningMessage(msg)) logsRateLimitLog.Printf("Sleeping for rate limit reset: duration=%s", waitDur) - time.Sleep(waitDur) - return nil + return sleepWithContext(ctx, waitDur) } if verbose { @@ -118,6 +147,5 @@ func checkAndWaitForRateLimit(verbose bool) error { } // Even when budget is healthy, apply the minimum inter-call cooldown. - time.Sleep(APICallCooldown) - return nil + return sleepWithContext(ctx, APICallCooldown) } diff --git a/pkg/cli/logs_rate_limit_test.go b/pkg/cli/logs_rate_limit_test.go index cab86429a6f..cae157ef644 100644 --- a/pkg/cli/logs_rate_limit_test.go +++ b/pkg/cli/logs_rate_limit_test.go @@ -3,7 +3,9 @@ package cli import ( + "context" "encoding/json" + stderrors "errors" "strconv" "testing" "time" @@ -83,3 +85,100 @@ func TestRateLimitResourceIsBelowThreshold(t *testing.T) { func jsonInt(n int64) string { return strconv.FormatInt(n, 10) } + +// TestSleepWithContextCancellation verifies that sleepWithContext returns ctx.Err() +// immediately when the context is cancelled before the timer fires. +func TestSleepWithContextCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() // cancel immediately + + start := time.Now() + err := sleepWithContext(ctx, 10*time.Second) + elapsed := time.Since(start) + + require.ErrorIs(t, err, context.Canceled) + assert.Less(t, elapsed, time.Second, "sleepWithContext should return quickly when context is already cancelled") +} + +// TestSleepWithContextDeadlineExceeded verifies that sleepWithContext respects a +// deadline that expires before the sleep duration. +func TestSleepWithContextDeadlineExceeded(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + + start := time.Now() + err := sleepWithContext(ctx, 10*time.Second) + elapsed := time.Since(start) + + require.ErrorIs(t, err, context.DeadlineExceeded) + assert.Less(t, elapsed, time.Second, "sleepWithContext should return as soon as the deadline expires") +} + +// TestSleepWithContextTimerFires verifies that sleepWithContext returns nil when the +// timer fires before context cancellation. +func TestSleepWithContextTimerFires(t *testing.T) { + ctx := context.Background() + start := time.Now() + err := sleepWithContext(ctx, 5*time.Millisecond) + elapsed := time.Since(start) + require.NoError(t, err, "sleepWithContext should return nil when timer fires normally") + assert.Less(t, elapsed, time.Second, "timer should have fired and returned promptly") +} + +func TestSleepWithContextNilContext(t *testing.T) { + var nilCtx context.Context + start := time.Now() + err := sleepWithContext(nilCtx, 2*time.Millisecond) + elapsed := time.Since(start) + require.NoError(t, err, "nil context should fall back to background context") + assert.Less(t, elapsed, time.Second, "nil context should not block longer than timer duration") +} + +func TestSleepWithContextAlreadyCanceled(t *testing.T) { + canceledCtx, cancel := context.WithCancel(context.Background()) + cancel() + start := time.Now() + err := sleepWithContext(canceledCtx, 2*time.Millisecond) + elapsed := time.Since(start) + require.ErrorIs(t, err, context.Canceled) + assert.Less(t, elapsed, time.Second, "already-canceled context should return promptly") +} + +func TestCheckAndWaitForRateLimitContextCancelled(t *testing.T) { + oldFetchRateLimitFunc := fetchRateLimitFunc + fetchRateLimitFunc = func() (rateLimitResource, error) { + return rateLimitResource{ + Limit: 5000, + Remaining: 0, + Reset: time.Now().Add(10 * time.Minute).Unix(), + }, nil + } + t.Cleanup(func() { fetchRateLimitFunc = oldFetchRateLimitFunc }) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + start := time.Now() + err := checkAndWaitForRateLimit(ctx, false) + elapsed := time.Since(start) + + require.ErrorIs(t, err, context.Canceled) + assert.Less(t, elapsed, 100*time.Millisecond, "cancelled context should interrupt rate-limit wait promptly") +} + +func TestCheckAndWaitForRateLimitFetchErrorAndContextDone(t *testing.T) { + oldFetchRateLimitFunc := fetchRateLimitFunc + expectedFetchErr := stderrors.New("fetch failure") + fetchRateLimitFunc = func() (rateLimitResource, error) { + return rateLimitResource{}, expectedFetchErr + } + t.Cleanup(func() { fetchRateLimitFunc = oldFetchRateLimitFunc }) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := checkAndWaitForRateLimit(ctx, false) + require.Error(t, err) + require.ErrorIs(t, err, expectedFetchErr) + require.ErrorIs(t, err, context.Canceled) +}