fix(server): missing query cache

This commit is contained in:
DarkSky
2026-07-15 04:38:49 +08:00
parent 3bbc890bcb
commit c61cc6a86f
6 changed files with 132 additions and 8 deletions

View File

@@ -138,9 +138,24 @@ test('should keep owned workspace writable when quota is within limit', async t
const state = await t.context.policy.reconcileWorkspaceQuotaState( const state = await t.context.policy.reconcileWorkspaceQuotaState(
workspace.id workspace.id
); );
const quotaState = Reflect.get(
t.context.policy,
'quotaState'
) as QuotaStateService;
const reconcile = Sinon.spy(quotaState, 'reconcileWorkspaceQuotaState');
t.false(state.isReadonly); t.false(state.isReadonly);
t.deepEqual(state.readonlyReasons, []); 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 => { test('should report readonly state when fallback owner member quota overflows', async t => {

View File

@@ -37,8 +37,21 @@ export class WorkspacePolicyService {
) {} ) {}
async getWorkspaceState(workspaceId: string): Promise<WorkspaceState> { async getWorkspaceState(workspaceId: string): Promise<WorkspaceState> {
const current = await this.quotaState.getWorkspaceQuotaState(workspaceId);
const quota = 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<QuotaStateService['reconcileWorkspaceQuotaState']>
>
): WorkspaceState {
const quotaSnapshot = quota as WorkspaceQuotaSnapshot; const quotaSnapshot = quota as WorkspaceQuotaSnapshot;
const readonlyReasons = quotaSnapshot.readonlyReasons; const readonlyReasons = quotaSnapshot.readonlyReasons;
@@ -67,7 +80,9 @@ export class WorkspacePolicyService {
} }
async reconcileWorkspaceQuotaState(workspaceId: string) { async reconcileWorkspaceQuotaState(workspaceId: string) {
return await this.getWorkspaceState(workspaceId); return this.toWorkspaceState(
await this.quotaState.reconcileWorkspaceQuotaState(workspaceId)
);
} }
async handleTeamPlanCanceled(workspaceId: string) { async handleTeamPlanCanceled(workspaceId: string) {

View File

@@ -2,6 +2,7 @@ import { randomUUID } from 'node:crypto';
import { PrismaClient } from '@prisma/client'; import { PrismaClient } from '@prisma/client';
import ava, { ExecutionContext, TestFn } from 'ava'; import ava, { ExecutionContext, TestFn } from 'ava';
import Sinon from 'sinon';
import { import {
createTestingModule, createTestingModule,
@@ -80,11 +81,19 @@ test('workspace quota state uses current member rows', async t => {
}, },
}); });
const state = await t.context.state.reconcileWorkspaceQuotaState( const resolveWorkspaceEntitlement = Sinon.spy(
workspace.id 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(state.memberCount, 1);
t.is(resolveWorkspaceEntitlement.callCount, 2);
resolveWorkspaceEntitlement.restore();
}); });
test('quota service exposes history period in seconds', async t => { 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); await t.context.state.reconcileUserQuotaState(user.id);
t.is(changes, 1); t.is(changes, 1);
t.is(getActiveEntitlements.callCount, 2);
getActiveEntitlements.restore();
}); });
test('workspace quota state requires owner from new permission table', async t => { test('workspace quota state requires owner from new permission table', async t => {

View File

@@ -1,5 +1,9 @@
import { Injectable } from '@nestjs/common'; 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 { EventBus, OnEvent } from '../../base';
import { EntitlementService } from '../entitlement'; import { EntitlementService } from '../entitlement';
@@ -23,6 +27,15 @@ declare global {
@Injectable() @Injectable()
export class QuotaStateService { export class QuotaStateService {
private readonly userReconciliations = new Map<
string,
Promise<EffectiveUserQuotaState>
>();
private readonly workspaceReconciliations = new Map<
string,
Promise<EffectiveWorkspaceQuotaState>
>();
constructor( constructor(
private readonly db: PrismaClient, private readonly db: PrismaClient,
private readonly entitlement: EntitlementService, private readonly entitlement: EntitlementService,
@@ -38,6 +51,16 @@ export class QuotaStateService {
async reconcileUserQuotaState( async reconcileUserQuotaState(
userId: string, userId: string,
options: { emit?: boolean } = {} 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] = const [previous, entitlement, entitlements, resolved, usedStorageQuota] =
await Promise.all([ await Promise.all([
@@ -80,6 +103,37 @@ export class QuotaStateService {
async reconcileWorkspaceQuotaState( async reconcileWorkspaceQuotaState(
workspaceId: string, workspaceId: string,
options: { emit?: boolean } = {} options: { emit?: boolean } = {}
) {
const key = `${workspaceId}:${options.emit ?? true}`;
return await this.reconcileOnce(this.workspaceReconciliations, key, () =>
this.reconcileWorkspaceQuotaStateNow(workspaceId, options)
);
}
private async reconcileOnce<T>(
reconciliations: Map<string, Promise<T>>,
key: string,
reconcile: () => Promise<T>
) {
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 owner = await this.getWorkspaceOwner(workspaceId);
const [ const [

View File

@@ -11,7 +11,7 @@ export type CreateBlobInput = Prisma.BlobUncheckedCreateInput;
@Injectable() @Injectable()
export class BlobModel extends BaseModel { export class BlobModel extends BaseModel {
async upsert(blob: CreateBlobInput) { async upsert(blob: CreateBlobInput) {
return await this.db.blob.upsert({ const result = await this.db.blob.upsert({
where: { where: {
workspaceId_key: { workspaceId_key: {
workspaceId: blob.workspaceId, workspaceId: blob.workspaceId,
@@ -33,6 +33,8 @@ export class BlobModel extends BaseModel {
uploadId: blob.uploadId, uploadId: blob.uploadId,
}, },
}); });
await this.markQuotaStateStale(blob.workspaceId);
return result;
} }
async delete(workspaceId: string, key: string, permanently = false) { async delete(workspaceId: string, key: string, permanently = false) {
@@ -43,6 +45,7 @@ export class BlobModel extends BaseModel {
key, key,
}, },
}); });
await this.markQuotaStateStale(workspaceId);
this.logger.log(`deleted blob ${workspaceId}/${key} permanently`); this.logger.log(`deleted blob ${workspaceId}/${key} permanently`);
return; return;
} }
@@ -58,6 +61,7 @@ export class BlobModel extends BaseModel {
deletedAt: new Date(), deletedAt: new Date(),
}, },
}); });
await this.markQuotaStateStale(workspaceId);
} }
async get(workspaceId: string, key: string) { async get(workspaceId: string, key: string) {
@@ -109,4 +113,11 @@ export class BlobModel extends BaseModel {
return sum._sum.size ?? 0; return sum._sum.size ?? 0;
} }
private async markQuotaStateStale(workspaceId: string) {
await this.db.effectiveWorkspaceQuotaState.updateMany({
where: { workspaceId },
data: { stale: true },
});
}
} }

View File

@@ -12,7 +12,7 @@ export type CreateCommentAttachmentInput =
@Injectable() @Injectable()
export class CommentAttachmentModel extends BaseModel { export class CommentAttachmentModel extends BaseModel {
async upsert(input: CreateCommentAttachmentInput) { async upsert(input: CreateCommentAttachmentInput) {
return await this.db.commentAttachment.upsert({ const result = await this.db.commentAttachment.upsert({
where: { where: {
workspaceId_docId_key: { workspaceId_docId_key: {
workspaceId: input.workspaceId, workspaceId: input.workspaceId,
@@ -35,6 +35,8 @@ export class CommentAttachmentModel extends BaseModel {
createdBy: input.createdBy, createdBy: input.createdBy,
}, },
}); });
await this.markQuotaStateStale(input.workspaceId);
return result;
} }
async delete(workspaceId: string, docId: string, key: string) { async delete(workspaceId: string, docId: string, key: string) {
@@ -45,6 +47,7 @@ export class CommentAttachmentModel extends BaseModel {
key, key,
}, },
}); });
await this.markQuotaStateStale(workspaceId);
this.logger.log(`deleted comment attachment ${workspaceId}/${key}`); 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 },
});
}
} }