feat(usage): add box usage metering ledger (Phase 1 Usage)#969
feat(usage): add box usage metering ledger (Phase 1 Usage)#969law-chain-hot wants to merge 2 commits into
Conversation
📝 WalkthroughWalkthroughAdds active and archived usage-period entities, database migration, lifecycle event handling, scheduled rollover and archiving jobs, NestJS module wiring, and comprehensive service tests. ChangesUsage tracking
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
16e5c5e to
2febc70
Compare
📦 BoxLite review — couldn't completepowered by BoxLite |
| case BoxState.STARTED: { | ||
| await this.closeUsagePeriod(event.box.id) | ||
| await this.createUsagePeriod(event) | ||
| break | ||
| } | ||
| case BoxState.STOPPING: | ||
| await this.closeUsagePeriod(event.box.id) | ||
| await this.createUsagePeriod(event, true) | ||
| break | ||
| // Safeguard if STOPPING state is skipped. | ||
| case BoxState.STOPPED: { | ||
| const cpuUsagePeriod = await this.usagePeriodRepository.findOne({ | ||
| where: { | ||
| boxId: event.box.id, | ||
| endAt: IsNull(), | ||
| cpu: Not(0), | ||
| }, | ||
| }) | ||
| if (cpuUsagePeriod) { | ||
| await this.closeUsagePeriod(event.box.id) | ||
| await this.createUsagePeriod(event, true) | ||
| } | ||
| break | ||
| } |
There was a problem hiding this comment.
🛑 close+reopen not transactional in event handlers
unlike closeAndReopenUsagePeriods (line 187) which wraps close+create in one DB transaction, handleBoxStateUpdate calls closeUsagePeriod then createUsagePeriod as two separate saves; if the second save throws (DB blip), the box is left with zero open usage periods and no automatic recovery — cron rollover only rewrites existing open rows (endAt IsNull), so it can't detect a box that lost its open row, silently under-billing until the box's next STATE_UPDATED/DESIRED_STATE_UPDATED event, which may never come while the box just keeps running.
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (4)
apps/api/src/usage/entities/usage-period.entity.ts (1)
48-49: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winType
regionTypeasRegionTypeinstead ofstring.
boxClassis typed as theBoxClassenum, butregionTypeis typed asstringdespite importing and usingRegionType. This inconsistency reduces type safety — callers could assign arbitrary strings. The same issue exists inUsagePeriodArchive.♻️ Proposed fix
- `@Column`({ type: 'character varying', default: RegionType.SHARED }) - regionType: string = RegionType.SHARED + `@Column`({ type: 'character varying', default: RegionType.SHARED }) + regionType: RegionType = RegionType.SHARED🤖 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/usage/entities/usage-period.entity.ts` around lines 48 - 49, Change the regionType property in both UsagePeriod and UsagePeriodArchive from string to RegionType, preserving the existing RegionType.SHARED default and column configuration.apps/api/src/migrations/pre-deploy/1782700000000-migration.ts (1)
29-45: 🚀 Performance & Scalability | 🔵 TrivialConsider adding indexes to the archive table.
The active table has indexes on
("boxId", "endAt")and a partial unique index, butbox_usage_period_archivehas none. SinceorganizationIdis explicitly described as a redundant column to "optimize billing queries," the archive table — which accumulates closed periods and will grow unbounded — would benefit from indexes on commonly queried columns (e.g.,organizationId,boxId,startAt/endAt). Without them, billing queries scanning archived periods will degrade as the table grows.🤖 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/migrations/pre-deploy/1782700000000-migration.ts` around lines 29 - 45, Add indexes to the archive table created in the migration, prioritizing billing query access through organizationId and common box/time lookups using boxId with endAt and/or startAt. Define the indexes alongside the CREATE TABLE migration and ensure the down migration removes them when dropping box_usage_period_archive.apps/api/src/usage/usage.service.spec.ts (1)
216-324: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd test coverage for the STOPPED safeguard and remaining terminal states.
The service's
handleBoxStateUpdateincludes a safeguard for whenSTOPPINGis skipped (checking for open CPU periods onSTOPPED), and handlesERROR,ARCHIVED, andDESTROYINGstates by closing the period. None of these paths are tested — onlyDESTROYEDis exercised indirectly via the archive test. The STOPPED safeguard is particularly important as it's a recovery path for unexpected state transitions.🤖 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/usage/usage.service.spec.ts` around lines 216 - 324, Add tests for handleBoxStateUpdate covering a direct STOPPED transition that closes any open CPU period, plus ERROR, ARCHIVED, and DESTROYING transitions that close the active period. Assert the resulting period has an endAt timestamp and no open CPU allocation, and verify the STOPPED recovery path works when STOPPING was skipped.apps/api/src/usage/services/usage.service.ts (1)
227-248: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winArchiving up to 5000 rows every 5s in one transaction.
A single transaction that reads, deletes, and re-inserts up to 5000 rows on a 5-second cadence can hold row/index locks long enough to contend with lifecycle-event writes and bloat WAL/undo. Consider a smaller batch size with the loop draining across runs, and ensure
endAt(the archive selection predicate) andstartAt(ordering) are indexed on the active table to keep thefindcheap.🤖 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/usage/services/usage.service.ts` around lines 227 - 248, Reduce the batch size used by the usage-period archival transaction and keep each run bounded so subsequent executions drain remaining rows; update the transaction around the archive logic in the usage-period service accordingly. Add or verify indexes on UsagePeriod.endAt for the selection predicate and UsagePeriod.startAt for ordering to keep archival queries efficient.
🤖 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 `@apps/api/src/usage/entities/usage-period-archive.entity.ts`:
- Line 56: Replace the unsafe cast in the usage-period archive mapping with an
explicit null check for usagePeriod.endAt, throwing a clear error when it is
null before assigning usagePeriodEntity.endAt. Keep the archive entity’s
non-null Date contract intact and reference the surrounding archive creation
logic for the guard placement.
In `@apps/api/src/usage/services/usage.service.ts`:
- Around line 217-252: archiveUsagePeriods can leave the Redis lock held when
its transaction fails. Wrap the transaction and related work in a try/finally
block within archiveUsagePeriods, ensuring redisLockProvider.unlock(lockKey)
always executes after a successful lock acquisition.
- Around line 43-49: Bound the wait in UsageService.onApplicationShutdown so
shutdown cannot hang indefinitely on stuck activeJobs. Add a deadline or maximum
wait duration, continue polling while jobs remain only until that limit, then
log that shutdown is proceeding with unfinished jobs and return.
- Around line 153-215: Ensure the global Redis lock acquired in
closeAndReopenUsagePeriods is always released by wrapping all work after lock
acquisition, including usage-period lookup and iteration, in a try/finally
block. Move unlock('close-and-reopen-usage-periods') into that finally block,
while preserving the per-usage-period cleanup and lock release behavior.
- Around line 175-206: After acquiring the lock in the usage-period processing
loop, refetch the current open period for the box instead of using the pre-lock
usagePeriod snapshot. Update and save it only when it still exists with a null
endAt; otherwise skip rollover, and use the freshly loaded period when creating
the next period inside the transaction.
In `@apps/api/src/usage/usage.service.spec.ts`:
- Around line 133-149: Update FakeTransactionManager.save to distinguish
UsagePeriodArchive from UsagePeriod for both single entities and arrays; route
single archives and archive arrays through archivedRows, while retaining
active-table handling for UsagePeriod values. Use an explicit entity-type check
rather than the current !Array.isArray(value) condition, and preserve archive ID
assignment and transaction tracking.
---
Nitpick comments:
In `@apps/api/src/migrations/pre-deploy/1782700000000-migration.ts`:
- Around line 29-45: Add indexes to the archive table created in the migration,
prioritizing billing query access through organizationId and common box/time
lookups using boxId with endAt and/or startAt. Define the indexes alongside the
CREATE TABLE migration and ensure the down migration removes them when dropping
box_usage_period_archive.
In `@apps/api/src/usage/entities/usage-period.entity.ts`:
- Around line 48-49: Change the regionType property in both UsagePeriod and
UsagePeriodArchive from string to RegionType, preserving the existing
RegionType.SHARED default and column configuration.
In `@apps/api/src/usage/services/usage.service.ts`:
- Around line 227-248: Reduce the batch size used by the usage-period archival
transaction and keep each run bounded so subsequent executions drain remaining
rows; update the transaction around the archive logic in the usage-period
service accordingly. Add or verify indexes on UsagePeriod.endAt for the
selection predicate and UsagePeriod.startAt for ordering to keep archival
queries efficient.
In `@apps/api/src/usage/usage.service.spec.ts`:
- Around line 216-324: Add tests for handleBoxStateUpdate covering a direct
STOPPED transition that closes any open CPU period, plus ERROR, ARCHIVED, and
DESTROYING transitions that close the active period. Assert the resulting period
has an endAt timestamp and no open CPU allocation, and verify the STOPPED
recovery path works when STOPPING was skipped.
🪄 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: 3de64576-9f77-4d10-80cf-b05cdec48cf6
📒 Files selected for processing (7)
apps/api/src/app.module.tsapps/api/src/migrations/pre-deploy/1782700000000-migration.tsapps/api/src/usage/entities/usage-period-archive.entity.tsapps/api/src/usage/entities/usage-period.entity.tsapps/api/src/usage/services/usage.service.tsapps/api/src/usage/usage.module.tsapps/api/src/usage/usage.service.spec.ts
| usagePeriodEntity.boxId = usagePeriod.boxId | ||
| usagePeriodEntity.organizationId = usagePeriod.organizationId | ||
| usagePeriodEntity.startAt = usagePeriod.startAt | ||
| usagePeriodEntity.endAt = usagePeriod.endAt as Date |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Unsafe as Date cast on endAt masks null values.
usagePeriod.endAt is Date | null, but the archive column is NOT NULL. The as Date cast silences the compiler; if an open period is accidentally passed, it will fail at the DB level with a constraint violation rather than being caught at compile time. Consider a runtime guard or explicit error.
🛡️ Proposed fix
public static fromUsagePeriod(usagePeriod: UsagePeriod) {
+ if (!usagePeriod.endAt) {
+ throw new Error('Cannot archive a usage period with null endAt')
+ }
const usagePeriodEntity = new UsagePeriodArchive()
usagePeriodEntity.boxId = usagePeriod.boxId
usagePeriodEntity.organizationId = usagePeriod.organizationId
usagePeriodEntity.startAt = usagePeriod.startAt
- usagePeriodEntity.endAt = usagePeriod.endAt as Date
+ usagePeriodEntity.endAt = usagePeriod.endAt📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| usagePeriodEntity.endAt = usagePeriod.endAt as Date | |
| if (!usagePeriod.endAt) { | |
| throw new Error('Cannot archive a usage period with null endAt') | |
| } | |
| usagePeriodEntity.endAt = usagePeriod.endAt |
🤖 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/usage/entities/usage-period-archive.entity.ts` at line 56,
Replace the unsafe cast in the usage-period archive mapping with an explicit
null check for usagePeriod.endAt, throwing a clear error when it is null before
assigning usagePeriodEntity.endAt. Keep the archive entity’s non-null Date
contract intact and reference the surrounding archive creation logic for the
guard placement.
| async onApplicationShutdown() { | ||
| // Wait for all active jobs to finish. | ||
| while (this.activeJobs.size > 0) { | ||
| this.logger.log(`Waiting for ${this.activeJobs.size} active jobs to finish`) | ||
| await sleep(1000) | ||
| } | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Add a bounded deadline to the shutdown wait.
onApplicationShutdown loops until activeJobs drains with no upper bound. A single stuck job (hung DB/Redis call) will block graceful shutdown indefinitely until the orchestrator force-kills the pod, defeating in-flight completion and delaying rollout.
🛡️ Proposed bounded wait
async onApplicationShutdown() {
- // Wait for all active jobs to finish.
- while (this.activeJobs.size > 0) {
- this.logger.log(`Waiting for ${this.activeJobs.size} active jobs to finish`)
- await sleep(1000)
- }
+ // Wait for all active jobs to finish, up to a bounded deadline.
+ const deadline = Date.now() + 30_000
+ while (this.activeJobs.size > 0 && Date.now() < deadline) {
+ this.logger.log(`Waiting for ${this.activeJobs.size} active jobs to finish`)
+ await sleep(1000)
+ }
+ if (this.activeJobs.size > 0) {
+ this.logger.warn(`Shutdown proceeding with ${this.activeJobs.size} unfinished jobs`)
+ }
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| async onApplicationShutdown() { | |
| // Wait for all active jobs to finish. | |
| while (this.activeJobs.size > 0) { | |
| this.logger.log(`Waiting for ${this.activeJobs.size} active jobs to finish`) | |
| await sleep(1000) | |
| } | |
| } | |
| async onApplicationShutdown() { | |
| // Wait for all active jobs to finish, up to a bounded deadline. | |
| const deadline = Date.now() + 30_000 | |
| while (this.activeJobs.size > 0 && Date.now() < deadline) { | |
| this.logger.log(`Waiting for ${this.activeJobs.size} active jobs to finish`) | |
| await sleep(1000) | |
| } | |
| if (this.activeJobs.size > 0) { | |
| this.logger.warn(`Shutdown proceeding with ${this.activeJobs.size} unfinished jobs`) | |
| } | |
| } |
🤖 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/usage/services/usage.service.ts` around lines 43 - 49, Bound the
wait in UsageService.onApplicationShutdown so shutdown cannot hang indefinitely
on stuck activeJobs. Add a deadline or maximum wait duration, continue polling
while jobs remain only until that limit, then log that shutdown is proceeding
with unfinished jobs and return.
| @Cron(CronExpression.EVERY_MINUTE, { name: 'close-and-reopen-usage-periods' }) | ||
| @TrackJobExecution() | ||
| @LogExecution('close-and-reopen-usage-periods') | ||
| @WithInstrumentation() | ||
| async closeAndReopenUsagePeriods() { | ||
| if (!(await this.redisLockProvider.lock('close-and-reopen-usage-periods', 60))) { | ||
| return | ||
| } | ||
|
|
||
| const usagePeriods = await this.usagePeriodRepository.find({ | ||
| where: { | ||
| endAt: IsNull(), | ||
| // 1 day ago | ||
| startAt: LessThan(new Date(Date.now() - 1000 * 60 * 60 * 24)), | ||
| organizationId: Not(BOX_WARM_POOL_UNASSIGNED_ORGANIZATION), | ||
| }, | ||
| order: { | ||
| startAt: 'ASC', | ||
| }, | ||
| take: 100, | ||
| }) | ||
|
|
||
| for (const usagePeriod of usagePeriods) { | ||
| if (!(await this.aquireLock(usagePeriod.boxId))) { | ||
| continue | ||
| } | ||
|
|
||
| try { | ||
| const box = await this.boxRepository.findOne({ | ||
| where: { | ||
| id: usagePeriod.boxId, | ||
| }, | ||
| }) | ||
|
|
||
| await this.usagePeriodRepository.manager.transaction(async (transactionalEntityManager) => { | ||
| const closeTime = new Date() | ||
| usagePeriod.endAt = closeTime | ||
| await transactionalEntityManager.save(usagePeriod) | ||
|
|
||
| if ( | ||
| box && | ||
| (box.state === BoxState.STARTED || box.state === BoxState.STOPPED || box.state === BoxState.STOPPING) | ||
| ) { | ||
| const newUsagePeriod = UsagePeriod.fromUsagePeriod(usagePeriod) | ||
| newUsagePeriod.startAt = closeTime | ||
| newUsagePeriod.endAt = null | ||
| if (box.state === BoxState.STOPPED) { | ||
| newUsagePeriod.cpu = 0 | ||
| newUsagePeriod.gpu = 0 | ||
| newUsagePeriod.mem = 0 | ||
| } | ||
| await transactionalEntityManager.save(newUsagePeriod) | ||
| } | ||
| }) | ||
| } catch (error) { | ||
| this.logger.error(`Error closing and reopening usage period ${usagePeriod.boxId}`, error) | ||
| } finally { | ||
| await this.releaseLock(usagePeriod.boxId) | ||
| } | ||
| } | ||
|
|
||
| await this.redisLockProvider.unlock('close-and-reopen-usage-periods') | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Global Redis lock is not released on error paths.
The close-and-reopen-usage-periods lock is only unlocked at Line 214. If find (Line 162) throws, or releaseLock in the per-iteration finally (Line 210) rejects and propagates, Line 214 is skipped and the lock is held until its 60s TTL expires, stalling this job across all replicas for up to a minute. Wrap the body in try/finally.
🔒 Guard the global unlock
if (!(await this.redisLockProvider.lock('close-and-reopen-usage-periods', 60))) {
return
}
-
- const usagePeriods = await this.usagePeriodRepository.find({
- ...
- })
- ...
- await this.redisLockProvider.unlock('close-and-reopen-usage-periods')
+ try {
+ const usagePeriods = await this.usagePeriodRepository.find({
+ ...
+ })
+ // ... loop ...
+ } finally {
+ await this.redisLockProvider.unlock('close-and-reopen-usage-periods')
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| @Cron(CronExpression.EVERY_MINUTE, { name: 'close-and-reopen-usage-periods' }) | |
| @TrackJobExecution() | |
| @LogExecution('close-and-reopen-usage-periods') | |
| @WithInstrumentation() | |
| async closeAndReopenUsagePeriods() { | |
| if (!(await this.redisLockProvider.lock('close-and-reopen-usage-periods', 60))) { | |
| return | |
| } | |
| const usagePeriods = await this.usagePeriodRepository.find({ | |
| where: { | |
| endAt: IsNull(), | |
| // 1 day ago | |
| startAt: LessThan(new Date(Date.now() - 1000 * 60 * 60 * 24)), | |
| organizationId: Not(BOX_WARM_POOL_UNASSIGNED_ORGANIZATION), | |
| }, | |
| order: { | |
| startAt: 'ASC', | |
| }, | |
| take: 100, | |
| }) | |
| for (const usagePeriod of usagePeriods) { | |
| if (!(await this.aquireLock(usagePeriod.boxId))) { | |
| continue | |
| } | |
| try { | |
| const box = await this.boxRepository.findOne({ | |
| where: { | |
| id: usagePeriod.boxId, | |
| }, | |
| }) | |
| await this.usagePeriodRepository.manager.transaction(async (transactionalEntityManager) => { | |
| const closeTime = new Date() | |
| usagePeriod.endAt = closeTime | |
| await transactionalEntityManager.save(usagePeriod) | |
| if ( | |
| box && | |
| (box.state === BoxState.STARTED || box.state === BoxState.STOPPED || box.state === BoxState.STOPPING) | |
| ) { | |
| const newUsagePeriod = UsagePeriod.fromUsagePeriod(usagePeriod) | |
| newUsagePeriod.startAt = closeTime | |
| newUsagePeriod.endAt = null | |
| if (box.state === BoxState.STOPPED) { | |
| newUsagePeriod.cpu = 0 | |
| newUsagePeriod.gpu = 0 | |
| newUsagePeriod.mem = 0 | |
| } | |
| await transactionalEntityManager.save(newUsagePeriod) | |
| } | |
| }) | |
| } catch (error) { | |
| this.logger.error(`Error closing and reopening usage period ${usagePeriod.boxId}`, error) | |
| } finally { | |
| await this.releaseLock(usagePeriod.boxId) | |
| } | |
| } | |
| await this.redisLockProvider.unlock('close-and-reopen-usage-periods') | |
| } | |
| `@Cron`(CronExpression.EVERY_MINUTE, { name: 'close-and-reopen-usage-periods' }) | |
| `@TrackJobExecution`() | |
| `@LogExecution`('close-and-reopen-usage-periods') | |
| `@WithInstrumentation`() | |
| async closeAndReopenUsagePeriods() { | |
| if (!(await this.redisLockProvider.lock('close-and-reopen-usage-periods', 60))) { | |
| return | |
| } | |
| try { | |
| const usagePeriods = await this.usagePeriodRepository.find({ | |
| where: { | |
| endAt: IsNull(), | |
| // 1 day ago | |
| startAt: LessThan(new Date(Date.now() - 1000 * 60 * 60 * 24)), | |
| organizationId: Not(BOX_WARM_POOL_UNASSIGNED_ORGANIZATION), | |
| }, | |
| order: { | |
| startAt: 'ASC', | |
| }, | |
| take: 100, | |
| }) | |
| for (const usagePeriod of usagePeriods) { | |
| if (!(await this.aquireLock(usagePeriod.boxId))) { | |
| continue | |
| } | |
| try { | |
| const box = await this.boxRepository.findOne({ | |
| where: { | |
| id: usagePeriod.boxId, | |
| }, | |
| }) | |
| await this.usagePeriodRepository.manager.transaction(async (transactionalEntityManager) => { | |
| const closeTime = new Date() | |
| usagePeriod.endAt = closeTime | |
| await transactionalEntityManager.save(usagePeriod) | |
| if ( | |
| box && | |
| (box.state === BoxState.STARTED || box.state === BoxState.STOPPED || box.state === BoxState.STOPPING) | |
| ) { | |
| const newUsagePeriod = UsagePeriod.fromUsagePeriod(usagePeriod) | |
| newUsagePeriod.startAt = closeTime | |
| newUsagePeriod.endAt = null | |
| if (box.state === BoxState.STOPPED) { | |
| newUsagePeriod.cpu = 0 | |
| newUsagePeriod.gpu = 0 | |
| newUsagePeriod.mem = 0 | |
| } | |
| await transactionalEntityManager.save(newUsagePeriod) | |
| } | |
| }) | |
| } catch (error) { | |
| this.logger.error(`Error closing and reopening usage period ${usagePeriod.boxId}`, error) | |
| } finally { | |
| await this.releaseLock(usagePeriod.boxId) | |
| } | |
| } | |
| } finally { | |
| await this.redisLockProvider.unlock('close-and-reopen-usage-periods') | |
| } | |
| } |
🤖 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/usage/services/usage.service.ts` around lines 153 - 215, Ensure
the global Redis lock acquired in closeAndReopenUsagePeriods is always released
by wrapping all work after lock acquisition, including usage-period lookup and
iteration, in a try/finally block. Move unlock('close-and-reopen-usage-periods')
into that finally block, while preserving the per-usage-period cleanup and lock
release behavior.
| for (const usagePeriod of usagePeriods) { | ||
| if (!(await this.aquireLock(usagePeriod.boxId))) { | ||
| continue | ||
| } | ||
|
|
||
| try { | ||
| const box = await this.boxRepository.findOne({ | ||
| where: { | ||
| id: usagePeriod.boxId, | ||
| }, | ||
| }) | ||
|
|
||
| await this.usagePeriodRepository.manager.transaction(async (transactionalEntityManager) => { | ||
| const closeTime = new Date() | ||
| usagePeriod.endAt = closeTime | ||
| await transactionalEntityManager.save(usagePeriod) | ||
|
|
||
| if ( | ||
| box && | ||
| (box.state === BoxState.STARTED || box.state === BoxState.STOPPED || box.state === BoxState.STOPPING) | ||
| ) { | ||
| const newUsagePeriod = UsagePeriod.fromUsagePeriod(usagePeriod) | ||
| newUsagePeriod.startAt = closeTime | ||
| newUsagePeriod.endAt = null | ||
| if (box.state === BoxState.STOPPED) { | ||
| newUsagePeriod.cpu = 0 | ||
| newUsagePeriod.gpu = 0 | ||
| newUsagePeriod.mem = 0 | ||
| } | ||
| await transactionalEntityManager.save(newUsagePeriod) | ||
| } | ||
| }) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm both cron and event handlers serialize on the same box lock key
rg -nP "usage-period-\$\{boxId\}|usage-period-" apps/api/src/usage/services/usage.service.ts
# Check whether any DB constraint prevents two open (endAt IS NULL) periods per box
fd -e ts . apps/api/src/migrations apps/api/src/usage/entities --exec rg -nP "endAt|IsNull|UNIQUE|Index" {}Repository: boxlite-ai/boxlite
Length of output: 6478
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant service section and surrounding helpers
sed -n '130,290p' apps/api/src/usage/services/usage.service.ts
# Find other code paths that manipulate usage periods or use the same lock
rg -n "UsagePeriod\.fromUsagePeriod|usage-period-\$\{boxId\}|aquireLock|releaseLock|endAt = null|box_usage_period_one_open_per_box_idx" apps/api/src -g'*.ts'Repository: boxlite-ai/boxlite
Length of output: 6259
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the event handlers and the initial query that populates usagePeriods
sed -n '1,140p' apps/api/src/usage/services/usage.service.ts
# Inspect the usage-period entity for relevant constraints/columns
sed -n '1,120p' apps/api/src/usage/entities/usage-period.entity.tsRepository: boxlite-ai/boxlite
Length of output: 7263
Refetch the open usage period inside the box lock
The partial unique index blocks duplicate open rows, but this cron still works from a pre-lock snapshot. If another handler closes the period first, this save can overwrite endAt on an already-closed row. Reload the current open period after acquiring the lock and only roll it forward if it is still open.
🤖 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/usage/services/usage.service.ts` around lines 175 - 206, After
acquiring the lock in the usage-period processing loop, refetch the current open
period for the box instead of using the pre-lock usagePeriod snapshot. Update
and save it only when it still exists with a null endAt; otherwise skip
rollover, and use the freshly loaded period when creating the next period inside
the transaction.
| @Cron(CronExpression.EVERY_5_SECONDS, { name: 'archive-usage-periods' }) | ||
| @TrackJobExecution() | ||
| @LogExecution('archive-usage-periods') | ||
| @WithInstrumentation() | ||
| async archiveUsagePeriods() { | ||
| const lockKey = 'archive-usage-periods' | ||
| if (!(await this.redisLockProvider.lock(lockKey, 60))) { | ||
| return | ||
| } | ||
|
|
||
| await this.usagePeriodRepository.manager.transaction(async (transactionalEntityManager) => { | ||
| const usagePeriods = await transactionalEntityManager.find(UsagePeriod, { | ||
| where: { | ||
| endAt: Not(IsNull()), | ||
| }, | ||
| order: { | ||
| startAt: 'ASC', | ||
| }, | ||
| take: 5000, | ||
| }) | ||
|
|
||
| if (usagePeriods.length === 0) { | ||
| return | ||
| } | ||
|
|
||
| this.logger.debug(`Found ${usagePeriods.length} usage periods to archive`) | ||
|
|
||
| await transactionalEntityManager.delete( | ||
| UsagePeriod, | ||
| usagePeriods.map((usagePeriod) => usagePeriod.id), | ||
| ) | ||
| await transactionalEntityManager.save(usagePeriods.map(UsagePeriodArchive.fromUsagePeriod)) | ||
| }) | ||
|
|
||
| await this.redisLockProvider.unlock(lockKey) | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Global lock leaks if the archive transaction throws.
Same pattern as closeAndReopenUsagePeriods: if the transaction (Lines 227-249) rejects, unlock(lockKey) at Line 251 never runs and the lock is held for the full 60s TTL. Since this job runs every 5s, one failure suppresses ~12 subsequent runs across replicas. Wrap in try/finally.
🔒 Guard the global unlock
const lockKey = 'archive-usage-periods'
if (!(await this.redisLockProvider.lock(lockKey, 60))) {
return
}
-
- await this.usagePeriodRepository.manager.transaction(async (transactionalEntityManager) => {
- ...
- })
-
- await this.redisLockProvider.unlock(lockKey)
+ try {
+ await this.usagePeriodRepository.manager.transaction(async (transactionalEntityManager) => {
+ ...
+ })
+ } finally {
+ await this.redisLockProvider.unlock(lockKey)
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| @Cron(CronExpression.EVERY_5_SECONDS, { name: 'archive-usage-periods' }) | |
| @TrackJobExecution() | |
| @LogExecution('archive-usage-periods') | |
| @WithInstrumentation() | |
| async archiveUsagePeriods() { | |
| const lockKey = 'archive-usage-periods' | |
| if (!(await this.redisLockProvider.lock(lockKey, 60))) { | |
| return | |
| } | |
| await this.usagePeriodRepository.manager.transaction(async (transactionalEntityManager) => { | |
| const usagePeriods = await transactionalEntityManager.find(UsagePeriod, { | |
| where: { | |
| endAt: Not(IsNull()), | |
| }, | |
| order: { | |
| startAt: 'ASC', | |
| }, | |
| take: 5000, | |
| }) | |
| if (usagePeriods.length === 0) { | |
| return | |
| } | |
| this.logger.debug(`Found ${usagePeriods.length} usage periods to archive`) | |
| await transactionalEntityManager.delete( | |
| UsagePeriod, | |
| usagePeriods.map((usagePeriod) => usagePeriod.id), | |
| ) | |
| await transactionalEntityManager.save(usagePeriods.map(UsagePeriodArchive.fromUsagePeriod)) | |
| }) | |
| await this.redisLockProvider.unlock(lockKey) | |
| } | |
| `@Cron`(CronExpression.EVERY_5_SECONDS, { name: 'archive-usage-periods' }) | |
| `@TrackJobExecution`() | |
| `@LogExecution`('archive-usage-periods') | |
| `@WithInstrumentation`() | |
| async archiveUsagePeriods() { | |
| const lockKey = 'archive-usage-periods' | |
| if (!(await this.redisLockProvider.lock(lockKey, 60))) { | |
| return | |
| } | |
| try { | |
| await this.usagePeriodRepository.manager.transaction(async (transactionalEntityManager) => { | |
| const usagePeriods = await transactionalEntityManager.find(UsagePeriod, { | |
| where: { | |
| endAt: Not(IsNull()), | |
| }, | |
| order: { | |
| startAt: 'ASC', | |
| }, | |
| take: 5000, | |
| }) | |
| if (usagePeriods.length === 0) { | |
| return | |
| } | |
| this.logger.debug(`Found ${usagePeriods.length} usage periods to archive`) | |
| await transactionalEntityManager.delete( | |
| UsagePeriod, | |
| usagePeriods.map((usagePeriod) => usagePeriod.id), | |
| ) | |
| await transactionalEntityManager.save(usagePeriods.map(UsagePeriodArchive.fromUsagePeriod)) | |
| }) | |
| } finally { | |
| await this.redisLockProvider.unlock(lockKey) | |
| } | |
| } |
🤖 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/usage/services/usage.service.ts` around lines 217 - 252,
archiveUsagePeriods can leave the Redis lock held when its transaction fails.
Wrap the transaction and related work in a try/finally block within
archiveUsagePeriods, ensuring redisLockProvider.unlock(lockKey) always executes
after a successful lock acquisition.
| async save( | ||
| entityOrInput: typeof UsagePeriod | typeof UsagePeriodArchive | UsagePeriod | UsagePeriod[] | UsagePeriodArchive[], | ||
| input?: UsagePeriod | UsagePeriod[] | UsagePeriodArchive[], | ||
| ): Promise<UsagePeriod | UsagePeriod[] | UsagePeriodArchive[]> { | ||
| this.periods.transactionOperations.push('save') | ||
| const value = input ?? (entityOrInput as UsagePeriod | UsagePeriod[] | UsagePeriodArchive[]) | ||
| if (!Array.isArray(value) || value[0] instanceof UsagePeriod) { | ||
| return this.periods.save(value as UsagePeriod | UsagePeriod[]) | ||
| } | ||
|
|
||
| const rows = value as UsagePeriodArchive[] | ||
| rows.forEach((row) => { | ||
| row.id = row.id ?? `archive-${this.periods.archivedRows.length + 1}` | ||
| this.periods.archivedRows.push(row) | ||
| }) | ||
| return rows | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
FakeTransactionManager.save misroutes single UsagePeriodArchive instances to the active table.
The condition !Array.isArray(value) || value[0] instanceof UsagePeriod routes any non-array value to this.periods.save() (the active rows array). A single UsagePeriodArchive passed without array-wrapping would be silently saved to the active table instead of archivedRows. This is currently masked because the service passes arrays, but it's a latent trap if the service ever saves a single archive row.
🛡️ Proposed fix: route by entity type instead of array membership
async save(
entityOrInput: typeof UsagePeriod | typeof UsagePeriodArchive | UsagePeriod | UsagePeriod[] | UsagePeriodArchive[],
input?: UsagePeriod | UsagePeriod[] | UsagePeriodArchive[],
): Promise<UsagePeriod | UsagePeriod[] | UsagePeriodArchive[]> {
this.periods.transactionOperations.push('save')
const value = input ?? (entityOrInput as UsagePeriod | UsagePeriod[] | UsagePeriodArchive[])
- if (!Array.isArray(value) || value[0] instanceof UsagePeriod) {
+ const first = Array.isArray(value) ? value[0] : value
+ if (first instanceof UsagePeriodArchive) {
+ const rows = Array.isArray(value) ? (value as UsagePeriodArchive[]) : [value as UsagePeriodArchive]
+ rows.forEach((row) => {
+ row.id = row.id ?? `archive-${this.periods.archivedRows.length + 1}`
+ this.periods.archivedRows.push(row)
+ })
+ return rows
+ }
- return this.periods.save(value as UsagePeriod | UsagePeriod[])
- }
-
- const rows = value as UsagePeriodArchive[]
- rows.forEach((row) => {
- row.id = row.id ?? `archive-${this.periods.archivedRows.length + 1}`
- this.periods.archivedRows.push(row)
- })
- return rows
+ return this.periods.save(value as UsagePeriod | UsagePeriod[])
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| async save( | |
| entityOrInput: typeof UsagePeriod | typeof UsagePeriodArchive | UsagePeriod | UsagePeriod[] | UsagePeriodArchive[], | |
| input?: UsagePeriod | UsagePeriod[] | UsagePeriodArchive[], | |
| ): Promise<UsagePeriod | UsagePeriod[] | UsagePeriodArchive[]> { | |
| this.periods.transactionOperations.push('save') | |
| const value = input ?? (entityOrInput as UsagePeriod | UsagePeriod[] | UsagePeriodArchive[]) | |
| if (!Array.isArray(value) || value[0] instanceof UsagePeriod) { | |
| return this.periods.save(value as UsagePeriod | UsagePeriod[]) | |
| } | |
| const rows = value as UsagePeriodArchive[] | |
| rows.forEach((row) => { | |
| row.id = row.id ?? `archive-${this.periods.archivedRows.length + 1}` | |
| this.periods.archivedRows.push(row) | |
| }) | |
| return rows | |
| } | |
| async save( | |
| entityOrInput: typeof UsagePeriod | typeof UsagePeriodArchive | UsagePeriod | UsagePeriod[] | UsagePeriodArchive[], | |
| input?: UsagePeriod | UsagePeriod[] | UsagePeriodArchive[], | |
| ): Promise<UsagePeriod | UsagePeriod[] | UsagePeriodArchive[]> { | |
| this.periods.transactionOperations.push('save') | |
| const value = input ?? (entityOrInput as UsagePeriod | UsagePeriod[] | UsagePeriodArchive[]) | |
| const first = Array.isArray(value) ? value[0] : value | |
| if (first instanceof UsagePeriodArchive) { | |
| const rows = Array.isArray(value) ? (value as UsagePeriodArchive[]) : [value as UsagePeriodArchive] | |
| rows.forEach((row) => { | |
| row.id = row.id ?? `archive-${this.periods.archivedRows.length + 1}` | |
| this.periods.archivedRows.push(row) | |
| }) | |
| return rows | |
| } | |
| return this.periods.save(value as UsagePeriod | UsagePeriod[]) | |
| } |
🤖 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/usage/usage.service.spec.ts` around lines 133 - 149, Update
FakeTransactionManager.save to distinguish UsagePeriodArchive from UsagePeriod
for both single entities and arrays; route single archives and archive arrays
through archivedRows, while retaining active-table handling for UsagePeriod
values. Use an explicit entity-type check rather than the current
!Array.isArray(value) condition, and preserve archive ID assignment and
transaction tracking.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2febc7052c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| const cpuUsagePeriod = await this.usagePeriodRepository.findOne({ | ||
| where: { | ||
| boxId: event.box.id, | ||
| endAt: IsNull(), | ||
| cpu: Not(0), | ||
| }, |
There was a problem hiding this comment.
Refresh disk-only usage on STOPPED resize events
When a stopped box is cold-resized, BoxService.resize updates disk and emits a STATE_UPDATED event back to STOPPED (see apps/api/src/box/services/box.service.ts around the V0 disk: newDisk, state: previousState update). If the box already has an open disk-only usage period, this branch ignores the event because it only matches open periods with cpu != 0, so the ledger keeps charging the old disk size until some later start/stop/destroy/rollover event. The STOPPED handler needs to also close/reopen when the open disk-only period's metered allocation no longer matches the updated box.
Useful? React with 👍 / 👎.
| await this.closeUsagePeriod(event.box.id) | ||
| await this.createUsagePeriod(event) | ||
| break | ||
| } | ||
| case BoxState.STOPPING: | ||
| await this.closeUsagePeriod(event.box.id) | ||
| await this.createUsagePeriod(event, true) |
There was a problem hiding this comment.
Make lifecycle period rotation atomic
For STARTED/STOPPING transitions that already have an open period, the old period is committed closed before the replacement period is inserted. If createUsagePeriod then fails or the process exits between these two awaits, the box is left with no open usage row and future usage is silently unmetered until another lifecycle event happens. The close and create operations should be performed in one DB transaction, as the daily rollover path already does.
Useful? React with 👍 / 👎.
| import { RegionType } from '../../region/enums/region-type.enum' | ||
|
|
||
| @Injectable() | ||
| export class UsageService implements TrackableJobExecutions, OnApplicationShutdown { |
There was a problem hiding this comment.
Seed usage periods for boxes already running at deploy
After this module is enabled on an existing environment, boxes that are already STARTED or STOPPED will not emit a new state event, and the daily rollover job only processes rows that already exist in box_usage_period. Those boxes therefore have no open ledger row from deploy time until their next lifecycle transition, losing running or disk-only usage for that interval. Add a bootstrap/backfill path that creates open periods for current non-terminal boxes when the service starts.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
🧹 Nitpick comments (1)
apps/api/src/usage/entities/box-usage-period.entity.ts (1)
48-49: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winType
regionTypeasRegionTypeinstead ofstring.
boxClassis typed as theBoxClassenum whileregionTypeis typed asstring, despite both usingcharacter varyingcolumns and enum defaults. TypingregionTypeasRegionTypewould provide compile-time safety consistent withboxClass.♻️ Proposed type fix
- `@Column`({ type: 'character varying', default: RegionType.SHARED }) - regionType: string = RegionType.SHARED + `@Column`({ type: 'character varying', default: RegionType.SHARED }) + regionType: RegionType = RegionType.SHARED🤖 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/usage/entities/box-usage-period.entity.ts` around lines 48 - 49, Change the regionType property in the box usage period entity from string to RegionType, preserving its existing default and database column configuration so it has compile-time enum safety consistent with boxClass.
🤖 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.
Nitpick comments:
In `@apps/api/src/usage/entities/box-usage-period.entity.ts`:
- Around line 48-49: Change the regionType property in the box usage period
entity from string to RegionType, preserving its existing default and database
column configuration so it has compile-time enum safety consistent with
boxClass.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 361fda75-c1c3-49f9-9263-0576747590fc
📒 Files selected for processing (5)
apps/api/src/usage/entities/box-usage-period-archive.entity.tsapps/api/src/usage/entities/box-usage-period.entity.tsapps/api/src/usage/services/usage.service.tsapps/api/src/usage/usage.module.tsapps/api/src/usage/usage.service.spec.ts
✅ Files skipped from review due to trivial changes (1)
- apps/api/src/usage/entities/box-usage-period-archive.entity.ts
🚧 Files skipped from review as they are similar to previous changes (3)
- apps/api/src/usage/usage.module.ts
- apps/api/src/usage/services/usage.service.ts
- apps/api/src/usage/usage.service.spec.ts
Summary
box_usage_periodandbox_usage_period_archiveas the usage ledgerArchitecture
Implementation notes
The lifecycle writer is intentionally small: Box state events open or close periods, Redis locks serialize writers, rollover bounds open-period duration, and archiving keeps the active table lightweight.
Review order
apps/api/src/usage/services/usage.service.ts: lifecycle, locks, rollover, archiveapps/api/src/usage/entities/: ledger schemaapps/api/src/migrations/pre-deploy/1782700000000-migration.ts: physical tables and indexesapps/api/src/usage/usage.service.spec.ts: lifecycle and persistence behaviorValidation
Supersedes #947.
Summary by CodeRabbit
New Features
Tests