feat(server): distinguash workspace quota calc (#9338)

fix AF-2021

packages/backend/server/src/core/quota/service.ts

@@ -270,6 +270,23 @@ export class QuotaService {
       .then(count => count > 0);
   }
 
+  /// check if workspaces have quota
+  /// return workspaces's id that have quota
+  async hasWorkspacesQuota(
+    workspaces: string[],
+    quota?: QuotaType
+  ): Promise<string[]> {
+    const workspaceIds = await this.prisma.workspaceFeature.findMany({
+      where: {
+        workspaceId: { in: workspaces },
+        feature: { feature: quota, type: FeatureKind.Quota },
+        activated: true,
+      },
+      select: { workspaceId: true },
+    });
+    return Array.from(new Set(workspaceIds.map(w => w.workspaceId)));
+  }
+
   async getWorkspaceConfig<Q extends QuotaType>(
     workspaceId: string,
     type: Q

packages/backend/server/src/core/quota/storage.ts

@@ -79,15 +79,20 @@ export class QuotaManagementService {
 
   async getUserStorageUsage(userId: string) {
     const workspaces = await this.permissions.getOwnedWorkspaces(userId);
+    const workspacesWithQuota = await this.quota.hasWorkspacesQuota(workspaces);
 
     const sizes = await Promise.allSettled(
-      workspaces.map(workspace => this.storage.totalSize(workspace))
+      workspaces
+        .filter(w => !workspacesWithQuota.includes(w))
+        .map(workspace => this.storage.totalSize(workspace))
     );
 
     return sizes.reduce((total, size) => {
       if (size.status === 'fulfilled') {
-        if (Number.isSafeInteger(size.value)) {
-          return total + size.value;
+        // ensure that size is within the safe range of gql
+        const totalSize = total + size.value;
+        if (Number.isSafeInteger(totalSize)) {
+          return totalSize;
         } else {
           this.logger.error(`Workspace size is invalid: ${size.value}`);
         }
@@ -98,6 +103,17 @@ export class QuotaManagementService {
     }, 0);
   }
 
+  async getWorkspaceStorageUsage(workspaceId: string) {
+    const totalSize = this.storage.totalSize(workspaceId);
+    // ensure that size is within the safe range of gql
+    if (Number.isSafeInteger(totalSize)) {
+      return totalSize;
+    } else {
+      this.logger.error(`Workspace size is invalid: ${totalSize}`);
+    }
+    return 0;
+  }
+
   private generateQuotaCalculator(
     quota: number,
     blobLimit: number,
@@ -146,17 +162,23 @@ export class QuotaManagementService {
     );
   }
 
-  private async getWorkspaceQuota(userId: string, workspaceId: string) {
+  private async getWorkspaceQuota(
+    userId: string,
+    workspaceId: string
+  ): Promise<{ quota: QuotaConfig; fromUser: boolean }> {
     const { feature: workspaceQuota } =
       (await this.quota.getWorkspaceQuota(workspaceId)) || {};
     const { feature: userQuota } = await this.quota.getUserQuota(userId);
     if (workspaceQuota) {
-      return workspaceQuota.withOverride({
-        // override user quota with workspace quota
-        copilotActionLimit: userQuota.copilotActionLimit,
-      });
+      return {
+        quota: workspaceQuota.withOverride({
+          // override user quota with workspace quota
+          copilotActionLimit: userQuota.copilotActionLimit,
+        }),
+        fromUser: false,
+      };
     }
-    return userQuota;
+    return { quota: userQuota, fromUser: true };
   }
 
   async checkWorkspaceSeat(workspaceId: string, excludeSelf = false) {
@@ -173,17 +195,24 @@ export class QuotaManagementService {
     const memberCount =
       await this.permissions.getWorkspaceMemberCount(workspaceId);
     const {
-      name,
-      blobLimit,
-      businessBlobLimit,
-      historyPeriod,
-      memberLimit,
-      storageQuota,
-      copilotActionLimit,
-      humanReadable,
+      quota: {
+        name,
+        blobLimit,
+        businessBlobLimit,
+        historyPeriod,
+        memberLimit,
+        storageQuota,
+        copilotActionLimit,
+        humanReadable,
+      },
+      fromUser,
     } = await this.getWorkspaceQuota(owner.id, workspaceId);
-    // get all workspaces size of owner used
-    const usedSize = await this.getUserStorageUsage(owner.id);
+
+    const usedSize = fromUser
+      ? // get all workspaces size of owner used
+        await this.getUserStorageUsage(owner.id)
+      : // get workspace size
+        await this.getWorkspaceStorageUsage(workspaceId);
     // relax restrictions if workspace has unlimited feature
     // todo(@darkskygit): need a mechanism to allow feature as a middleware to edit quota
     const unlimited = await this.feature.hasWorkspaceFeature(

packages/backend/server/tests/quota.spec.ts

@@ -13,7 +13,7 @@ import {
 } from '../src/core/quota';
 import { OneGB, OneMB } from '../src/core/quota/constant';
 import { FreePlan, ProPlan } from '../src/core/quota/schema';
-import { StorageModule } from '../src/core/storage';
+import { StorageModule, WorkspaceBlobStorage } from '../src/core/storage';
 import { WorkspaceResolver } from '../src/core/workspaces/resolvers';
 import { createTestingModule } from './utils';
 import { WorkspaceResolverMock } from './utils/feature';
@@ -23,6 +23,7 @@ const test = ava as TestFn<{
   quota: QuotaService;
   quotaManager: QuotaManagementService;
   workspace: WorkspaceResolver;
+  workspaceBlob: WorkspaceBlobStorage;
   module: TestingModule;
 }>;
 
@@ -37,16 +38,12 @@ test.beforeEach(async t => {
     },
   });
 
-  const quota = module.get(QuotaService);
-  const quotaManager = module.get(QuotaManagementService);
-  const workspace = module.get(WorkspaceResolver);
-  const auth = module.get(AuthService);
-
   t.context.module = module;
-  t.context.quota = quota;
-  t.context.quotaManager = quotaManager;
-  t.context.workspace = workspace;
-  t.context.auth = auth;
+  t.context.auth = module.get(AuthService);
+  t.context.quota = module.get(QuotaService);
+  t.context.quotaManager = module.get(QuotaManagementService);
+  t.context.workspace = module.get(WorkspaceResolver);
+  t.context.workspaceBlob = module.get(WorkspaceBlobStorage);
 });
 
 test.afterEach.always(async t => {
@@ -165,3 +162,33 @@ test('should be able to override quota', async t => {
   t.is(wq3.storageQuota, 140 * OneGB, 'should be override to 120GB');
   t.is(wq3.memberLimit, 2, 'should be override to 1');
 });
+
+test('should be able to check with workspace quota', async t => {
+  const { auth, quotaManager, workspace, workspaceBlob } = t.context;
+
+  const u1 = await auth.signUp('test@affine.pro', '123456');
+  const w1 = await workspace.createWorkspace(u1, null);
+  const w2 = await workspace.createWorkspace(u1, null);
+  await quotaManager.addTeamWorkspace(w2.id, 'test');
+
+  {
+    const wq = await quotaManager.getWorkspaceUsage(w1.id);
+    t.is(wq.usedSize, 0, 'should be 0');
+  }
+
+  {
+    await workspaceBlob.put(w1.id, 'test', Buffer.from([0, 0]));
+    const wq1 = await quotaManager.getWorkspaceUsage(w1.id);
+    t.is(wq1.usedSize, 2, 'should be 2');
+    const wq2 = await quotaManager.getWorkspaceUsage(w2.id);
+    t.is(wq2.usedSize, 0, 'should be 0');
+  }
+
+  {
+    await workspaceBlob.put(w2.id, 'test', Buffer.from([0, 0, 0]));
+    const wq1 = await quotaManager.getWorkspaceUsage(w1.id);
+    t.is(wq1.usedSize, 2, 'should be 2');
+    const wq2 = await quotaManager.getWorkspaceUsage(w2.id);
+    t.is(wq2.usedSize, 3, 'should be 0');
+  }
+});