feat(retry): quarantine per-(chat, requester) inbound retry receipts - #982
feat(retry): quarantine per-(chat, requester) inbound retry receipts#982Salientekill wants to merge 4 commits into
Conversation
In large groups a cohort of members whose pairwise sessions never establish (dead registrations, exhausted prekeys, re-registered LIDs) sends a retry receipt for every single group message. Each receipt pays markForgetSenderKey (a DB write plus a whole-group sender-key cache invalidation), bundle processing and possibly a resend - all upstream of the per-chat resend cap (oxidezap#871), which only bounds the resend itself. Observed in production: ~58k "Marked ... for fresh SKDM" in 3.5 days from 468 distinct members of a single 1012-participant group (~99.7% of all marks), with the per-chat cap blocking ~75% of the resends but none of the repair work. This bounds the whole repair path per member with a token bucket keyed by (chat user, requester user): burst 2, refill 2/day. One mark is enough to repair a healthy member (the next send carries the SKDM), so past the burst further receipts from the same pair are dropped before any work happens; the refill keeps genuine recovery possible. Device is excluded from the key on purpose - WA Web re-targets the whole user when the primary goes cold, so all devices of a broken account share one budget. - RetryMarkQuarantine reuses the existing TokenBucket/Cache machinery from the resend limiter (lazy monotonic refill, capacity-only cache) - burst 0 disables; live-tunable via Client::set_retry_mark_quarantine - observability: StatsSnapshot.retry_receipts_quarantined counter and retry_mark_quarantine_pairs in MemoryDiagnostics
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI (base), Organization UI (inherited) Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR adds a per-(chat, requester) retry-receipt quarantine, wires it into client setup and retry handling, and exports quarantine counters through stats and memory reporting. ChangesRetry receipt quarantine
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Source as Retry receipt source
participant Handler as handle_retry_receipt
participant Quarantine as RetryMarkQuarantine
participant Repair as session repair/resend path
Source->>Handler: group/status retry receipt
Handler->>Quarantine: try_acquire(chat, requester)
alt permit granted
Quarantine-->>Handler: true
Handler->>Repair: continue repair/resend logic
Repair-->>Handler: result
else throttled
Quarantine-->>Handler: false
Handler-->>Source: Ok(()) with debug log
end
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
|
| Filename | Overview |
|---|---|
| src/resend_rate_limiter.rs | Adds RetryMarkQuarantine struct reusing the existing TokenBucket/Cache machinery, with two new constants, correct set_rate/try_acquire/observability methods, and thorough unit + concurrency tests. |
| src/retry.rs | Inserts the quarantine gate after the message-cache lookup and the is_peer check; correctly protects the group-info fetch, rotateKey path, and markForgetSenderKey. Minor inefficiency: resolve_encryption_jid still runs before the gate for all quarantined pairs, though resolved_jid is not consumed until after it. Inline comment violates the project's explain-why-not-what style rule. |
| src/cache_config.rs | Adds a dedicated retry_mark_quarantine_capacity field (default 32,768) with correct wiring into Debug and Default; capacity comment accurately explains the O(groups x storm-members) key space. |
| src/client.rs | Adds retry_mark_quarantine: RetryMarkQuarantine field to Client and retry_mark_quarantine_pairs: u64 to MemoryReport with correct Display formatting. |
| src/client/accessors.rs | Adds set_retry_mark_quarantine public tuning method and wires retry_receipts_quarantined into the stats snapshot and retry_mark_quarantine_pairs into the memory report. |
| src/client/lifecycle.rs | Initializes retry_mark_quarantine with the dedicated retry_mark_quarantine_capacity config field and the two new default constants; previous concern about shared capacity no longer applies. |
| wacore/src/stats.rs | Adds retry_receipts_quarantined: u64 to StatsSnapshot with a correct zero initializer in SessionStats::snapshot(). |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[handle_retry_receipt] --> B{retry_count >= MAX?}
B -- yes --> Z1[drop]
B -- no --> C[pending_retries dedup check]
C -- duplicate in-flight --> Z2[drop]
C -- new --> D[has_device? schedule_unknown_device_sync if not]
D --> E[message cache lookup]
E -- not found --> Z3[drop]
E -- found --> F[resolve_encryption_jid]
F --> G[should_drop_unknown_device_retry?]
G -- yes --> Z4[drop]
G -- no --> H[compute is_peer]
H --> I{is_group_or_status and not is_peer?}
I -- no --> J[proceed]
I -- yes --> K[RetryMarkQuarantine.try_acquire]
K -- token available --> J
K -- budget exhausted --> Z5[quarantine drop retry_receipts_quarantined++]
J --> L[fetch group info]
L --> M{rotateKey path unknown participant?}
M -- yes --> N[clear_sender_key_devices DB write + cache invalidate]
N --> O[update_local_signal_session markForgetSenderKey + processKeyBundle]
M -- no --> O
O --> P[resend_rate_limiter.try_acquire]
P -- throttled --> Z6[drop resend]
P -- allowed --> Q[resend message]
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
A[handle_retry_receipt] --> B{retry_count >= MAX?}
B -- yes --> Z1[drop]
B -- no --> C[pending_retries dedup check]
C -- duplicate in-flight --> Z2[drop]
C -- new --> D[has_device? schedule_unknown_device_sync if not]
D --> E[message cache lookup]
E -- not found --> Z3[drop]
E -- found --> F[resolve_encryption_jid]
F --> G[should_drop_unknown_device_retry?]
G -- yes --> Z4[drop]
G -- no --> H[compute is_peer]
H --> I{is_group_or_status and not is_peer?}
I -- no --> J[proceed]
I -- yes --> K[RetryMarkQuarantine.try_acquire]
K -- token available --> J
K -- budget exhausted --> Z5[quarantine drop retry_receipts_quarantined++]
J --> L[fetch group info]
L --> M{rotateKey path unknown participant?}
M -- yes --> N[clear_sender_key_devices DB write + cache invalidate]
N --> O[update_local_signal_session markForgetSenderKey + processKeyBundle]
M -- no --> O
O --> P[resend_rate_limiter.try_acquire]
P -- throttled --> Z6[drop resend]
P -- allowed --> Q[resend message]
Reviews (4): Last reviewed commit: "fix(retry): exempt own companion devices..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/client/lifecycle.rs`:
- Around line 197-202: `RetryMarkQuarantine` is currently reusing
`resend_rate_limiter_capacity`, which couples it to the per-chat bucket budget.
Add a dedicated capacity setting for the quarantine cache in the
lifecycle/config path, wire it through `Lifecycle` initialization, and use that
new knob when constructing `RetryMarkQuarantine::new` instead of
`cache_config.resend_rate_limiter_capacity`. Keep the existing burst/refill
settings unchanged and update the relevant config structure/constructor so the
quarantine capacity can be tuned independently.
In `@src/resend_rate_limiter.rs`:
- Around line 231-252: Add a multi-threaded concurrency test for
RetryMarkQuarantine similar to the existing ResendRateLimiter contention test.
In the tests module for RetryMarkQuarantine, create a tokio multi_thread test
that spawns many concurrent try_acquire calls against the same Jid pair using
Arc and an atomic counter, then assert the allowed count never exceeds the
configured burst. Use the existing RetryMarkQuarantine::new and try_acquire
symbols so the race behavior is verified under real contention.
In `@src/retry.rs`:
- Around line 412-431: Move the quarantine gate earlier in the retry handling
flow so quarantined (chat, requester) pairs skip the expensive work before it
runs. The current check on retry_mark_quarantine only bypasses
update_local_signal_session, while query_info() and the unknown-device rotation
path still execute; refactor the surrounding logic so the try_acquire() guard in
this retry handler returns before those DB/cache operations, or adjust the
comment if the intended scope is only the later session update.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI (base), Organization UI (inherited)
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: a597b675-2074-4431-981f-31353d651cbc
📒 Files selected for processing (6)
src/client.rssrc/client/accessors.rssrc/client/lifecycle.rssrc/resend_rate_limiter.rssrc/retry.rswacore/src/stats.rs
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
3 issues found across 6 files
Confidence score: 3/5
- In
src/resend_rate_limiter.rs(try_acquire), allocating twoStrings per receipt before cache lookup adds allocator churn on every retry event, including cache hits; during receipt storms this can raise CPU/latency and blunt the limiter’s intended protection — avoid per-call allocations (e.g., reuse/borrow keys) before merging. - In
src/client/lifecycle.rs,RetryMarkQuarantineappears keyed by(chat, requester)but initialized fromresend_rate_limiter_capacity(typically per-chat), which can evict active pair buckets in large-group storms and let retries slip through sooner than intended — size this cache for pair cardinality or split capacity settings before merging. src/resend_rate_limiter.rscurrently has only single-threaded coverage for the new quarantine limiter, so concurrent acquires for the same(chat, requester)remain unverified and could hide burst-enforcement regressions — add a multi-threaded contention test to de-risk behavior under load.
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="src/resend_rate_limiter.rs">
<violation number="1" location="src/resend_rate_limiter.rs:202">
P2: `try_acquire` allocates two `String`s on every inbound retry receipt before the cache lookup. In the high-volume receipt storms this feature targets, that introduces per-receipt allocator churn even on cache hits, contradicting the 'allocation only on miss' design principle applied to the nearby `ResendRateLimiter`. Consider structuring the cache lookup so the key can be passed by reference (e.g., using a borrowed key type or equivalent lookup) to avoid repeated allocations for the same (chat, requester) pair.</violation>
</file>
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
| let burst = burst as f64; | ||
| let refill_per_sec = self.refill_per_day.load(Ordering::Relaxed) as f64 / 86_400.0; | ||
|
|
||
| let key = (chat.user.to_string(), requester.user.to_string()); |
There was a problem hiding this comment.
P2: try_acquire allocates two Strings on every inbound retry receipt before the cache lookup. In the high-volume receipt storms this feature targets, that introduces per-receipt allocator churn even on cache hits, contradicting the 'allocation only on miss' design principle applied to the nearby ResendRateLimiter. Consider structuring the cache lookup so the key can be passed by reference (e.g., using a borrowed key type or equivalent lookup) to avoid repeated allocations for the same (chat, requester) pair.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At src/resend_rate_limiter.rs, line 202:
<comment>`try_acquire` allocates two `String`s on every inbound retry receipt before the cache lookup. In the high-volume receipt storms this feature targets, that introduces per-receipt allocator churn even on cache hits, contradicting the 'allocation only on miss' design principle applied to the nearby `ResendRateLimiter`. Consider structuring the cache lookup so the key can be passed by reference (e.g., using a borrowed key type or equivalent lookup) to avoid repeated allocations for the same (chat, requester) pair.</comment>
<file context>
@@ -137,11 +137,120 @@ impl ResendRateLimiter {
+ let burst = burst as f64;
+ let refill_per_sec = self.refill_per_day.load(Ordering::Relaxed) as f64 / 86_400.0;
+
+ let key = (chat.user.to_string(), requester.user.to_string());
+ let bucket = self
+ .buckets
</file context>
- cache_config.retry_mark_quarantine_capacity (default 32768): the quarantine keyspace is O(groups x broken members), not O(chats), and evicting an ACTIVE pair refunds its burst - sharing the resend limiter's per-chat capacity (4096) undersized it (greptile P1) - gate moved before the group-info fetch and the rotateKey path, so a quarantined pair no longer pays query_info nor the unknown-sender rotateKey (own sender-key deletion + whole-group invalidation) at full rate; still after the message-cache lookup so expired-message no-ops don't burn the repair budget (greptile P2) - documented the owned-key allocation tradeoff in try_acquire
|
Addressed the review in the latest commit:
|
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/resend_rate_limiter.rs (1)
100-127: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winWe're literally running the same rate-limiter code twice — let's connect these, not duplicate them.
RetryMarkQuarantine::try_acquire(Lines 194-222) is nearly a byte-for-byte copy ofResendRateLimiter::try_acquire(Lines 100-127): same burst-check short-circuit, same get-or-create-bucket dance, same lock+try_take+throttled_total bump. Only the key type (Jidvs(String, String)) and the refill divisor (60.0 vs 86,400.0) differ. Two copies of concurrency-sensitive token-bucket logic means two places that can silently drift out of sync the next time someone tunes this. At our scale, that's not a hypothetical — it's a "why did burst=0 behave differently in prod" bug waiting to happen.Consider parameterizing a single generic
TokenBucketLimiter<K>over the key type and refill-seconds divisor, then have bothResendRateLimiterandRetryMarkQuarantinewrap it.♻️ Sketch of a shared generic limiter
-pub(crate) async fn try_acquire(&self, chat: &Jid) -> bool { - let burst = self.burst.load(Ordering::Relaxed); - if burst == 0 { return true; } - ... -} +struct TokenBucketLimiter<K: Hash + Eq + Clone + Send + Sync + 'static> { + buckets: Cache<K, Arc<Mutex<TokenBucket>>>, + burst: AtomicU32, + refill_per_window: AtomicU32, + window_secs: f64, + throttled_total: AtomicU64, +} + +impl<K: Hash + Eq + Clone + Send + Sync + 'static> TokenBucketLimiter<K> { + async fn try_acquire_with(&self, key: K) -> bool { /* shared logic once */ } +}Also applies to: 174-233
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/resend_rate_limiter.rs` around lines 100 - 127, The rate-limiter logic is duplicated between ResendRateLimiter::try_acquire and RetryMarkQuarantine::try_acquire, so refactor the shared token-bucket flow into a single reusable limiter. Introduce a generic TokenBucketLimiter<K> (or equivalent shared helper) that encapsulates the burst check, get-or-create bucket, mutex-protected try_take, and throttled_total update, and parameterize only the key type and refill-per-second calculation. Then have both try_acquire methods delegate to that shared implementation instead of maintaining two copies.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@src/resend_rate_limiter.rs`:
- Around line 100-127: The rate-limiter logic is duplicated between
ResendRateLimiter::try_acquire and RetryMarkQuarantine::try_acquire, so refactor
the shared token-bucket flow into a single reusable limiter. Introduce a generic
TokenBucketLimiter<K> (or equivalent shared helper) that encapsulates the burst
check, get-or-create bucket, mutex-protected try_take, and throttled_total
update, and parameterize only the key type and refill-per-second calculation.
Then have both try_acquire methods delegate to that shared implementation
instead of maintaining two copies.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI (base), Organization UI (inherited)
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: f0e5f488-4203-49c8-a2e8-8cb3d78b3ed4
📒 Files selected for processing (4)
src/cache_config.rssrc/client/lifecycle.rssrc/resend_rate_limiter.rssrc/retry.rs
There was a problem hiding this comment.
0 issues found across 4 files (changes from recent commits).
Requires human review: Auto-approval blocked by 2 unresolved issues from previous reviews.
Re-trigger cubic
Mirrors the ResendRateLimiter contention test: 40 concurrent receipts for the same (chat, requester) pair, asserting exactly burst pass and the rest are throttled — proving burst enforcement holds under the concurrent-receipt storm the quarantine guards (coderabbit/cubic).
|
Second round — thanks @coderabbitai and @cubic-dev-ai:
|
@Salientekill Looks good across the board. A few thoughts: Concurrency test — solid addition. Mirroring the existing Capacity — the dedicated Key allocation — I buy your reasoning. The tradeoff is cleanly stated (0.19 receipts/s vs. a DB write + whole-group invalidation), and One thing to consider: greptile's point about if is_group_or_status && !is_peer && !self.retry_mark_quarantine.try_acquire(...).awaitWorth deciding whether companion devices share or have their own budget, and either adding the guard or documenting why sharing is intentional. Tip: get faster answers by chatting with cubic’s review copilot in the review UI |
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
fd resend_rate_limiter.rsLength of output: 203 🏁 Script executed: #!/bin/bash
rg -n "quarantine_concurrent_acquires_for_one_pair_do_not_exceed_burst" -A 40 src/resend_rate_limiter.rsLength of output: 1828
Capacity fix confirmed from the prior round, no further comment there. On the key allocation call — I'll weigh in since you asked. Your math checks out: ~0.19 receipts/s even at the pathological volume that motivated this PR, two small string allocs, and only on the branch that's already replacing a DB write plus whole-group cache invalidation. That's not a hot path, it's a rare-event path with a much bigger cost already baked in. Building a custom key newtype purely to enable a borrowed lookup here would be optimizing something that isn't the bottleneck — that's the kind of complexity-for-complexity's-sake I don't want in the codebase. Keep it as documented. If someone profiles this in production and it actually shows up, that's the time to revisit, not now. Nice work closing the loop on all three points. If you found this review helpful, would you consider giving us a shout-out on X? Thank you for using CodeRabbit! |
There was a problem hiding this comment.
0 issues found across 1 file (changes from recent commits).
Requires human review: Auto-approval blocked by 1 unresolved issue from previous reviews.
Re-trigger cubic
is_peer (our own account's secondary devices) was subject to the quarantine gate: a companion that was offline or rotated keys can legitimately need many group-session repair cycles, and dropping those after the burst would block group decryption on our own device. The storm this guards is third-party members, never our own. (greptile)
|
Good catch @greptile-apps — that was a real regression. Fixed in 32056d6: added |
|
The fix is correctly applied — the updated gate in the PR diff already reads: if is_group_or_status
&& !is_peer
&& !self
.retry_mark_quarantine
.try_acquire(&info.chat, &info.requester)
.awaitThe exemption is the right call. The inline comment explaining the exemption rationale is also clear and sufficient. The regression is fully addressed. No further changes needed on this. Tip: You can customize Greptile's behavior for this repo with |
There was a problem hiding this comment.
0 issues found across 1 file (changes from recent commits).
Requires human review: Auto-approval blocked by 1 unresolved issue from previous reviews.
Re-trigger cubic
|
Thanks for the very thorough work here @Salientekill, and especially for the production data. The diagnosis is correct and the engineering (token bucket reuse, the contention test, the I want to be transparent about where I landed. I cross-checked the whole retry path against the captured WA Web bundle, and the blocker for me is parity rather than code quality:
So a per-(chat, requester) quarantine that drops eligible repair requests is a deliberate divergence from WA Web. That's a perfectly reasonable trade-off for a bot at your scale, but not something I want the library doing for everyone by default. Rather than reject the idea, I opened #985 as a compliant, decoupled version: an opt-in I also folded in two things from reviewing your PR: the per-receipt Would you be up for trying #985 in your setup? Register Generated by Claude Code |
|
Superseded by #985 — @jlucaso1's opt-in |
3 participants
Problem
In large groups, a cohort of members whose pairwise sessions never establish (dead registrations, exhausted prekeys, re-registered LIDs) sends a retry receipt for every single group message. Each receipt pays
markForgetSenderKey(a DB write + a whole-group sender-key cache invalidation), bundle processing and possibly a resend — all upstream of the per-chat resend cap (#871), which only bounds the resend itself.Production evidence (bot in a 1012-participant group, 3.5 days of debug logs):
Marked <jid> for fresh SKDM ... due to retry receiptlines — 99.7% from this one group, coming from 468 distinct members (~68 marks/member/day ≈ one per bot message: they fail on every message, i.e. their sessions never repair).handle_retry_receipt.has_key=true, including devices whose distribution failed (406 / no bundle) —wacore/src/send/group.rs— relying on the retry-receipt path to repair, which produces a permanent warm/cold oscillation for this cohort. (A follow-up PR could add a per-device re-target cooldown; this PR bounds the receipt side.)Mechanism
RetryMarkQuarantine: a token bucket keyed by (chat user, requester user) — burst 2, refill 2/day. One mark is enough to repair a healthy member (the next send carries the SKDM), so past the burst, further receipts from the same pair are dropped before any repair work; the refill keeps genuine recovery possible. Device is excluded from the key on purpose: WA Web re-targets the whole user when the primary goes cold, so all devices of a broken account share one budget.With the observed numbers this cuts the repair-path work from ~16.7k/day to ~470×3/day worst case (~97%), while leaving first-time repairs untouched.
Implementation
TokenBucket+ capacity-onlyCachemachinery from the resend limiter (lazy monotonic refill, no timers, allocation only on miss).handle_retry_receiptright beforeupdate_local_signal_session, after the cheap early-outs and the concurrency guard; early return releases the pending-retry scopeguard normally.burst = 0disables; live-tunable viaClient::set_retry_mark_quarantine(burst, refill_per_day).StatsSnapshot.retry_receipts_quarantined+retry_mark_quarantine_pairsinMemoryDiagnostics.Relation to #924: that PR bounds outbound retry receipts (our own decrypt failures); this one bounds the inbound direction, which is where the group-storm cost actually lands.