fix(server): doc upsert race condition (#5755)

packages/backend/server/src/core/doc/manager.ts

@@ -19,6 +19,7 @@ import {
 
 import {
   Cache,
+  CallTimer,
   Config,
   EventEmitter,
   type EventPayload,
@@ -463,6 +464,7 @@ export class DocManager implements OnModuleInit, OnModuleDestroy {
     });
   }
 
+  @CallTimer('doc', 'upsert')
   private async upsert(
     workspaceId: string,
     guid: string,
@@ -472,7 +474,6 @@ export class DocManager implements OnModuleInit, OnModuleDestroy {
     updatedAt: Date,
     initialSeq?: number
   ) {
-    return this.lockSnapshotForUpsert(workspaceId, guid, async () => {
     const blob = Buffer.from(encodeStateAsUpdate(doc));
 
     if (isEmptyBuffer(blob)) {
@@ -481,21 +482,31 @@ export class DocManager implements OnModuleInit, OnModuleDestroy {
 
     const state = Buffer.from(encodeStateVector(doc));
 
-      return await this.db.$transaction(async db => {
+    await this.db.$queryRaw`BEGIN;`;
-        const snapshot = await db.snapshot.findUnique({
+    let committed = false;
-          where: {
+    const commit = async () => {
-            id_workspaceId: {
+      if (!committed) {
-              id: guid,
+        committed = true;
-              workspaceId,
+        await this.db.$queryRaw`COMMIT;`;
-            },
+      }
-          },
+    };
-        });
+    try {
+      const [snapshot]: {
+        workspace_id: string;
+        id: string;
+        blob: Buffer;
+        state?: Buffer;
+      }[] = await this.db.$queryRaw`
+        -- LOCK TABLE "Snapshot" IN SHARE ROW EXCLUSIVE MODE;
+        SELECT * FROM snapshots WHERE workspace_id = ${workspaceId} AND guid = ${guid} limit 1
+        FOR UPDATE;
+      `;
 
       // update
       if (snapshot) {
         // only update if state is newer
         if (isStateNewer(snapshot.state ?? Buffer.from([0]), state)) {
-            await db.snapshot.update({
+          await this.db.snapshot.update({
             select: {
               seq: true,
             },
@@ -518,7 +529,9 @@ export class DocManager implements OnModuleInit, OnModuleDestroy {
         }
       } else {
         // create
-          await db.snapshot.create({
+        // no record exists, should commit the previous row lock first
+        await commit();
+        await this.db.snapshot.create({
           select: {
             seq: true,
           },
@@ -535,10 +548,13 @@ export class DocManager implements OnModuleInit, OnModuleDestroy {
 
         return true;
       }
-      });
+    } catch (e) {
-    });
+      await this.db.$queryRaw`ROLLBACK;`;
+      throw e;
+    } finally {
+      await commit();
+    }
   }
-
   private async _get(
     workspaceId: string,
     guid: string
@@ -559,6 +575,7 @@ export class DocManager implements OnModuleInit, OnModuleDestroy {
    * Squash updates into a single update and save it as snapshot,
    * and delete the updates records at the same time.
    */
+  @CallTimer('doc', 'squash')
   private async squash(updates: Update[], snapshot: Snapshot | null) {
     if (!updates.length) {
       throw new Error('No updates to squash');
@@ -761,18 +778,6 @@ export class DocManager implements OnModuleInit, OnModuleDestroy {
     );
   }
 
-  async lockSnapshotForUpsert<T>(
-    workspaceId: string,
-    guid: string,
-    job: () => Promise<T>
-  ) {
-    return this.doWithLock(
-      'doc:manager:snapshot',
-      `${workspaceId}::${guid}`,
-      job
-    );
-  }
-
   @Cron(CronExpression.EVERY_MINUTE)
   async reportUpdatesQueueCount() {
     metrics.doc