fix(server): wrap read-modify-write apis with distributed lock (#6142)

packages/backend/server/src/core/features/feature.ts

@@ -1,5 +1,4 @@
-import { PrismaClient } from '@prisma/client';
-
+import { PrismaTransaction } from '../../fundamentals';
 import { Feature, FeatureSchema, FeatureType } from './types';
 
 class FeatureConfig {
@@ -67,7 +66,7 @@ export type FeatureConfigType<F extends FeatureType> = InstanceType<
 
 const FeatureCache = new Map<number, FeatureConfigType<FeatureType>>();
 
-export async function getFeature(prisma: PrismaClient, featureId: number) {
+export async function getFeature(prisma: PrismaTransaction, featureId: number) {
   const cachedQuota = FeatureCache.get(featureId);
 
   if (cachedQuota) {

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

@@ -1,5 +1,4 @@
-import { PrismaClient } from '@prisma/client';
-
+import { PrismaTransaction } from '../../fundamentals';
 import { formatDate, formatSize, Quota, QuotaSchema } from './types';
 
 const QuotaCache = new Map<number, QuotaConfig>();
@@ -7,14 +6,14 @@ const QuotaCache = new Map<number, QuotaConfig>();
 export class QuotaConfig {
   readonly config: Quota;
 
-  static async get(prisma: PrismaClient, featureId: number) {
+  static async get(tx: PrismaTransaction, featureId: number) {
     const cachedQuota = QuotaCache.get(featureId);
 
     if (cachedQuota) {
       return cachedQuota;
     }
 
-    const quota = await prisma.features.findFirst({
+    const quota = await tx.features.findFirst({
       where: {
         id: featureId,
       },

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

@@ -1,13 +1,15 @@
 import { Injectable } from '@nestjs/common';
 import { PrismaClient } from '@prisma/client';
 
-import { type EventPayload, OnEvent } from '../../fundamentals';
+import {
+  type EventPayload,
+  OnEvent,
+  PrismaTransaction,
+} from '../../fundamentals';
 import { FeatureKind } from '../features';
 import { QuotaConfig } from './quota';
 import { QuotaType } from './types';
 
-type Transaction = Parameters<Parameters<PrismaClient['$transaction']>[0]>[0];
-
 @Injectable()
 export class QuotaService {
   constructor(private readonly prisma: PrismaClient) {}
@@ -140,8 +142,8 @@ export class QuotaService {
     });
   }
 
-  async hasQuota(userId: string, quota: QuotaType, transaction?: Transaction) {
-    const executor = transaction ?? this.prisma;
+  async hasQuota(userId: string, quota: QuotaType, tx?: PrismaTransaction) {
+    const executor = tx ?? this.prisma;
 
     return executor.userFeatures
       .count({

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

@@ -54,7 +54,7 @@ export class UserService {
 
     return this.createUser({
       email,
-      name: 'Unnamed',
+      name: email.split('@')[0],
       ...data,
     });
   }

packages/backend/server/src/core/workspaces/resolvers/workspace.ts

@@ -25,7 +25,9 @@ import {
   EventEmitter,
   type FileUpload,
   MailService,
+  MutexService,
   Throttle,
+  TooManyRequestsException,
 } from '../../../fundamentals';
 import { CurrentUser, Public } from '../../auth';
 import { QuotaManagementService, QuotaQueryType } from '../../quota';
@@ -58,7 +60,8 @@ export class WorkspaceResolver {
     private readonly quota: QuotaManagementService,
     private readonly users: UserService,
     private readonly event: EventEmitter,
-    private readonly blobStorage: WorkspaceBlobStorage
+    private readonly blobStorage: WorkspaceBlobStorage,
+    private readonly mutex: MutexService
   ) {}
 
   @ResolveField(() => Permission, {
@@ -336,74 +339,87 @@ export class WorkspaceResolver {
       throw new ForbiddenException('Cannot change owner');
     }
 
-    // member limit check
-    const [memberCount, quota] = await Promise.all([
-      this.prisma.workspaceUserPermission.count({
-        where: { workspaceId },
-      }),
-      this.quota.getWorkspaceUsage(workspaceId),
-    ]);
-    if (memberCount >= quota.memberLimit) {
-      throw new PayloadTooLargeException('Workspace member limit reached.');
-    }
+    try {
+      // lock to prevent concurrent invite
+      const lockFlag = `invite:${workspaceId}`;
+      await using lock = await this.mutex.lock(lockFlag);
+      if (!lock) {
+        return new TooManyRequestsException('Server is busy');
+      }
 
-    let target = await this.users.findUserByEmail(email);
-    if (target) {
-      const originRecord = await this.prisma.workspaceUserPermission.findFirst({
-        where: {
-          workspaceId,
-          userId: target.id,
-        },
-      });
-      // only invite if the user is not already in the workspace
-      if (originRecord) return originRecord.id;
-    } else {
-      target = await this.users.createAnonymousUser(email, {
-        registered: false,
-      });
-    }
+      // member limit check
+      const [memberCount, quota] = await Promise.all([
+        this.prisma.workspaceUserPermission.count({
+          where: { workspaceId },
+        }),
+        this.quota.getWorkspaceUsage(workspaceId),
+      ]);
+      if (memberCount >= quota.memberLimit) {
+        return new PayloadTooLargeException('Workspace member limit reached.');
+      }
 
-    const inviteId = await this.permissions.grant(
-      workspaceId,
-      target.id,
-      permission
-    );
-    if (sendInviteMail) {
-      const inviteInfo = await this.getInviteInfo(inviteId);
-
-      try {
-        await this.mailer.sendInviteEmail(email, inviteId, {
-          workspace: {
-            id: inviteInfo.workspace.id,
-            name: inviteInfo.workspace.name,
-            avatar: inviteInfo.workspace.avatar,
-          },
-          user: {
-            avatar: inviteInfo.user?.avatarUrl || '',
-            name: inviteInfo.user?.name || '',
-          },
+      let target = await this.users.findUserByEmail(email);
+      if (target) {
+        const originRecord =
+          await this.prisma.workspaceUserPermission.findFirst({
+            where: {
+              workspaceId,
+              userId: target.id,
+            },
+          });
+        // only invite if the user is not already in the workspace
+        if (originRecord) return originRecord.id;
+      } else {
+        target = await this.users.createAnonymousUser(email, {
+          registered: false,
         });
-      } catch (e) {
-        const ret = await this.permissions.revokeWorkspace(
-          workspaceId,
-          target.id
-        );
+      }
 
-        if (!ret) {
-          this.logger.fatal(
-            `failed to send ${workspaceId} invite email to ${email} and failed to revoke permission: ${inviteId}, ${e}`
+      const inviteId = await this.permissions.grant(
+        workspaceId,
+        target.id,
+        permission
+      );
+      if (sendInviteMail) {
+        const inviteInfo = await this.getInviteInfo(inviteId);
+
+        try {
+          await this.mailer.sendInviteEmail(email, inviteId, {
+            workspace: {
+              id: inviteInfo.workspace.id,
+              name: inviteInfo.workspace.name,
+              avatar: inviteInfo.workspace.avatar,
+            },
+            user: {
+              avatar: inviteInfo.user?.avatarUrl || '',
+              name: inviteInfo.user?.name || '',
+            },
+          });
+        } catch (e) {
+          const ret = await this.permissions.revokeWorkspace(
+            workspaceId,
+            target.id
           );
-        } else {
-          this.logger.warn(
-            `failed to send ${workspaceId} invite email to ${email}, but successfully revoked permission: ${e}`
+
+          if (!ret) {
+            this.logger.fatal(
+              `failed to send ${workspaceId} invite email to ${email} and failed to revoke permission: ${inviteId}, ${e}`
+            );
+          } else {
+            this.logger.warn(
+              `failed to send ${workspaceId} invite email to ${email}, but successfully revoked permission: ${e}`
+            );
+          }
+          return new InternalServerErrorException(
+            'Failed to send invite email. Please try again.'
           );
         }
-        return new InternalServerErrorException(
-          'Failed to send invite email. Please try again.'
-        );
       }
+      return inviteId;
+    } catch (e) {
+      this.logger.error('failed to invite user', e);
+      return new TooManyRequestsException('Server is busy');
     }
-    return inviteId;
   }
 
   @Throttle({