Conversation
|
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
📝 WalkthroughWalkthroughAdds a context-aware retry API and migrates several call sites to use it: introduces context-aware retry functions in the retry and db packages, updates tests for context behavior, and replaces prior db.WithRetry(...) calls with db.WithRetryContext(ctx, ...) in handlers and services. Changes
Sequence Diagram(s)sequenceDiagram
participant Caller
participant WithRetryContext as DB.WithRetryContext
participant Backoff as backoffStrategy
participant Operation as fn()
participant Ctx as ctx
Caller->>WithRetryContext: call WithRetryContext(ctx, fn)
loop up to DefaultAttempts
WithRetryContext->>Ctx: check ctx.Done()
alt ctx canceled / deadline
WithRetryContext-->>Caller: return ctx.Err()
else
WithRetryContext->>Operation: invoke fn()
Operation-->>WithRetryContext: (result, err)
alt err == nil
WithRetryContext-->>Caller: return (result, nil)
else if shouldRetryError(err)
WithRetryContext->>Backoff: compute delay
Backoff-->>WithRetryContext: delay
WithRetryContext->>Ctx: sleep delay (cancelable)
alt ctx canceled during sleep
WithRetryContext-->>Caller: return ctx.Err()
end
else
WithRetryContext-->>Caller: return (result, err)
end
end
end
WithRetryContext-->>Caller: return final (result, err)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes
Possibly related PRs
Suggested labels
Pre-merge checks and finishing touches❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
@Flo4604 we won't able to add harness because its creating circular dependency. Harness importing db, db importing harness. |
|
ah yeah thats fine then |
ogzhanolguncu
requested review from
chronark,
imeyer and
perkinsjr
as code owners
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
go/internal/services/usagelimiter/redis.go (1)
156-158: Compile error: cannot range over int; start N workers properlyfor range replayWorkers is invalid and won’t compile. Use an indexed loop.
Apply:
-for range replayWorkers { - go s.replayRequests() -} +for i := 0; i < replayWorkers; i++ { + go s.replayRequests() +}
♻️ Duplicate comments (2)
go/pkg/db/retry_test.go (1)
50-63: Skip retry on not foundAsserting IsNotFound and single call addresses prior feedback about specificity.
go/pkg/retry/retry.go (1)
225-233: Good fix: time.NewTimer + select avoids time.After leak.Thanks for switching away from time.After and stopping the timer on cancel. This addresses the leak pattern raised earlier.
🧹 Nitpick comments (2)
go/pkg/db/retry.go (1)
49-59: Make backoff scalable and reduce herd with jitterCurrent table caps at 3 and returns 50ms for n>len(delays). Prefer a scalable exponential with optional cap and jitter to spread retries.
Example:
-func backoffStrategy(n int) time.Duration { - delays := []time.Duration{DefaultBackoff, DefaultBackoff * 2, DefaultBackoff * 4} - if n <= 0 || n > len(delays) { - return DefaultBackoff - } - return delays[n-1] -} +func backoffStrategy(n int) time.Duration { + if n < 1 { + return DefaultBackoff + } + // exp backoff with cap + d := DefaultBackoff << (n - 1) // 50ms, 100ms, 200ms, 400ms, ... + if d > 2*time.Second { + d = 2 * time.Second + } + // add ±20% jitter + j := d / 5 + return d - j + time.Duration(rand.Int63n(int64(2*j+1))) +}Remember to seed a rand source or use math/rand/v2 as appropriate.
go/pkg/retry/retry.go (1)
207-234: Inject a timer factory for testability; DoContext currently bypasses r.sleep.DoContext can’t leverage r.sleep, making tests slower/flaky. Inject a timer factory (similar to sleep hook) to stub timers in tests.
Apply this diff:
diff --git a/go/pkg/retry/retry.go b/go/pkg/retry/retry.go @@ type retry struct { // overwrite time.Sleep to speed up tests sleep func(d time.Duration) + // newTimer allows injecting a fake timer in tests + newTimer func(d time.Duration) *time.Timer } @@ func New(applies ...Apply) *retry { shouldRetry: nil, // nil means all errors are retryable - sleep: time.Sleep, + sleep: time.Sleep, + newTimer: time.NewTimer, } @@ func (r *retry) DoContext(ctx context.Context, fn func() error) error { - timer := time.NewTimer(r.backoff(i)) + timer := r.newTimer(r.backoff(i)) select { case <-ctx.Done(): timer.Stop() return ctx.Err() case <-timer.C: }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (9)
go/apps/api/routes/v2_ratelimit_delete_override/handler.go(2 hunks)go/apps/api/routes/v2_ratelimit_limit/handler.go(2 hunks)go/internal/services/keys/get.go(1 hunks)go/internal/services/usagelimiter/limit.go(1 hunks)go/internal/services/usagelimiter/redis.go(1 hunks)go/pkg/db/retry.go(2 hunks)go/pkg/db/retry_test.go(13 hunks)go/pkg/retry/retry.go(2 hunks)go/pkg/retry/retry_test.go(2 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-08-21T15:54:45.198Z
Learnt from: chronark
PR: unkeyed/unkey#3825
File: go/internal/services/usagelimiter/limit.go:38-0
Timestamp: 2025-08-21T15:54:45.198Z
Learning: In go/internal/services/usagelimiter/limit.go, the UpdateKeyCreditsDecrement operation cannot be safely wrapped with db.WithRetry due to the lack of idempotency mechanisms in the current tech stack. Retrying this non-idempotent write operation risks double-charging users if the first attempt commits but the client sees a transient error.
Applied to files:
go/internal/services/usagelimiter/limit.gogo/internal/services/usagelimiter/redis.go
🧬 Code graph analysis (8)
go/internal/services/usagelimiter/limit.go (1)
go/pkg/db/retry.go (1)
WithRetryContext(36-46)
go/pkg/db/retry.go (3)
go/pkg/retry/retry.go (5)
DoWithResultContext(242-250)New(64-75)Attempts(82-87)Backoff(92-97)ShouldRetry(116-121)go/pkg/db/handle_err_no_rows.go (1)
IsNotFound(8-10)go/pkg/db/handle_err_duplicate_key.go (1)
IsDuplicateKeyError(7-13)
go/pkg/retry/retry_test.go (1)
go/pkg/retry/retry.go (4)
New(64-75)Attempts(82-87)Backoff(92-97)DoWithResultContext(242-250)
go/internal/services/usagelimiter/redis.go (1)
go/pkg/db/retry.go (1)
WithRetryContext(36-46)
go/pkg/db/retry_test.go (5)
go/pkg/db/retry.go (1)
WithRetryContext(36-46)go/pkg/retry/retry.go (1)
New(64-75)go/pkg/db/key_find_for_verification.sql_generated.go (1)
FindKeyForVerificationRow(92-119)go/pkg/uid/uid.go (1)
KeyPrefix(16-16)go/pkg/db/key_insert.sql_generated.go (1)
InsertKeyParams(51-67)
go/internal/services/keys/get.go (2)
go/pkg/db/retry.go (1)
WithRetryContext(36-46)go/pkg/db/key_find_for_verification.sql_generated.go (1)
FindKeyForVerificationRow(92-119)
go/apps/api/routes/v2_ratelimit_delete_override/handler.go (3)
go/apps/api/openapi/gen.go (2)
V2RatelimitDeleteOverrideRequestBody(1726-1739)V2RatelimitDeleteOverrideResponseBody(1742-1748)go/pkg/db/retry.go (1)
WithRetryContext(36-46)go/pkg/db/ratelimit_namespace_find.sql_generated.go (1)
FindRatelimitNamespaceRow(37-45)
go/apps/api/routes/v2_ratelimit_limit/handler.go (3)
go/apps/api/openapi/gen.go (2)
V2RatelimitLimitRequestBody(1785-1813)V2RatelimitLimitResponseBody(1816-1821)go/pkg/db/retry.go (1)
WithRetryContext(36-46)go/pkg/db/ratelimit_namespace_find.sql_generated.go (1)
FindRatelimitNamespaceRow(37-45)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
- GitHub Check: Test API / API Test Local
- GitHub Check: Build / Build
- GitHub Check: Test Go API Local / Test
🔇 Additional comments (21)
go/internal/services/usagelimiter/limit.go (1)
16-18: Good swap to context-aware retry; keep decrement non-retriedUsing db.WithRetryContext here is correct and preserves cancellation semantics. Thanks for keeping UpdateKeyCreditsDecrement outside any retry wrapper to avoid double-charging on transient failures. Based on learnings
go/internal/services/keys/get.go (1)
73-75: Context-aware DB fetch in cache loader looks gooddb.WithRetryContext(ctx, …) ensures cancellation/deadline cut through the retry loop. No further changes needed.
go/internal/services/usagelimiter/redis.go (1)
241-244: Context-aware retry on DB init is correctSwitch to db.WithRetryContext(ctx, …) is appropriate; cancellation will stop backoff and avoid hanging on shutdown paths.
go/apps/api/routes/v2_ratelimit_delete_override/handler.go (2)
26-29: Type block consolidation is fineGrouping Request/Response aliases improves readability without behavior change.
64-70: Context-aware namespace lookupdb.WithRetryContext(ctx, …) threads request ctx into retries; good alignment with the new retry API.
go/pkg/retry/retry_test.go (2)
298-389: DoContext tests cover key cancellation scenarios wellCovers pre-cancel, cancel during backoff, and deadline-exceeded with precise call counts. The ctx-aware fake sleep is a solid pattern.
391-450: DoWithResultContext semantics verifiedAsserting last successful result before ctx error is great; matches wrapper behavior. Test timings/backoff align with expectations.
go/apps/api/routes/v2_ratelimit_limit/handler.go (2)
32-35: Type block consolidation looks goodNo behavior change; improves readability.
78-86: Loader now respects context during retriesdb.WithRetryContext(ctx, …) is the right call; cancellation will stop backoff quickly.
go/pkg/db/retry_test.go (10)
18-30: Happy-path succeeds in one callBaseline test is clear; validates zero retry overhead.
32-48: Retries transient errors, then succeedsGood assertion on call count and result.
65-79: Skip retry on duplicate keySingle attempt with IsDuplicateKeyError is correct.
81-95: Exhausts retries path validatedMatches default 3 attempts and error propagation.
97-123: Generic result types coveredBoth int and struct cases validate type-safe usage.
125-190: Context cancellation unit tests look solidPre-cancel, cancel during backoff, and deadline during backoff are precisely asserted with call counts.
192-269: Integration: transient failure then successUses real sqlc ops; good realism and assertions.
271-307: Integration: duplicate key not retriedCorrectly asserts single call and duplicate detection.
309-321: Integration: not found not retriedGood single-call assertion with IsNotFound.
322-368: Integration: cancellation halts further DB opsStrong check that second insert never runs after cancel; good guard against post-cancel side effects.
go/pkg/db/retry.go (1)
36-45: Wiring to retry.DoWithResultContext is correctThe wrapper applies attempts/backoff/shouldRetry and passes ctx through; aligns with the new API.
go/pkg/retry/retry.go (1)
239-250: LGTM: DoWithResultContext mirrors DoContext semantics and returns last result.Implementation is correct and consistent with DoWithResult.
|
Thank you for following the naming conventions for pull request titles! 🙏 |
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
go/pkg/retry/retry_test.go (1)
367-374: Remove dead code: mocked sleep is never called by DoContext.The mocked sleep function is never invoked because
DoContextusestime.NewTimerwith aselectstatement directly (lines 238-246 in retry.go) rather than callingr.sleep(). This mock is dead code and could confuse future maintainers.Consider removing this mock:
- // Use fake sleep that respects context cancellation - tt.retry.sleep = func(d time.Duration) { - select { - case <-ctx.Done(): - return - case <-time.After(d): - return - } - } - if tt.cancelAfter > 0 {
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
go/pkg/retry/retry.go(2 hunks)go/pkg/retry/retry_test.go(3 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
go/pkg/retry/retry_test.go (1)
go/pkg/retry/retry.go (4)
New(65-76)Attempts(83-88)Backoff(93-98)DoWithResultContext(255-263)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (5)
- GitHub Check: Build / Build
- GitHub Check: Test Go API Local / Test
- GitHub Check: Test API / API Test Local
- GitHub Check: autofix
- GitHub Check: Analyze (javascript-typescript)
🔇 Additional comments (5)
go/pkg/retry/retry_test.go (3)
299-390: LGTM: comprehensive context cancellation test coverage.The test cases properly cover pre-cancelled contexts, cancellation during backoff, and deadline exceeded scenarios. The tests correctly verify that context cancellation interrupts retry attempts as expected.
392-451: LGTM: DoWithResultContext correctly validates result handling.The tests verify that partial results are returned when context is cancelled at various points in the retry sequence, ensuring result preservation across all cancellation scenarios.
453-499: LGTM: validates non-retryable context errors per requirements.Tests confirm that context errors (both direct and wrapped) returned by
fn()are treated as non-retryable and short-circuit immediately without backoff, addressing the linked issue objectives.go/pkg/retry/retry.go (2)
203-250: LGTM: DoContext correctly implements context-aware retry semantics.The implementation properly addresses all requirements from the linked issue and past review feedback:
- Context checked before each attempt (lines 210-213)
- Context errors from
fn()treated as non-retryable (lines 228-230)- Backoff uses
time.NewTimerwith select to avoid memory leaks (lines 238-246)- Timer properly stopped on context cancellation (line 242)
255-263: LGTM: DoWithResultContext correctly captures and returns results.The wrapper correctly delegates to
DoContextwhile preserving the last result value, matching the behavior of the non-contextDoWithResultfunction.
|
|
Graphite Automations"Post a GIF when PR approved" took an action on this PR • (10/28/25)1 gif was posted to this PR based on Andreas Thomas's automation. |
3 participants
What does this PR do?
This PR adds context aware
DoWithResult-DoWithResultContext. To make this as robust as possible we've introduced two new methods topkg/retrycalledDoContextandDoWithResultContext. Reason for introducing new methods are we need to be reactive to context changes(cancels and deadlines) with regular still its impossible to cut the flow when cancel invoked or deadlines passed.Due to duplication on
pkg/db/retry.gowe also had to make it reusable so we can share the same initialization logic.Fixes #3828
If there is not an issue for this, please create one first. This is used to tracking purposes and also helps use understand why this PR exists
Type of change
How should this be tested?
Checklist
Required
pnpm buildpnpm fmtconsole.logsgit pull origin mainAppreciated