From c61cc6a86f5f8364732296f0bb8393b37e0f70b3 Mon Sep 17 00:00:00 2001 From: DarkSky Date: Wed, 15 Jul 2026 04:38:49 +0800 Subject: [PATCH] fix(server): missing query cache --- .../core/permission/__tests__/policy.spec.ts | 15 +++++ .../server/src/core/permission/policy.ts | 19 ++++++- .../src/core/quota/__tests__/state.spec.ts | 25 ++++++++- .../backend/server/src/core/quota/state.ts | 56 ++++++++++++++++++- packages/backend/server/src/models/blob.ts | 13 ++++- .../server/src/models/comment-attachment.ts | 12 +++- 6 files changed, 132 insertions(+), 8 deletions(-) diff --git a/packages/backend/server/src/core/permission/__tests__/policy.spec.ts b/packages/backend/server/src/core/permission/__tests__/policy.spec.ts index 0eec8dbb3a..35336bf565 100644 --- a/packages/backend/server/src/core/permission/__tests__/policy.spec.ts +++ b/packages/backend/server/src/core/permission/__tests__/policy.spec.ts @@ -138,9 +138,24 @@ test('should keep owned workspace writable when quota is within limit', async t const state = await t.context.policy.reconcileWorkspaceQuotaState( workspace.id ); + const quotaState = Reflect.get( + t.context.policy, + 'quotaState' + ) as QuotaStateService; + const reconcile = Sinon.spy(quotaState, 'reconcileWorkspaceQuotaState'); t.false(state.isReadonly); t.deepEqual(state.readonlyReasons, []); + + await t.context.policy.getWorkspaceState(workspace.id); + t.is(reconcile.callCount, 0); + + await t.context.db.effectiveWorkspaceQuotaState.update({ + where: { workspaceId: workspace.id }, + data: { stale: true }, + }); + await t.context.policy.getWorkspaceState(workspace.id); + t.is(reconcile.callCount, 1); }); test('should report readonly state when fallback owner member quota overflows', async t => { diff --git a/packages/backend/server/src/core/permission/policy.ts b/packages/backend/server/src/core/permission/policy.ts index 13a21ffd22..f12d9d4631 100644 --- a/packages/backend/server/src/core/permission/policy.ts +++ b/packages/backend/server/src/core/permission/policy.ts @@ -37,8 +37,21 @@ export class WorkspacePolicyService { ) {} async getWorkspaceState(workspaceId: string): Promise { + const current = await this.quotaState.getWorkspaceQuotaState(workspaceId); const quota = - await this.quotaState.reconcileWorkspaceQuotaState(workspaceId); + current?.known && + !current.stale && + (!current.staleAfter || current.staleAfter > new Date()) + ? current + : await this.quotaState.reconcileWorkspaceQuotaState(workspaceId); + return this.toWorkspaceState(quota); + } + + private toWorkspaceState( + quota: Awaited< + ReturnType + > + ): WorkspaceState { const quotaSnapshot = quota as WorkspaceQuotaSnapshot; const readonlyReasons = quotaSnapshot.readonlyReasons; @@ -67,7 +80,9 @@ export class WorkspacePolicyService { } async reconcileWorkspaceQuotaState(workspaceId: string) { - return await this.getWorkspaceState(workspaceId); + return this.toWorkspaceState( + await this.quotaState.reconcileWorkspaceQuotaState(workspaceId) + ); } async handleTeamPlanCanceled(workspaceId: string) { diff --git a/packages/backend/server/src/core/quota/__tests__/state.spec.ts b/packages/backend/server/src/core/quota/__tests__/state.spec.ts index 97037d4f3c..0c4afdce4e 100644 --- a/packages/backend/server/src/core/quota/__tests__/state.spec.ts +++ b/packages/backend/server/src/core/quota/__tests__/state.spec.ts @@ -2,6 +2,7 @@ import { randomUUID } from 'node:crypto'; import { PrismaClient } from '@prisma/client'; import ava, { ExecutionContext, TestFn } from 'ava'; +import Sinon from 'sinon'; import { createTestingModule, @@ -80,11 +81,19 @@ test('workspace quota state uses current member rows', async t => { }, }); - const state = await t.context.state.reconcileWorkspaceQuotaState( - workspace.id + const resolveWorkspaceEntitlement = Sinon.spy( + t.context.entitlement, + 'resolveWorkspaceEntitlement' ); + const [state] = await Promise.all([ + t.context.state.reconcileWorkspaceQuotaState(workspace.id), + t.context.state.reconcileWorkspaceQuotaState(workspace.id), + t.context.state.reconcileWorkspaceQuotaState(workspace.id), + ]); t.is(state.memberCount, 1); + t.is(resolveWorkspaceEntitlement.callCount, 2); + resolveWorkspaceEntitlement.restore(); }); test('quota service exposes history period in seconds', async t => { @@ -131,10 +140,20 @@ test('quota state reconcile does not publish unchanged snapshots', async t => { } }); - await t.context.state.reconcileUserQuotaState(user.id); + const getActiveEntitlements = Sinon.spy( + t.context.entitlement, + 'getActiveEntitlements' + ); + await Promise.all([ + t.context.state.reconcileUserQuotaState(user.id), + t.context.state.reconcileUserQuotaState(user.id), + t.context.state.reconcileUserQuotaState(user.id), + ]); await t.context.state.reconcileUserQuotaState(user.id); t.is(changes, 1); + t.is(getActiveEntitlements.callCount, 2); + getActiveEntitlements.restore(); }); test('workspace quota state requires owner from new permission table', async t => { diff --git a/packages/backend/server/src/core/quota/state.ts b/packages/backend/server/src/core/quota/state.ts index e02e550cd9..b7da502342 100644 --- a/packages/backend/server/src/core/quota/state.ts +++ b/packages/backend/server/src/core/quota/state.ts @@ -1,5 +1,9 @@ import { Injectable } from '@nestjs/common'; -import { PrismaClient } from '@prisma/client'; +import { + type EffectiveUserQuotaState, + type EffectiveWorkspaceQuotaState, + PrismaClient, +} from '@prisma/client'; import { EventBus, OnEvent } from '../../base'; import { EntitlementService } from '../entitlement'; @@ -23,6 +27,15 @@ declare global { @Injectable() export class QuotaStateService { + private readonly userReconciliations = new Map< + string, + Promise + >(); + private readonly workspaceReconciliations = new Map< + string, + Promise + >(); + constructor( private readonly db: PrismaClient, private readonly entitlement: EntitlementService, @@ -38,6 +51,16 @@ export class QuotaStateService { async reconcileUserQuotaState( userId: string, options: { emit?: boolean } = {} + ) { + const key = `${userId}:${options.emit ?? true}`; + return await this.reconcileOnce(this.userReconciliations, key, () => + this.reconcileUserQuotaStateNow(userId, options) + ); + } + + private async reconcileUserQuotaStateNow( + userId: string, + options: { emit?: boolean } ) { const [previous, entitlement, entitlements, resolved, usedStorageQuota] = await Promise.all([ @@ -80,6 +103,37 @@ export class QuotaStateService { async reconcileWorkspaceQuotaState( workspaceId: string, options: { emit?: boolean } = {} + ) { + const key = `${workspaceId}:${options.emit ?? true}`; + return await this.reconcileOnce(this.workspaceReconciliations, key, () => + this.reconcileWorkspaceQuotaStateNow(workspaceId, options) + ); + } + + private async reconcileOnce( + reconciliations: Map>, + key: string, + reconcile: () => Promise + ) { + const existing = reconciliations.get(key); + if (existing) { + return await existing; + } + + const reconciliation = reconcile(); + reconciliations.set(key, reconciliation); + try { + return await reconciliation; + } finally { + if (reconciliations.get(key) === reconciliation) { + reconciliations.delete(key); + } + } + } + + private async reconcileWorkspaceQuotaStateNow( + workspaceId: string, + options: { emit?: boolean } ) { const owner = await this.getWorkspaceOwner(workspaceId); const [ diff --git a/packages/backend/server/src/models/blob.ts b/packages/backend/server/src/models/blob.ts index 5ae22df41e..c2b76ece0c 100644 --- a/packages/backend/server/src/models/blob.ts +++ b/packages/backend/server/src/models/blob.ts @@ -11,7 +11,7 @@ export type CreateBlobInput = Prisma.BlobUncheckedCreateInput; @Injectable() export class BlobModel extends BaseModel { async upsert(blob: CreateBlobInput) { - return await this.db.blob.upsert({ + const result = await this.db.blob.upsert({ where: { workspaceId_key: { workspaceId: blob.workspaceId, @@ -33,6 +33,8 @@ export class BlobModel extends BaseModel { uploadId: blob.uploadId, }, }); + await this.markQuotaStateStale(blob.workspaceId); + return result; } async delete(workspaceId: string, key: string, permanently = false) { @@ -43,6 +45,7 @@ export class BlobModel extends BaseModel { key, }, }); + await this.markQuotaStateStale(workspaceId); this.logger.log(`deleted blob ${workspaceId}/${key} permanently`); return; } @@ -58,6 +61,7 @@ export class BlobModel extends BaseModel { deletedAt: new Date(), }, }); + await this.markQuotaStateStale(workspaceId); } async get(workspaceId: string, key: string) { @@ -109,4 +113,11 @@ export class BlobModel extends BaseModel { return sum._sum.size ?? 0; } + + private async markQuotaStateStale(workspaceId: string) { + await this.db.effectiveWorkspaceQuotaState.updateMany({ + where: { workspaceId }, + data: { stale: true }, + }); + } } diff --git a/packages/backend/server/src/models/comment-attachment.ts b/packages/backend/server/src/models/comment-attachment.ts index db98ae2581..4f6c65c626 100644 --- a/packages/backend/server/src/models/comment-attachment.ts +++ b/packages/backend/server/src/models/comment-attachment.ts @@ -12,7 +12,7 @@ export type CreateCommentAttachmentInput = @Injectable() export class CommentAttachmentModel extends BaseModel { async upsert(input: CreateCommentAttachmentInput) { - return await this.db.commentAttachment.upsert({ + const result = await this.db.commentAttachment.upsert({ where: { workspaceId_docId_key: { workspaceId: input.workspaceId, @@ -35,6 +35,8 @@ export class CommentAttachmentModel extends BaseModel { createdBy: input.createdBy, }, }); + await this.markQuotaStateStale(input.workspaceId); + return result; } async delete(workspaceId: string, docId: string, key: string) { @@ -45,6 +47,7 @@ export class CommentAttachmentModel extends BaseModel { key, }, }); + await this.markQuotaStateStale(workspaceId); this.logger.log(`deleted comment attachment ${workspaceId}/${key}`); } @@ -68,4 +71,11 @@ export class CommentAttachmentModel extends BaseModel { }, }); } + + private async markQuotaStateStale(workspaceId: string) { + await this.db.effectiveWorkspaceQuotaState.updateMany({ + where: { workspaceId }, + data: { stale: true }, + }); + } }