Fix aggregate storage entitlement usage across buckets - #1010
Conversation
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
📝 WalkthroughWalkthroughStorage write entitlement checks now aggregate D1 usage and durable-object estimates across inventoried user storage buckets. Tests cover over-entitlement rejection, ownership removal, boundary writes, and retained bucket data. ChangesStorage entitlement aggregation
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant WriteAuthorization
participant StorageBucketService
participant StorageRunnerRPC
participant D1Storage
WriteAuthorization->>D1Storage: readUserD1StorageBytes
WriteAuthorization->>StorageBucketService: listUserStorageBucketIds
WriteAuthorization->>StorageRunnerRPC: getEstimatedBytes for inventoried storage ids
StorageRunnerRPC-->>WriteAuthorization: estimated durable-object bytes
D1Storage-->>WriteAuthorization: current D1 bytes
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 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 |
|
@coderabbitai review |
|
🔎 Preview deployed: https://kody-pr-1010.kody-a99.workers.dev Worker: Mocks:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
packages/worker/src/storage-runner.workers.test.ts (2)
243-296: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest math is sound but silently depends on
estimateA > 0andrawSize: 0.
targetD1Bytes - initialD1Bytesonly lands ontargetD1Bytesbecause the seeded row starts withraw_size = 0, and the denial at Line 292 only holds ifestimateA > 0. Both are true today but neither is asserted, so a future change togetEstimatedBytesor the seed helper would make this test pass vacuously. Addingexpect(estimateA).toBeGreaterThan(0)(and asserting the post-update D1 total) would make the intent explicit.♻️ Suggested assertions
const estimateA = (await runnerA.getEstimatedBytes()).estimatedBytes const estimateB = (await runnerB.getEstimatedBytes()).estimatedBytes + expect(estimateA).toBeGreaterThan(0) + expect(estimateB).toBeGreaterThan(0) const initialD1Bytes = await readUserD1StorageBytes({ db: env.APP_DB, userId, }) const targetD1Bytes = limit - estimateB - 1 await env.APP_DB.prepare( `UPDATE email_messages SET raw_size = ? WHERE user_id = ?`, ) .bind(targetD1Bytes - initialD1Bytes, userId) .run() + await expect( + readUserD1StorageBytes({ db: env.APP_DB, userId }), + ).resolves.toBe(targetD1Bytes)🤖 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 `@packages/worker/src/storage-runner.workers.test.ts` around lines 243 - 296, Add explicit assertions in the storage entitlement test after calculating estimateA and after updating email_messages: assert estimateA is greater than zero, and read/assert the resulting D1 storage total equals targetD1Bytes. Keep the existing raw_size: 0 setup and aggregate denial assertions unchanged.
277-291: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer Vitest's rejection matchers over manual
.then(null, capture).
await expect(...).rejects.toBeInstanceOf(EntitlementLimitError)plus a captured error would express the same intent without the hand-rolled sentinel andthrow new Error(...).🤖 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 `@packages/worker/src/storage-runner.workers.test.ts` around lines 277 - 291, Replace the manual rejection capture and instanceof check around assertStorageRunnerWriteWithinEntitlement with Vitest’s expect(...).rejects.toBeInstanceOf(EntitlementLimitError) matcher. Remove the aggregateDenied sentinel and custom throw while preserving the existing entitlement-rejection assertion.packages/worker/src/storage-runner.ts (1)
597-621: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy liftPer-write DO fan-out now scales with bucket count.
Every storage write triggers one inventory query plus
NDO round-trips, serialized across batches of 16. For users who accumulate manyexec:*buckets this becomes the dominant latency of the write path. Consider a short-TTL memo of the aggregate (keyed byuserId) or persisting per-bucket estimates inuser_storage_bucketsso the entitlement check reads D1 only.🤖 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 `@packages/worker/src/storage-runner.ts` around lines 597 - 621, The storage write path currently performs a DO estimate RPC for every registered storage bucket; reduce this fan-out by caching the aggregate durable-object byte estimate with a short TTL keyed by userId, or persist per-bucket estimates in user_storage_buckets and read the aggregate from D1. Update the code surrounding durableObjectBytes and preserve the entitlement calculation while avoiding per-write DO round-trips.
🤖 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 `@packages/worker/src/storage-runner.ts`:
- Around line 606-616: Update the concurrent estimate reads in the getCurrent
flow around storageRunnerRpc().getEstimatedBytes() so one failing StorageRunner
does not reject the entire Promise.all operation. Use Promise.allSettled and
apply each unreadable bucket’s last known estimate, or, if the intended policy
is fail-closed, catch and rethrow an explicit entitlement-related error;
preserve successful estimates and the existing batching behavior.
---
Nitpick comments:
In `@packages/worker/src/storage-runner.ts`:
- Around line 597-621: The storage write path currently performs a DO estimate
RPC for every registered storage bucket; reduce this fan-out by caching the
aggregate durable-object byte estimate with a short TTL keyed by userId, or
persist per-bucket estimates in user_storage_buckets and read the aggregate from
D1. Update the code surrounding durableObjectBytes and preserve the entitlement
calculation while avoiding per-write DO round-trips.
In `@packages/worker/src/storage-runner.workers.test.ts`:
- Around line 243-296: Add explicit assertions in the storage entitlement test
after calculating estimateA and after updating email_messages: assert estimateA
is greater than zero, and read/assert the resulting D1 storage total equals
targetD1Bytes. Keep the existing raw_size: 0 setup and aggregate denial
assertions unchanged.
- Around line 277-291: Replace the manual rejection capture and instanceof check
around assertStorageRunnerWriteWithinEntitlement with Vitest’s
expect(...).rejects.toBeInstanceOf(EntitlementLimitError) matcher. Remove the
aggregateDenied sentinel and custom throw while preserving the existing
entitlement-rejection assertion.
🪄 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: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 558961a3-8237-42d1-a512-d7d38ddb552a
📒 Files selected for processing (2)
packages/worker/src/storage-runner.tspackages/worker/src/storage-runner.workers.test.ts
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
3 participants
Summary
storage_bytesTesting
npm run test:workers -- --run packages/worker/src/storage-runner.workers.test.ts(8 tests passed)npm run validate(passed before and after AI review fixes)System recap — extends an existing primitive (medium risk)
Mode: recap · Base:
main@6620d1b3· Head:69ff2f95Classification: extends — storage entitlement enforcement now aggregates every bucket in the per-user durable storage inventory.
Primitives touched
durable-storageSystem map
A storage write reads the acting user's bucket inventory, obtains bounded-concurrency estimates from those user-scoped Durable Objects, and combines them with D1 usage before enforcement.
Legend: green = composes (wiring only) · amber = extended by this PR · red = new primitive · gray = context (unchanged, included only when an edge crosses it).
Invariants
userId.Conductor Report
STATUS: done
What changed: storage-byte entitlement checks now sum D1 bytes plus every inventoried per-user StorageRunner bucket, with the current bucket included and deduplicated during registration races. Durable Object estimate reads run in batches of at most 16 and fail closed with an explicit verification error if unreadable.
Tests run: focused Workers suite passed (8/8).
npm run validatepassed twice, including after AI review fixes; formatting, lint, typecheck, Node/Workers tests, Playwright E2E, MCP E2E, backup build, primitives, and migrations checks are green.Blocker: none.
Scope spill: none.
Summary by CodeRabbit