feat(api): add org-wide box quota and concurrency limit - #1028
Conversation
📦 BoxLite review — 1 issue ·
|
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds organization quota persistence and enforcement with Redis-backed reservations and usage tracking. Box creation, start, proxy auto-resume, and resize flows reserve or roll back quota usage. Apps API tests are added to CI and pre-commit path detection. ChangesOrganization quota enforcement
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant BoxService
participant OrganizationUsageService
participant Redis
participant Database
BoxService->>OrganizationUsageService: validate quota and reserve pending usage
OrganizationUsageService->>Redis: read cached current and pending usage
OrganizationUsageService->>Database: load quota limits or refresh usage
OrganizationUsageService-->>BoxService: return reservation
BoxService->>OrganizationUsageService: finalize lifecycle event or roll back
OrganizationUsageService->>Redis: transfer pending usage to current or decrement pending usage
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
apps/api/src/organization/services/organization-usage.service.integration.spec.ts (1)
74-79: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest name claims concurrency but reserves run sequentially.
Both
validateOrganizationQuotascalls areawait-ed in sequence, so this exercises pending-counter accumulation, not concurrent execution. The reservation design's key property is that the reserve-then-check window is safe under true concurrency; issuing the two reserves viaPromise.all(and asserting at most one succeeds) would actually cover that path.🤖 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 `@apps/api/src/organization/services/organization-usage.service.integration.spec.ts` around lines 74 - 79, The test named “serializes concurrent reserves: the first takes the last slot, the second is rejected” currently awaits reservations sequentially; update it to start both validateOrganizationQuotas calls concurrently and await them together, asserting that exactly one succeeds and the other rejects with the CPU-limit error. Preserve the pending CPU assertion to verify only the successful reservation remains held.
🤖 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 @.github/workflows/test.yml:
- Around line 299-318: Add job-level permissions for the apps job, granting only
contents: read, and configure the checkout step to use persist-credentials:
false. Update the jobs.apps definition and its actions/checkout step so
repository-controlled install and test code cannot reuse a persisted or broader
GitHub token.
- Around line 23-24: Update the path filters in the workflow trigger sections of
test.yml to match all changes under apps/, including dependency, lockfile, and
configuration files, instead of limiting matches to apps/api/** and
apps/libs/**; keep the existing Apps test workflow behavior unchanged for
matching changes.
In `@apps/api/src/box/services/box.service.ts`:
- Around line 1080-1089: Move the box state transition that calls updateWhere()
into the same try block as the quota reservation and subsequent resize flow,
ensuring any failure after validateResizeQuota() triggers the existing
reservation rollback path. Preserve the current RESIZING transition behavior and
use the existing error handling rather than adding a separate rollback
mechanism.
In `@apps/api/src/organization/services/box-usage.ts`:
- Around line 25-32: Update boxUsageContribution to accept sufficient lifecycle
context and treat RESIZING boxes headed to STARTED as compute-consuming, while
cold resizes continue contributing zero compute resources; preserve existing
disk-state handling. In apps/api/src/organization/services/box-usage.ts#L25-L32,
update the compute-state decision using the established resize context symbols.
In apps/api/src/organization/services/box-usage.spec.ts#L15-L35, add assertions
covering both hot-resize and cold-resize contribution behavior.
---
Nitpick comments:
In
`@apps/api/src/organization/services/organization-usage.service.integration.spec.ts`:
- Around line 74-79: The test named “serializes concurrent reserves: the first
takes the last slot, the second is rejected” currently awaits reservations
sequentially; update it to start both validateOrganizationQuotas calls
concurrently and await them together, asserting that exactly one succeeds and
the other rejects with the CPU-limit error. Preserve the pending CPU assertion
to verify only the successful reservation remains held.
🪄 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: c3618019-4254-4c61-bf7f-9c4f7cddad56
📒 Files selected for processing (16)
.github/workflows/test.yml.pre-commit-config.yamlapps/api/src/box/services/box.service.spec.tsapps/api/src/box/services/box.service.tsapps/api/src/boxlite-rest/box-auto-resume.service.spec.tsapps/api/src/boxlite-rest/boxlite-ws-proxy.service.spec.tsapps/api/src/migrations/pre-deploy/1784260000000-add-organization-quota-migration.tsapps/api/src/organization/constants/box-consuming-states.constant.tsapps/api/src/organization/entities/organization-quota.entity.tsapps/api/src/organization/organization.module.tsapps/api/src/organization/services/box-usage.spec.tsapps/api/src/organization/services/box-usage.tsapps/api/src/organization/services/org-quota.spec.tsapps/api/src/organization/services/org-quota.tsapps/api/src/organization/services/organization-usage.service.integration.spec.tsapps/api/src/organization/services/organization-usage.service.ts
9603b8d to
fc24f2b
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
.github/workflows/test.yml (1)
307-312: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winRemove the recursive submodule checkout from the apps API test job.
This job only runs NestJS/Jest unit tests under
apps/api, andapps/apionly references thelibkrunname in comments/comms—not submodule paths. Use a plainactions/checkoutwithpersist-credentials: falseto avoid unnecessary submodule work.🤖 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 @.github/workflows/test.yml around lines 307 - 312, Update the checkout step in the apps API test job to remove the recursive submodule configuration, while retaining actions/checkout and persist-credentials: false. Keep the rest of the job unchanged.
🤖 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 @.github/workflows/test.yml:
- Around line 296-302: Add a Redis container health check to the services.redis
configuration, using redis-cli ping with suitable interval, timeout, and retry
settings so the apps job waits until Redis accepts connections before running
tests. Preserve the existing image, port mapping, and REDIS_HOST configuration.
---
Nitpick comments:
In @.github/workflows/test.yml:
- Around line 307-312: Update the checkout step in the apps API test job to
remove the recursive submodule configuration, while retaining actions/checkout
and persist-credentials: false. Keep the rest of the job unchanged.
🪄 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: 36e5edb3-c794-450c-8f9f-4675a284f879
📒 Files selected for processing (16)
.github/workflows/test.yml.pre-commit-config.yamlapps/api/src/box/services/box.service.spec.tsapps/api/src/box/services/box.service.tsapps/api/src/boxlite-rest/box-auto-resume.service.spec.tsapps/api/src/boxlite-rest/boxlite-ws-proxy.service.spec.tsapps/api/src/migrations/pre-deploy/1784260000000-add-organization-quota-migration.tsapps/api/src/organization/constants/box-consuming-states.constant.tsapps/api/src/organization/entities/organization-quota.entity.tsapps/api/src/organization/organization.module.tsapps/api/src/organization/services/box-usage.spec.tsapps/api/src/organization/services/box-usage.tsapps/api/src/organization/services/org-quota.spec.tsapps/api/src/organization/services/org-quota.tsapps/api/src/organization/services/organization-usage.service.integration.spec.tsapps/api/src/organization/services/organization-usage.service.ts
🚧 Files skipped from review as they are similar to previous changes (14)
- apps/api/src/boxlite-rest/box-auto-resume.service.spec.ts
- apps/api/src/organization/constants/box-consuming-states.constant.ts
- apps/api/src/organization/services/org-quota.ts
- apps/api/src/migrations/pre-deploy/1784260000000-add-organization-quota-migration.ts
- apps/api/src/organization/organization.module.ts
- apps/api/src/organization/services/org-quota.spec.ts
- apps/api/src/organization/services/box-usage.spec.ts
- apps/api/src/boxlite-rest/boxlite-ws-proxy.service.spec.ts
- apps/api/src/organization/entities/organization-quota.entity.ts
- apps/api/src/organization/services/box-usage.ts
- apps/api/src/box/services/box.service.spec.ts
- apps/api/src/organization/services/organization-usage.service.integration.spec.ts
- apps/api/src/box/services/box.service.ts
- apps/api/src/organization/services/organization-usage.service.ts
| private async excludeBoxFromUsage(overview: BoxUsageOverview, excludeBoxId: string): Promise<BoxUsageOverview> { | ||
| const box = await this.boxRepository.findOne({ where: { id: excludeBoxId } }) | ||
| if (!box) { | ||
| return overview | ||
| } | ||
|
|
||
| const contribution = boxUsageContribution(box) | ||
| return { | ||
| ...overview, | ||
| cpu: Math.max(0, overview.cpu - contribution.cpu), | ||
| memory: Math.max(0, overview.memory - contribution.memory), | ||
| disk: Math.max(0, overview.disk - contribution.disk), | ||
| gpu: Math.max(0, overview.gpu - contribution.gpu), | ||
| count: Math.max(0, overview.count - contribution.count), | ||
| } | ||
| } |
There was a problem hiding this comment.
excludeBoxFromUsage (line 211-226) subtracts a STOPPED box's already-counted disk from the projected 'current' total, while incrementPendingBoxUsage (line 320-349, reserveDisk=false at line 338 since the box is already in BOX_STATES_CONSUMING_DISK) never adds it back to 'pending'; net effect is that box's disk is counted zero times instead of once during start()/ensureStartedForProxy()'s quota check (box.service.ts:858,906), so a disk-quota violation caused by that box can silently pass.
fc24f2b to
ce1b6ce
Compare
| @OnEvent(BoxEvents.STATE_UPDATED) | ||
| async handleBoxStateUpdated(event: BoxStateUpdatedEvent): Promise<void> { | ||
| // Transitions into or out of a conditionally-consuming state (resize) are | ||
| // accounted for by the resize flow, not by simple set membership here. | ||
| if ( | ||
| BOX_STATES_CONDITIONALLY_CONSUMING_COMPUTE.includes(event.oldState) || | ||
| BOX_STATES_CONDITIONALLY_CONSUMING_COMPUTE.includes(event.newState) | ||
| ) { | ||
| return | ||
| } | ||
|
|
||
| const box = event.box | ||
| const lockKey = `box:${box.id}:quota-usage-update` | ||
| await this.redisLockProvider.waitForLock(lockKey, 60) | ||
| try { | ||
| // Warm-pool assignment re-emits STARTED -> STARTED to attribute an already | ||
| // running box to its new organization; the membership deltas would be zero. | ||
| if (event.oldState === event.newState && event.newState === BoxState.STARTED) { | ||
| await this.updateCurrentQuotaUsage(box.organizationId, 'cpu', box.cpu) | ||
| await this.updateCurrentQuotaUsage(box.organizationId, 'memory', box.mem) | ||
| await this.updateCurrentQuotaUsage(box.organizationId, 'disk', box.disk) | ||
| await this.updateCurrentQuotaUsage(box.organizationId, 'gpu', box.gpu) | ||
| await this.updateCurrentQuotaUsage(box.organizationId, 'count', 1) | ||
| return | ||
| } | ||
|
|
||
| const cpuDelta = stateTransitionDelta(box.cpu, event.oldState, event.newState, BOX_STATES_CONSUMING_COMPUTE) | ||
| const memoryDelta = stateTransitionDelta(box.mem, event.oldState, event.newState, BOX_STATES_CONSUMING_COMPUTE) | ||
| const diskDelta = stateTransitionDelta(box.disk, event.oldState, event.newState, BOX_STATES_CONSUMING_DISK) | ||
| const gpuDelta = stateTransitionDelta(box.gpu, event.oldState, event.newState, BOX_STATES_CONSUMING_COMPUTE) | ||
| const countDelta = stateTransitionDelta(1, event.oldState, event.newState, BOX_STATES_CONSUMING_COMPUTE) | ||
|
|
||
| await this.updateCurrentQuotaUsage(box.organizationId, 'cpu', cpuDelta) | ||
| await this.updateCurrentQuotaUsage(box.organizationId, 'memory', memoryDelta) | ||
| await this.updateCurrentQuotaUsage(box.organizationId, 'disk', diskDelta) | ||
| await this.updateCurrentQuotaUsage(box.organizationId, 'gpu', gpuDelta) | ||
| await this.updateCurrentQuotaUsage(box.organizationId, 'count', countDelta) | ||
| } catch (error) { | ||
| this.logger.warn(`Error updating cached box quota usage for organization ${box.organizationId}: ${error}`) | ||
| } finally { | ||
| await this.redisLockProvider.unlock(lockKey) | ||
| } | ||
| } |
There was a problem hiding this comment.
handleBoxStateUpdated (organization-usage.service.ts:450-455) returns early for any transition into/out of RESIZING, so resize (box.service.ts validateResizeQuota) never calls updateCurrentQuotaUsage to atomically drain pending into current; current-cache TTL and pending-reservation TTL are set independently, so if current expires and is re-fetched from the DB (already reflecting the resized cpu/mem) before pending's own 60s TTL lapses, getBoxUsageOverview adds the stale pending delta on top of the already-correct current, double counting the resize and spuriously rejecting unrelated creates/starts/resizes on that org for up to ~60s.
Enforce per-organization resource ceilings (cpu, memory, disk, gpu) and a maximum number of concurrently running boxes. Limits live in a dedicated OrganizationQuota table (1:1, get-or-default, backfilled for existing orgs); a ceiling of 0 denies that dimension. Enforcement runs on every path that adds or grows usage — create, start, warm-pool assignment, proxy auto-resume, and resize — via a Redis pending-reservation checked against the summed box-table usage, rejecting with 400 when a projected total would exceed a ceiling. Reservations use atomic Lua and carry a TTL so a crashed request self-heals. Also: - add pure unit tests (org-quota, box-usage) and a gated real-Redis integration spec for the reservation engine. - run the apps/api jest suite in CI (new required test job with a redis service) and on apps-only pushes (pre-push filter); fix two stale boxlite-rest AutoResume specs so the suite is green. Claude-Session: https://claude.ai/code/session_01Y7tNejVgFCQMQhkMvYPKhT
ce1b6ce to
a369386
Compare
| CONSTRAINT "organization_quota_organizationId_fk" FOREIGN KEY ("organizationId") REFERENCES "organization"("id") ON DELETE CASCADE ON UPDATE NO ACTION | ||
| )`, | ||
| ) | ||
| // Backfill one default-quota row per existing organization so every org has an |
There was a problem hiding this comment.
Backfill gives every existing org a flat default quota (10 cpu/40GB/30GB/10 boxes/0 gpu) unrelated to their actual current usage; any org already at/above these numbers, or using GPU boxes, will get BadRequestException on their very next create/start after this deploys, with no grandfathering.
Signed-off-by: dorianzheng <8065637+DorianZheng@users.noreply.github.com>
| import { BoxStoppedEvent } from '../events/box-stopped.event' | ||
| import { OrganizationService } from '../../organization/services/organization.service' | ||
| import { OrganizationUsageService, PendingBoxReservation } from '../../organization/services/organization-usage.service' | ||
| import { resizeQuotaDeltas } from '../../organization/services/box-usage' |
There was a problem hiding this comment.
🛑 Resize never enforces the new org quota
resizeQuotaDeltas is imported but never invoked in box.service.ts, and OrganizationUsageService.validateResizeQuota has no caller outside its own spec (grep -rn validateResizeQuota apps/api/src shows only the definition and the test) — resizing cpu/memory upward on a running box bypasses the org quota entirely.
#1037 removed box resource resize, leaving the quota feature's resize paths dead. Drop validateResizeQuota / resizeQuotaDeltas and the RESIZING usage-accounting (conditional-compute predicate + event-handler skip), and remove the now-unused resizeQuotaDeltas import flagged in review. A box's contribution is again a plain function of its state. Raise the default per-org ceilings from 10 vCPU / 40GB / 30GB / 10 boxes to 64 / 256GB / 512GB / 50 (GPU stays 0) across DEFAULT_ORG_QUOTA, the entity, and the backfill migration, so enabling the quota does not retroactively block existing organizations. Claude-Session: https://claude.ai/code/session_01Y7tNejVgFCQMQhkMvYPKhT
Custom kernels (#1041, #1051) and the capability policy both extend `AdvancedBoxOptions`, the CLI flag set, and option validation, so the two features are combined rather than either replacing the other. Capability name validation moves to `sanitize_common`, which main split out of `sanitize`: a capability list is request data, not a filesystem source, so it must also be checked on the persisted path. The warm-pool capability test now stubs `organizationUsageService`, which the org-quota work (#1028) made a required collaborator of `BoxService::create`.
Per-organization ceilings (cpu / memory / disk / gpu) plus a max concurrent-box count, stored in a dedicated
OrganizationQuotatable (1:1, get-or-default, backfilled for existing orgs; a ceiling of 0 denies that dimension).Enforcement runs on every path that adds or grows usage — create, start, warm-pool assignment, proxy auto-resume, and resize — via a Redis pending-reservation (atomic Lua, TTL self-heal) checked against the summed box-table usage; rejects with 400 when a projected total would exceed a ceiling.
Tests / CI
org-quota,box-usage) plus a gated real-Redis integration spec for the reservation engine.appsjest job intest.yml(with a redis service), andapps/added to the pre-push test filter.Deferred (not in this PR)
listAvailableRegionsdanglingregion_quotaSQL cleanup (needs its own test).Box.createdBycolumn + the org-vs-user precedence rule).https://claude.ai/code/session_01Y7tNejVgFCQMQhkMvYPKhT
Summary by CodeRabbit
New Features
Bug Fixes
Chores / Tests / CI