diff --git a/src/backend/services/metering/MeteringService.ts b/src/backend/services/metering/MeteringService.ts index 7813d7b62..c946a3789 100644 --- a/src/backend/services/metering/MeteringService.ts +++ b/src/backend/services/metering/MeteringService.ts @@ -146,6 +146,9 @@ export class MeteringService extends PuterService { static CREDIT_CACHE_MS = 15_000; static CREDIT_CACHE_LIMIT = 50_000; + /** Where "about to run out" starts, as a fraction of the month allowance. */ + static NEAR_LIMIT_FRACTION = 0.9; + /** * How long usage that isn't decided on may sit in memory before it is * written, and how many actor buckets are held at once. Egress and @@ -172,6 +175,12 @@ export class MeteringService extends PuterService { { policy: SubscriptionPolicy; expiresAt: number } >(); + /** Uuid → the last budget state announced, so a retry loop emits once. */ + private creditAlertState = new Map< + string, + 'ok' | 'near-limit' | 'exhausted' + >(); + /** Uuid → whether the actor had budget left. See CREDIT_CACHE_MS. */ private creditCache = new Map< string, @@ -1276,6 +1285,8 @@ export class MeteringService extends PuterService { /** Local-only drop. The announcement path is `invalidateActorCredits`. */ #dropCachedCredits(userUuid: string): void { this.creditCache.delete(userUuid); + // Added capacity re-arms the alert: the next exhaustion is news again. + this.creditAlertState.delete(userUuid); } async #refreshCredits(actor: Actor): Promise { @@ -1333,14 +1344,66 @@ export class MeteringService extends PuterService { this.rememberHasCredits(userId, true); return; } - this.rememberHasCredits( - userId, - MeteringService.remainingFrom( - allowanceUsed, - monthUsageAllowance, - addons, - ) > 0, + const remaining = MeteringService.remainingFrom( + allowanceUsed, + monthUsageAllowance, + addons, ); + this.rememberHasCredits(userId, remaining > 0); + this.#noteCreditState( + userId, + remaining, + allowanceUsed, + monthUsageAllowance, + addons, + ); + } + + /** Transitions only: a blocked actor retries, and every retry lands here. */ + #noteCreditState( + userUuid: string, + remaining: number, + allowanceUsed: number, + monthUsageAllowance: number, + addons: UsageAddons | null | undefined, + ): void { + // Purchased credits are spendable, so the allowance alone warns early. + const capacity = + (monthUsageAllowance || 0) + (addons?.purchasedCredits || 0); + const state = + remaining <= 0 + ? 'exhausted' + : remaining <= + capacity * (1 - MeteringService.NEAR_LIMIT_FRACTION) + ? 'near-limit' + : 'ok'; + + if (this.creditAlertState.get(userUuid) === state) return; + // Same FIFO bound as `creditCache`; this map has one entry per actor. + if ( + this.creditAlertState.size >= MeteringService.CREDIT_CACHE_LIMIT && + !this.creditAlertState.has(userUuid) + ) { + const oldest = this.creditAlertState.keys().next().value; + if (oldest !== undefined) this.creditAlertState.delete(oldest); + } + this.creditAlertState.set(userUuid, state); + if (state === 'ok') return; + + try { + this.clients.event.emit( + 'metering.credit-state', + { + user_uuid: userUuid, + state, + allowance_used: allowanceUsed, + month_usage_allowance: monthUsageAllowance, + }, + {}, + ); + } catch (e) { + console.warn('[metering] credit-state emit failed:', e); + } } private rememberHasCredits(userId: string, hasCredits: boolean): void { diff --git a/src/backend/services/team/CreditState.test.ts b/src/backend/services/team/CreditState.test.ts new file mode 100644 index 000000000..c67fbfaea --- /dev/null +++ b/src/backend/services/team/CreditState.test.ts @@ -0,0 +1,180 @@ +/** + * Copyright (C) 2024-present Puter Technologies Inc. + * + * This file is part of Puter. + * + * Puter is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published + * by the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +import { v4 as uuidv4 } from 'uuid'; +import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest'; +import type { Actor } from '../../core/actor'; +import { PuterServer } from '../../server.ts'; +import { setupTestServer } from '../../testUtil.ts'; + +describe('credit-state transitions', () => { + let server: PuterServer; + let service: PuterServer['services']['team']; + let owner: { id: number }; + + /** Resolved per seat: a seat is on the common tier like any account. */ + const allowanceOf = async (actor: Actor) => + (await server.services.metering.getActorSubscription(actor)) + .monthUsageAllowance; + + /** Credit-state events seen since the last reset. */ + const states: Array<{ state: string; user_uuid: string }> = []; + + const makeUser = async () => { + const username = `cr_${Math.random().toString(36).slice(2, 10)}`; + const created = (await server.stores.user.create({ + username, + uuid: uuidv4(), + password: null, + email: `${username}@test.local`, + })) as unknown as { id: number }; + return { id: created.id, username }; + }; + + const freeHandle = () => `cw-${Math.random().toString(36).slice(2, 10)}`; + + /** + * A team with its own owner and one real provisioned seat. The owner + * is per-test so mail is attributable: these handlers run detached, and a + * shared recipient would let a late send land in the next test's tally. + */ + const makeSeat = async () => { + const owner = await makeUser(); + const team = await service.createTeam(owner.id, { + name: 'Credit Co', + handle: freeHandle(), + }); + const username = `seat_${Math.random().toString(36).slice(2, 10)}`; + const created = await service.provisionAccount(team.uid, owner.id, { + username, + email: `${username}@test.local`, + }); + const user = await server.stores.user.getById(created.userId); + const ownerRow = await server.stores.user.getById(owner.id); + return { + team, + user: user!, + owner: ownerRow!, + actor: { user } as unknown as Actor, + }; + }; + + /** Spend `fraction` of the seat's own monthly allowance. */ + const spend = async (actor: Actor, fraction: number) => + server.services.metering.incrementUsage( + actor, + 'kv:read', + 1, + Math.floor((await allowanceOf(actor)) * fraction), + ); + + /** Forget what this process has announced, as a second node would not know. */ + const asAnotherNode = (uuid: string) => { + ( + server.services.metering as unknown as { + creditAlertState: Map; + } + ).creditAlertState.delete(uuid); + }; + + beforeAll(async () => { + // The cap has its own suite; these tests need many teams. + server = await setupTestServer({ + teams_enabled: true, + max_teams_per_user: 100, + } as never); + service = server.services.team; + owner = await makeUser(); + + server.clients.event.on('metering.credit-state', (( + _k: string, + d: { state: string; user_uuid: string }, + ) => { + states.push(d); + }) as never); + + }); + + beforeEach(() => { + states.length = 0; + }); + + afterAll(async () => { + vi.restoreAllMocks(); + await server?.shutdown(); + }); + + // -- the transition signal ---------------------------------------- + + it('announces near-limit when the member passes 90%', async () => { + const { actor, user } = await makeSeat(); + states.length = 0; + + await spend(actor, 0.95); + + const mine = states.filter((s) => s.user_uuid === user.uuid); + expect(mine.map((s) => s.state)).toEqual(['near-limit']); + }); + + it('does not warn a member whose purchased credits carry them past the allowance', async () => { + const { actor, user } = await makeSeat(); + const allowance = await allowanceOf(actor); + await server.services.metering.updateAddonCredit(user.uuid, allowance); + states.length = 0; + + await spend(actor, 0.95); + + expect(states.filter((s) => s.user_uuid === user.uuid)).toHaveLength(0); + }); + + it('says nothing while the member is comfortably inside the allowance', async () => { + const { actor, user } = await makeSeat(); + states.length = 0; + + await spend(actor, 0.5); + + expect(states.filter((s) => s.user_uuid === user.uuid)).toHaveLength(0); + }); + + it('announces exhausted once, however many times the member retries', async () => { + const { actor, user } = await makeSeat(); + await spend(actor, 1); + states.length = 0; + + // Increments, not cached reads: a cache hit never recomputes the + // state, so a read loop would pass this test without exercising it. + for (let i = 0; i < 50; i++) { + await spend(actor, 0.01); + } + + expect(states.filter((s) => s.user_uuid === user.uuid)).toHaveLength(0); + }); + + it('crosses both lines in order, one announcement each', async () => { + const { actor, user } = await makeSeat(); + states.length = 0; + + await spend(actor, 0.95); + await spend(actor, 0.1); + + const mine = states.filter((s) => s.user_uuid === user.uuid); + expect(mine.map((s) => s.state)).toEqual(['near-limit', 'exhausted']); + }); + +});