Skip to content

fix(limits): reserve message-send quota atomically at accept - #1048

Open
tunglambk wants to merge 1 commit into
tokencanopy:mainfrom
tunglambk:fix/messages-month-quota-race
Open

tunglambk wants to merge 1 commit into
tokencanopy:mainfrom
tunglambk:fix/messages-month-quota-race

Conversation

@tunglambk

Copy link
Copy Markdown
Contributor

Summary

CheckMessageSend reads the month's outbound total from usage_summaries, but that total is only written at the terminal outcome (meterSentTx), so an accepted-but-unsent message was invisible to the cap and a burst of immediate sends could all pass against the same pre-increment count. I reproduced it against real Postgres on main: 10 concurrent POST /v1/agents/{email}/messages with max_messages_month=3 all returned 202, and three sequential sends with the cap at 2 all passed.

The accepted message row is now the reservation. usage.Store.MessagesThisMonth and MessagesToday add the units of outbound rows in accepted/sending to the metered total, so a send consumes budget the moment it is accepted, and limits.DBEnforcer.ReserveMessageSendTx takes a per-user advisory lock (keyspace 3, distinct from the domain-name/domain-owner/agent keyspaces) and re-reads the flow caps on the accept transaction's own connection before the insert — the same shape as #942's max_agents fix and #901's max_domains fix. DeliverOutbound calls it for immediate sends, and the platform test send (acceptPlatformSend) calls it too. A reservation that would cross the cap rolls the whole accept transaction back and answers 402 limit_exceeded with the same details the pre-check already produces.

Scheduled and review-held sends are deliberately excluded from the reservation: the worker's fire-time gate judges those against the month they actually fire in, which is the month the cap should apply to. The checks that reject them at fire time now see in-flight immediate sends too, which is what that gate was already reaching for.

Operational risk

No schema change and no change to the /v1 request or response shape. A send that used to slip past the cap under a burst now gets the same 402 limit_exceeded the endpoint already returns for an over-cap request.

Two behavior notes worth a look. The account usage view (GetUsage → usage.messages_month) now includes accepted-but-unsent sends, because it reads the same counter; billing reconciliation reads usage_summaries directly and this change leaves that table alone. And a send that fails before the terminal meter frees its reservation, because failed sends were already unmetered — I didn't change that refund behavior, the reservation just follows it.

The loopback self-send path is not covered. It meters synchronously in the same request rather than through the queue-first accept transaction, so it has no accept-to-send window to close. Closing it would need a reservation inside the loopback transaction; say the word if you want that here.

Client surface checklist

Not applicable: internal concurrency fix, the request/response shape of the send endpoints is unchanged.

Test plan

  • New e2e tests (internal/e2e/messages_send_quota_race_e2e_test.go, -tags integration), against a real Postgres 16 container:
    • TestConcurrentImmediateSendsRespectMonthlyQuota: 10 concurrent immediate sends with the cap at 3. On main all 10 get 202; on this branch exactly 3 get 202 (3 accepted rows) and 7 get 402 with resource=messages_month, limit=3, current=3.
    • TestImmediateSendConsumesMonthlyQuotaAtAccept: cap 2, three sequential sends. On main all three pass; on this branch the third is 402 and the first two each consume budget.
    • TestImmediateSendConsumesDailyQuotaAtAccept: the optional max_messages_day cap takes the same accept-time reservation.
  • New DB-backed counter test (internal/usage/pending_quota_test.go): accepted sends count, scheduled sends and review holds do not, and a terminal send moves out of the pending count into usage_summaries without double counting.
  • go test -tags integration -race ./internal/e2e/ -run 'QuotaAtAccept|RespectMonthlyQuota' green.
  • go test ./internal/usage/ ./internal/limits/ ./internal/agent/ ./internal/httpapi/ green, and go test ./... green.
  • Full ./internal/e2e/ suite green except TestWebhooksE2E_RotationGrace_DualSig and TestWebhooksE2E_DisabledWebhookNoFire, which need Mailpit on 127.0.0.1:1025; both fail the same way on main here.
  • scripts/check-repository-text-integrity.sh green. make fmt-check reports the same three pre-existing gofmt-drift files as main under Go 1.27 (CI pins 1.26); my files are clean.
  • Not verified: contention across more than one process, or a saturated connection pool. The lock is per-user, so it only serializes one account's own immediate accepts, and the reservation reads on the accept transaction's connection so it never waits on a second pool connection while holding the lock.

Fixes #1005

CheckMessageSend read usage_summaries, which only gets written when a send
reaches its terminal outcome in meterSentTx, so every accepted-but-unsent
message was invisible to the cap. Reproduced against real Postgres on main:
10 concurrent POST /v1/agents/{email}/messages with max_messages_month=3 all
returned 202, and three sequential sends with the cap at 2 all passed.

The accepted message row is now the reservation. MessagesThisMonth and
MessagesToday count outbound rows in accepted/sending alongside the metered
total, and the accept transaction takes a per-user advisory lock (keyspace 3)
and re-reads the flow caps on its own connection before inserting, the same
shape as tokencanopy#942 and tokencanopy#901. Scheduled and review-held sends are excluded because
the worker's existing fire-time gate judges them against the month they fire
in.

Fixes tokencanopy#1005
@tunglambk
tunglambk requested a review from jiashuoz as a code owner September 27, 2026 02:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

CheckMessageSend has no atomic increment, so concurrent immediate sends can oversell MaxMessagesMonth

1 participant