fix(server): doc upsert without row lock (#5765)

patch cmdk based on https://github.com/pengx17/cmdk/tree/patch-1
fix https://github.com/toeverything/AFFiNE/issues/5548

.github/deployment/front/Dockerfile

@@ -1,6 +1,6 @@
 FROM openresty/openresty:1.21.4.3-0-buster
 WORKDIR /app
-COPY ./packages/frontend/core/dist/index.html ./dist/index.html
+COPY ./packages/frontend/core/dist ./dist
 COPY ./.github/deployment/front/nginx.conf /usr/local/openresty/nginx/conf/nginx.conf
 COPY ./.github/deployment/front/affine.nginx.conf /etc/nginx/conf.d/affine.nginx.conf
 

packages/backend/server/schema.prisma

@@ -265,7 +265,9 @@ model Snapshot {
   seq         Int      @default(0) @db.Integer
   state       Bytes?   @db.ByteA
   createdAt   DateTime @default(now()) @map("created_at") @db.Timestamptz(6)
-  updatedAt   DateTime @updatedAt @map("updated_at") @db.Timestamptz(6)
+  // the `updated_at` field will not record the time of record changed,
+  // but the created time of last seen update that has been merged into snapshot.
+  updatedAt   DateTime @map("updated_at") @db.Timestamptz(6)
 
   @@id([id, workspaceId])
   @@map("snapshots")

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

@@ -10,7 +10,6 @@ import { chunk } from 'lodash-es';
 import { defer, retry } from 'rxjs';
 import {
   applyUpdate,
-  decodeStateVector,
   Doc,
   encodeStateAsUpdate,
   encodeStateVector,
@@ -19,6 +18,7 @@ import {
 
 import {
   Cache,
+  CallTimer,
   Config,
   EventEmitter,
   type EventPayload,
@@ -45,36 +45,6 @@ function compare(yBinary: Buffer, jwstBinary: Buffer, strict = false): boolean {
   return compare(yBinary, yBinary2, true);
 }
 
-/**
- * Detect whether rhs state is newer than lhs state.
- *
- * How could we tell a state is newer:
- *
- *   i. if the state vector size is larger, it's newer
- *  ii. if the state vector size is same, compare each client's state
- */
-function isStateNewer(lhs: Buffer, rhs: Buffer): boolean {
-  const lhsVector = decodeStateVector(lhs);
-  const rhsVector = decodeStateVector(rhs);
-
-  if (lhsVector.size < rhsVector.size) {
-    return true;
-  }
-
-  for (const [client, state] of lhsVector) {
-    const rstate = rhsVector.get(client);
-    if (!rstate) {
-      return false;
-    }
-
-    if (state < rstate) {
-      return true;
-    }
-  }
-
-  return false;
-}
-
 export function isEmptyBuffer(buf: Buffer): boolean {
   return (
     buf.length === 0 ||
@@ -119,6 +89,7 @@ export class DocManager implements OnModuleInit, OnModuleDestroy {
     this.destroy();
   }
 
+  @CallTimer('doc', 'yjs_recover_updates_to_doc')
   private recoverDoc(...updates: Buffer[]): Promise<Doc> {
     const doc = new Doc();
     const chunks = chunk(updates, 10);
@@ -382,7 +353,7 @@ export class DocManager implements OnModuleInit, OnModuleDestroy {
     const updates = await this.getUpdates(workspaceId, guid);
 
     if (updates.length) {
-      const doc = await this.squash(updates, snapshot);
+      const doc = await this.squash(snapshot, updates);
       return Buffer.from(encodeStateVector(doc));
     }
 
@@ -463,80 +434,92 @@ export class DocManager implements OnModuleInit, OnModuleDestroy {
     });
   }
 
+  /**
+   * @returns whether the snapshot is updated to the latest, `undefined` means the doc to be upserted is outdated.
+   */
+  @CallTimer('doc', 'upsert')
   private async upsert(
     workspaceId: string,
     guid: string,
     doc: Doc,
     // we always delay the snapshot update to avoid db overload,
-    // so the value of `updatedAt` will not be accurate to user's real action time
+    // so the value of auto updated `updatedAt` by db will never be accurate to user's real action time
     updatedAt: Date,
-    initialSeq?: number
+    seq: number
   ) {
-    return this.lockSnapshotForUpsert(workspaceId, guid, async () => {
-      const blob = Buffer.from(encodeStateAsUpdate(doc));
+    const blob = Buffer.from(encodeStateAsUpdate(doc));
 
-      if (isEmptyBuffer(blob)) {
-        return false;
+    if (isEmptyBuffer(blob)) {
+      return undefined;
+    }
+
+    const state = Buffer.from(encodeStateVector(doc));
+
+    // CONCERNS:
+    //   i. Because we save the real user's last seen action time as `updatedAt`,
+    //      it's possible to simply compare the `updatedAt` to determine if the snapshot is older than the one we are going to save.
+    //
+    //  ii. Prisma doesn't support `upsert` with additional `where` condition along side unique constraint.
+    //      In our case, we need to manually check the `updatedAt` to avoid overriding the newer snapshot.
+    //      where: { id_workspaceId: {}, updatedAt: { lt: updatedAt } }
+    //                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+    //
+    // iii. Only set the seq number when creating the snapshot.
+    //      For updating scenario, the seq number will be updated when updates pushed to db.
+    try {
+      const result: { updatedAt: Date }[] = await this.db.$queryRaw`
+        INSERT INTO "snapshots" ("workspace_id", "guid", "blob", "state", "seq", "created_at", "updated_at")
+        VALUES (${workspaceId}, ${guid}, ${blob}, ${state}, ${seq}, DEFAULT, ${updatedAt})
+        ON CONFLICT ("workspace_id", "guid")
+        DO UPDATE SET "blob" = ${blob}, "state" = ${state}, "updated_at" = ${updatedAt}, "seq" = ${seq}
+        WHERE "snapshots"."workspace_id" = ${workspaceId} AND "snapshots"."guid" = ${guid} AND "snapshots"."updated_at" <= ${updatedAt}
+        RETURNING "snapshots"."workspace_id" as "workspaceId", "snapshots"."guid" as "id", "snapshots"."updated_at" as "updatedAt"
+      `;
+
+      // const result = await this.db.snapshot.upsert({
+      //   select: {
+      //     updatedAt: true,
+      //     seq: true,
+      //   },
+      //   where: {
+      //     id_workspaceId: {
+      //       workspaceId,
+      //       id: guid,
+      //     },
+      //     ⬇️ NOT SUPPORTED BY PRISMA YET
+      //     updatedAt: {
+      //       lt: updatedAt,
+      //     },
+      //   },
+      //   update: {
+      //     blob,
+      //     state,
+      //     updatedAt,
+      //   },
+      //   create: {
+      //     workspaceId,
+      //     id: guid,
+      //     blob,
+      //     state,
+      //     updatedAt,
+      //     seq,
+      //   },
+      // });
+
+      // if the condition `snapshot.updatedAt > updatedAt` is true, by which means the snapshot has already been updated by other process,
+      // the updates has been applied to current `doc` must have been seen by the other process as well.
+      // The `updatedSnapshot` will be `undefined` in this case.
+      const updatedSnapshot = result.at(0);
+
+      if (!updatedSnapshot) {
+        return undefined;
       }
 
-      const state = Buffer.from(encodeStateVector(doc));
-
-      return await this.db.$transaction(async db => {
-        const snapshot = await db.snapshot.findUnique({
-          where: {
-            id_workspaceId: {
-              id: guid,
-              workspaceId,
-            },
-          },
-        });
-
-        // update
-        if (snapshot) {
-          // only update if state is newer
-          if (isStateNewer(snapshot.state ?? Buffer.from([0]), state)) {
-            await db.snapshot.update({
-              select: {
-                seq: true,
-              },
-              where: {
-                id_workspaceId: {
-                  workspaceId,
-                  id: guid,
-                },
-              },
-              data: {
-                blob,
-                state,
-                updatedAt,
-              },
-            });
-
-            return true;
-          } else {
-            return false;
-          }
-        } else {
-          // create
-          await db.snapshot.create({
-            select: {
-              seq: true,
-            },
-            data: {
-              id: guid,
-              workspaceId,
-              blob,
-              state,
-              seq: initialSeq,
-              createdAt: updatedAt,
-              updatedAt,
-            },
-          });
-
-          return true;
-        }
-      });
-    });
+      return true;
+    } catch (e) {
+      this.logger.error('Failed to upsert snapshot', e);
+      return false;
+    }
   }
 
   private async _get(
@@ -548,7 +531,7 @@ export class DocManager implements OnModuleInit, OnModuleDestroy {
 
     if (updates.length) {
       return {
-        doc: await this.squash(updates, snapshot),
+        doc: await this.squash(snapshot, updates),
       };
     }
 
@@ -559,17 +542,17 @@ 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.
    */
-  private async squash(updates: Update[], snapshot: Snapshot | null) {
+  @CallTimer('doc', 'squash')
+  private async squash(snapshot: Snapshot | null, updates: Update[]) {
     if (!updates.length) {
       throw new Error('No updates to squash');
     }
-    const first = updates[0];
-    const last = updates[updates.length - 1];
 
-    const { id, workspaceId } = first;
+    const last = updates[updates.length - 1];
+    const { id, workspaceId } = last;
 
     const doc = await this.applyUpdates(
-      first.id,
+      id,
       snapshot ? snapshot.blob : Buffer.from([0, 0]),
       ...updates.map(u => u.blob)
     );
@@ -600,19 +583,24 @@ export class DocManager implements OnModuleInit, OnModuleDestroy {
       );
     }
 
-    // always delete updates
-    // the upsert will return false if the state is not newer, so we don't need to worry about it
-    const { count } = await this.db.update.deleteMany({
-      where: {
-        id,
-        workspaceId,
-        seq: {
-          in: updates.map(u => u.seq),
+    // we will keep the updates only if the upsert failed on unknown reason
+    // `done === undefined` means the updates is outdated(have already been merged by other process), safe to be deleted
+    // `done === true` means the upsert is successful, safe to be deleted
+    if (done !== false) {
+      // always delete updates
+      // the upsert will return false if the state is not newer, so we don't need to worry about it
+      const { count } = await this.db.update.deleteMany({
+        where: {
+          id,
+          workspaceId,
+          seq: {
+            in: updates.map(u => u.seq),
+          },
         },
-      },
-    });
+      });
 
-    await this.updateCachedUpdatesCount(workspaceId, id, -count);
+      await this.updateCachedUpdatesCount(workspaceId, id, -count);
+    }
 
     return doc;
   }
@@ -761,18 +749,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

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

@@ -277,6 +277,7 @@ export class WorkspaceResolver {
             id: workspace.id,
             workspaceId: workspace.id,
             blob: buffer,
+            updatedAt: new Date(),
           },
         });
       }

packages/backend/server/src/plugins/payment/resolver.ts

@@ -3,6 +3,7 @@ import {
   Args,
   Context,
   Field,
+  InputType,
   Int,
   Mutation,
   ObjectType,
@@ -125,6 +126,31 @@ class UserInvoiceType implements Partial<UserInvoice> {
   updatedAt!: Date;
 }
 
+@InputType()
+class CreateCheckoutSessionInput {
+  @Field(() => SubscriptionRecurring, {
+    nullable: true,
+    defaultValue: SubscriptionRecurring.Yearly,
+  })
+  recurring!: SubscriptionRecurring;
+
+  @Field(() => SubscriptionPlan, {
+    nullable: true,
+    defaultValue: SubscriptionPlan.Pro,
+  })
+  plan!: SubscriptionPlan;
+
+  @Field(() => String, { nullable: true })
+  coupon!: string | null;
+
+  @Field(() => String, { nullable: true })
+  successCallbackLink!: string | null;
+
+  // @FIXME(forehalo): we should put this field in the header instead of as a explicity args
+  @Field(() => String)
+  idempotencyKey!: string;
+}
+
 @Auth()
 @Resolver(() => UserSubscriptionType)
 export class SubscriptionResolver {
@@ -182,7 +208,11 @@ export class SubscriptionResolver {
     });
   }
 
+  /**
+   * @deprecated
+   */
   @Mutation(() => String, {
+    deprecationReason: 'use `createCheckoutSession` instead',
     description: 'Create a subscription checkout link of stripe',
   })
   async checkout(
@@ -193,6 +223,7 @@ export class SubscriptionResolver {
   ) {
     const session = await this.service.createCheckoutSession({
       user,
+      plan: SubscriptionPlan.Pro,
       recurring,
       redirectUrl: `${this.config.baseUrl}/upgrade-success`,
       idempotencyKey,
@@ -210,6 +241,36 @@ export class SubscriptionResolver {
     return session.url;
   }
 
+  @Mutation(() => String, {
+    description: 'Create a subscription checkout link of stripe',
+  })
+  async createCheckoutSession(
+    @CurrentUser() user: User,
+    @Args({ name: 'input', type: () => CreateCheckoutSessionInput })
+    input: CreateCheckoutSessionInput
+  ) {
+    const session = await this.service.createCheckoutSession({
+      user,
+      plan: input.plan,
+      recurring: input.recurring,
+      promotionCode: input.coupon,
+      redirectUrl:
+        input.successCallbackLink ?? `${this.config.baseUrl}/upgrade-success`,
+      idempotencyKey: input.idempotencyKey,
+    });
+
+    if (!session.url) {
+      throw new GraphQLError('Failed to create checkout session', {
+        extensions: {
+          status: HttpStatus[HttpStatus.BAD_GATEWAY],
+          code: HttpStatus.BAD_GATEWAY,
+        },
+      });
+    }
+
+    return session.url;
+  }
+
   @Mutation(() => String, {
     description: 'Create a stripe customer portal to manage payment methods',
   })

packages/backend/server/src/plugins/payment/service.ts

@@ -69,13 +69,15 @@ export class SubscriptionService {
   async createCheckoutSession({
     user,
     recurring,
+    plan,
+    promotionCode,
     redirectUrl,
     idempotencyKey,
-    plan = SubscriptionPlan.Pro,
   }: {
     user: User;
-    plan?: SubscriptionPlan;
     recurring: SubscriptionRecurring;
+    plan: SubscriptionPlan;
+    promotionCode?: string | null;
     redirectUrl: string;
     idempotencyKey: string;
   }) {
@@ -95,7 +97,28 @@ export class SubscriptionService {
       `${idempotencyKey}-getOrCreateCustomer`,
       user
     );
-    const coupon = await this.getAvailableCoupon(user, CouponType.EarlyAccess);
+
+    let discount: { coupon?: string; promotion_code?: string } | undefined;
+
+    if (promotionCode) {
+      const code = await this.getAvailablePromotionCode(
+        promotionCode,
+        customer.stripeCustomerId
+      );
+      if (code) {
+        discount ??= {};
+        discount.promotion_code = code;
+      }
+    } else {
+      const coupon = await this.getAvailableCoupon(
+        user,
+        CouponType.EarlyAccess
+      );
+      if (coupon) {
+        discount ??= {};
+        discount.coupon = coupon;
+      }
+    }
 
     return await this.stripe.checkout.sessions.create(
       {
@@ -108,13 +131,11 @@ export class SubscriptionService {
         tax_id_collection: {
           enabled: true,
         },
-        ...(coupon
+        ...(discount
           ? {
-              discounts: [{ coupon }],
+              discounts: [discount],
             }
-          : {
-              allow_promotion_codes: true,
-            }),
+          : { allow_promotion_codes: true }),
         mode: 'subscription',
         success_url: redirectUrl,
         customer: customer.stripeCustomerId,
@@ -643,4 +664,33 @@ export class SubscriptionService {
 
     return null;
   }
+
+  private async getAvailablePromotionCode(
+    userFacingPromotionCode: string,
+    customer?: string
+  ) {
+    const list = await this.stripe.promotionCodes.list({
+      code: userFacingPromotionCode,
+      active: true,
+      limit: 1,
+    });
+
+    const code = list.data[0];
+    if (!code) {
+      return null;
+    }
+
+    let available = false;
+
+    if (code.customer) {
+      available =
+        typeof code.customer === 'string'
+          ? code.customer === customer
+          : code.customer.id === customer;
+    } else {
+      available = true;
+    }
+
+    return available ? code.id : null;
+  }
 }

packages/backend/server/src/schema.gql

@@ -2,6 +2,14 @@
 # THIS FILE WAS AUTOMATICALLY GENERATED (DO NOT MODIFY)
 # ------------------------------------------------------
 
+input CreateCheckoutSessionInput {
+  coupon: String
+  idempotencyKey: String!
+  plan: SubscriptionPlan = Pro
+  recurring: SubscriptionRecurring = Yearly
+  successCallbackLink: String
+}
+
 """
 A date-time string at UTC, such as 2019-12-03T09:54:33Z, compliant with the date-time format.
 """
@@ -107,7 +115,10 @@ type Mutation {
   changePassword(newPassword: String!, token: String!): UserType!
 
   """Create a subscription checkout link of stripe"""
-  checkout(idempotencyKey: String!, recurring: SubscriptionRecurring!): String!
+  checkout(idempotencyKey: String!, recurring: SubscriptionRecurring!): String! @deprecated(reason: "use `createCheckoutSession` instead")
+
+  """Create a subscription checkout link of stripe"""
+  createCheckoutSession(input: CreateCheckoutSessionInput!): String!
 
   """Create a stripe customer portal to manage payment methods"""
   createCustomerPortal: String!

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

@@ -4,12 +4,7 @@ import { TestingModule } from '@nestjs/testing';
 import { PrismaClient } from '@prisma/client';
 import test from 'ava';
 import * as Sinon from 'sinon';
-import {
-  applyUpdate,
-  decodeStateVector,
-  Doc as YDoc,
-  encodeStateAsUpdate,
-} from 'yjs';
+import { applyUpdate, Doc as YDoc, encodeStateAsUpdate } from 'yjs';
 
 import { DocManager, DocModule } from '../src/core/doc';
 import { QuotaModule } from '../src/core/quota';
@@ -277,72 +272,120 @@ test('should throw if meet max retry times', async t => {
   t.is(stub.callCount, 5);
 });
 
-test('should not update snapshot if state is outdated', async t => {
-  const db = m.get(PrismaClient);
+test('should be able to insert the snapshot if it is new created', async t => {
   const manager = m.get(DocManager);
 
-  await db.snapshot.create({
-    data: {
-      id: '2',
-      workspaceId: '2',
-      blob: Buffer.from([0, 0]),
-      seq: 1,
-    },
-  });
   const doc = new YDoc();
   const text = doc.getText('content');
-  const updates: Buffer[] = [];
-
-  doc.on('update', update => {
-    updates.push(Buffer.from(update));
-  });
-
   text.insert(0, 'hello');
-  text.insert(5, 'world');
-  text.insert(5, ' ');
+  const update = encodeStateAsUpdate(doc);
 
-  await Promise.all(updates.map(update => manager.push('2', '2', update)));
+  await manager.push('1', '1', Buffer.from(update));
 
-  const updateWith3Records = await manager.getUpdates('2', '2');
-  text.insert(11, '!');
-  await manager.push('2', '2', updates[3]);
-  const updateWith4Records = await manager.getUpdates('2', '2');
-
-  // Simulation:
-  //   Node A get 3 updates and squash them at time 1, will finish at time 10
-  //   Node B get 4 updates and squash them at time 3, will finish at time 8
-  //   Node B finish the squash first, and update the snapshot
-  //   Node A finish the squash later, and update the snapshot to an outdated state
-  // Time: ---------------------->
-  //    A:   ^get           ^upsert
-  //    B:      ^get    ^upsert
-  //
-  // We should avoid such situation
+  const updates = await manager.getUpdates('1', '1');
+  t.is(updates.length, 1);
   // @ts-expect-error private
-  await manager.squash(updateWith4Records, null);
-  // @ts-expect-error private
-  await manager.squash(updateWith3Records, null);
+  const snapshot = await manager.squash(null, updates);
 
-  const result = await db.snapshot.findUnique({
+  t.truthy(snapshot);
+  t.is(snapshot.getText('content').toString(), 'hello');
+
+  const restUpdates = await manager.getUpdates('1', '1');
+
+  t.is(restUpdates.length, 0);
+});
+
+test('should be able to merge updates into snapshot', async t => {
+  const manager = m.get(DocManager);
+
+  const updates: Buffer[] = [];
+  {
+    const doc = new YDoc();
+    doc.on('update', data => {
+      updates.push(Buffer.from(data));
+    });
+
+    const text = doc.getText('content');
+    text.insert(0, 'hello');
+    text.insert(5, 'world');
+    text.insert(5, ' ');
+    text.insert(11, '!');
+  }
+
+  {
+    await manager.batchPush('1', '1', updates.slice(0, 2));
+    // do the merge
+    const doc = (await manager.get('1', '1'))!;
+
+    t.is(doc.getText('content').toString(), 'helloworld');
+  }
+
+  {
+    await manager.batchPush('1', '1', updates.slice(2));
+    const doc = (await manager.get('1', '1'))!;
+
+    t.is(doc.getText('content').toString(), 'hello world!');
+  }
+
+  const restUpdates = await manager.getUpdates('1', '1');
+
+  t.is(restUpdates.length, 0);
+});
+
+test('should not update snapshot if doc is outdated', async t => {
+  const manager = m.get(DocManager);
+  const db = m.get(PrismaClient);
+
+  const updates: Buffer[] = [];
+  {
+    const doc = new YDoc();
+    doc.on('update', data => {
+      updates.push(Buffer.from(data));
+    });
+
+    const text = doc.getText('content');
+    text.insert(0, 'hello');
+    text.insert(5, 'world');
+    text.insert(5, ' ');
+    text.insert(11, '!');
+  }
+
+  await manager.batchPush('2', '1', updates.slice(0, 2)); // 'helloworld'
+  // merge updates into snapshot
+  await manager.get('2', '1');
+  // fake the snapshot is a lot newer
+  await db.snapshot.update({
     where: {
       id_workspaceId: {
-        id: '2',
         workspaceId: '2',
+        id: '1',
       },
     },
+    data: {
+      updatedAt: new Date(Date.now() + 10000),
+    },
   });
 
-  if (!result) {
-    t.fail('snapshot not found');
-    return;
+  {
+    const snapshot = await manager.getSnapshot('2', '1');
+    await manager.batchPush('2', '1', updates.slice(2)); // 'hello world!'
+    const updateRecords = await manager.getUpdates('2', '1');
+
+    // @ts-expect-error private
+    const doc = await manager.squash(snapshot, updateRecords);
+
+    // all updated will merged into doc not matter it's timestamp is outdated or not,
+    // but the snapshot record will not be updated
+    t.is(doc.getText('content').toString(), 'hello world!');
   }
 
-  const state = decodeStateVector(result.state!);
-  t.is(state.get(doc.clientID), 12);
+  {
+    const doc = new YDoc();
+    applyUpdate(doc, (await manager.getSnapshot('2', '1'))!.blob);
+    // the snapshot will not get touched if the new doc's timestamp is outdated
+    t.is(doc.getText('content').toString(), 'helloworld');
 
-  const d = new YDoc();
-  applyUpdate(d, result.blob!);
-
-  const dtext = d.getText('content');
-  t.is(dtext.toString(), 'hello world!');
+    // the updates are known as outdated, so they will be deleted
+    t.is((await manager.getUpdates('2', '1')).length, 0);
+  }
 });

packages/frontend/core/src/components/affine/auth/sign-in.tsx

@@ -20,6 +20,7 @@ import type { AuthPanelProps } from './index';
 import * as style from './style.css';
 import { INTERNAL_BETA_URL, useAuth } from './use-auth';
 import { Captcha, useCaptcha } from './use-captcha';
+import { useSubscriptionSearch } from './use-subscription';
 
 function validateEmail(email: string) {
   return emailRegex.test(email);
@@ -34,6 +35,7 @@ export const SignIn: FC<AuthPanelProps> = ({
   const t = useAFFiNEI18N();
   const loginStatus = useCurrentLoginStatus();
   const [verifyToken, challenge] = useCaptcha();
+  const subscriptionData = useSubscriptionSearch();
 
   const {
     isMutating: isSigningIn,
@@ -81,7 +83,8 @@ export const SignIn: FC<AuthPanelProps> = ({
     if (verifyToken) {
       if (user) {
         // provider password sign-in if user has by default
-        if (user.hasPassword) {
+        //  If with payment, onl support email sign in to avoid redirect to affine app
+        if (user.hasPassword && !subscriptionData) {
           setAuthState('signInWithPassword');
         } else {
           const res = await signIn(email, verifyToken, challenge);
@@ -101,6 +104,7 @@ export const SignIn: FC<AuthPanelProps> = ({
       }
     }
   }, [
+    subscriptionData,
     challenge,
     email,
     setAuthEmail,

packages/frontend/core/src/components/affine/auth/subscription-redirect.tsx

@@ -3,10 +3,10 @@ import { Button } from '@affine/component/ui/button';
 import { Loading } from '@affine/component/ui/loading';
 import { AffineShapeIcon } from '@affine/core/components/page-list';
 import { useAsyncCallback } from '@affine/core/hooks/affine-async-hooks';
-import type { SubscriptionRecurring } from '@affine/graphql';
+import type { SubscriptionPlan, SubscriptionRecurring } from '@affine/graphql';
 import {
   changePasswordMutation,
-  checkoutMutation,
+  createCheckoutSessionMutation,
   subscriptionQuery,
 } from '@affine/graphql';
 import { useAFFiNEI18N } from '@affine/i18n/hooks';
@@ -30,18 +30,25 @@ const usePaymentRedirect = () => {
   }
 
   const recurring = searchData.recurring as SubscriptionRecurring;
+  const plan = searchData.plan as SubscriptionPlan;
+  const coupon = searchData.coupon;
   const idempotencyKey = useMemo(() => nanoid(), []);
   const { trigger: checkoutSubscription } = useMutation({
-    mutation: checkoutMutation,
+    mutation: createCheckoutSessionMutation,
   });
 
   return useAsyncCallback(async () => {
-    const { checkout } = await checkoutSubscription({
-      recurring,
-      idempotencyKey,
+    const { createCheckoutSession: checkoutUrl } = await checkoutSubscription({
+      input: {
+        recurring,
+        plan,
+        coupon,
+        idempotencyKey,
+        successCallbackLink: null,
+      },
     });
-    window.open(checkout, '_self', 'norefferer');
-  }, [recurring, idempotencyKey, checkoutSubscription]);
+    window.open(checkoutUrl, '_self', 'norefferer');
+  }, [recurring, plan, coupon, idempotencyKey, checkoutSubscription]);
 };
 
 const CenterLoading = () => {

packages/frontend/core/src/components/affine/auth/use-subscription.ts

@@ -4,6 +4,7 @@ import { useSearchParams } from 'react-router-dom';
 enum SubscriptionKey {
   Recurring = 'subscription_recurring',
   Plan = 'subscription_plan',
+  Coupon = 'coupon',
   SignUp = 'sign_up', // A new user with subscription journey: signup > set password > pay in stripe > go to app
   Token = 'token', // When signup, there should have a token to set password
 }
@@ -22,11 +23,13 @@ export function useSubscriptionSearch() {
 
     const recurring = searchParams.get(SubscriptionKey.Recurring);
     const plan = searchParams.get(SubscriptionKey.Plan);
+    const coupon = searchParams.get(SubscriptionKey.Coupon);
     const withSignUp = searchParams.get(SubscriptionKey.SignUp) === '1';
     const passwordToken = searchParams.get(SubscriptionKey.Token);
     return {
       recurring,
       plan,
+      coupon,
       withSignUp,
       passwordToken,
       getRedirectUrl(signUp?: boolean) {
@@ -35,6 +38,10 @@ export function useSubscriptionSearch() {
           [SubscriptionKey.Plan, plan ?? ''],
         ]);
 
+        if (coupon) {
+          paymentParams.set(SubscriptionKey.Coupon, coupon);
+        }
+
         if (signUp) {
           paymentParams.set(SubscriptionKey.SignUp, '1');
         }

packages/frontend/core/src/components/affine/setting-modal/general-setting/plans/plan-card.tsx

@@ -6,7 +6,7 @@ import type {
   SubscriptionMutator,
 } from '@affine/core/hooks/use-subscription';
 import {
-  checkoutMutation,
+  createCheckoutSessionMutation,
   SubscriptionPlan,
   SubscriptionRecurring,
   SubscriptionStatus,
@@ -359,7 +359,7 @@ const Upgrade = ({
 }) => {
   const t = useAFFiNEI18N();
   const { isMutating, trigger } = useMutation({
-    mutation: checkoutMutation,
+    mutation: createCheckoutSessionMutation,
   });
 
   const newTabRef = useRef<Window | null>(null);
@@ -383,13 +383,21 @@ const Upgrade = ({
       newTabRef.current.focus();
     } else {
       await trigger(
-        { recurring, idempotencyKey },
+        {
+          input: {
+            recurring,
+            idempotencyKey,
+            plan: SubscriptionPlan.Pro, // Only support prod plan now.
+            coupon: null,
+            successCallbackLink: null,
+          },
+        },
         {
           onSuccess: data => {
             // FIXME: safari prevents from opening new tab by window api
             // TODO(@xp): what if electron?
             const newTab = window.open(
-              data.checkout,
+              data.createCheckoutSession,
               '_blank',
               'noopener noreferrer'
             );

packages/frontend/core/src/hooks/__tests__/use-block-suite-workspace-helper.spec.ts

@@ -7,7 +7,7 @@ import { __unstableSchemas, AffineSchemas } from '@blocksuite/blocks/models';
 import { Schema, Workspace } from '@blocksuite/store';
 import { renderHook } from '@testing-library/react';
 import { initEmptyPage } from '@toeverything/infra/blocksuite';
-import { beforeEach, describe, expect, test } from 'vitest';
+import { beforeEach, describe, expect, test, vi } from 'vitest';
 
 import { useBlockSuitePageMeta } from '../use-block-suite-page-meta';
 import { useBlockSuiteWorkspaceHelper } from '../use-block-suite-workspace-helper';
@@ -17,18 +17,26 @@ let blockSuiteWorkspace: Workspace;
 const schema = new Schema();
 schema.register(AffineSchemas).register(__unstableSchemas);
 
+// todo: this module has some side-effects that will break the tests
+vi.mock('@affine/workspace-impl', () => ({
+  default: {},
+}));
+
 beforeEach(async () => {
   blockSuiteWorkspace = new Workspace({
     id: 'test',
     schema,
   });
+
+  blockSuiteWorkspace.doc.emit('sync', []);
+
   await initEmptyPage(blockSuiteWorkspace.createPage({ id: 'page0' }));
   await initEmptyPage(blockSuiteWorkspace.createPage({ id: 'page1' }));
   await initEmptyPage(blockSuiteWorkspace.createPage({ id: 'page2' }));
 });
 
 describe('useBlockSuiteWorkspaceHelper', () => {
-  test('should create page', () => {
+  test('should create page', async () => {
     expect(blockSuiteWorkspace.meta.pageMetas.length).toBe(3);
     const helperHook = renderHook(() =>
       useBlockSuiteWorkspaceHelper(blockSuiteWorkspace)
@@ -36,6 +44,7 @@ describe('useBlockSuiteWorkspaceHelper', () => {
     const pageMetaHook = renderHook(() =>
       useBlockSuitePageMeta(blockSuiteWorkspace)
     );
+    await new Promise(resolve => setTimeout(resolve));
     expect(pageMetaHook.result.current.length).toBe(3);
     expect(blockSuiteWorkspace.meta.pageMetas.length).toBe(3);
     const page = helperHook.result.current.createPage('page4');

packages/frontend/core/src/hooks/use-block-suite-page-meta.ts

@@ -5,9 +5,11 @@ import type { Atom } from 'jotai';
 import { atom, useAtomValue } from 'jotai';
 import { useMemo } from 'react';
 
+import { useJournalHelper } from './use-journal';
+
 const weakMap = new WeakMap<Workspace, Atom<PageMeta[]>>();
 
-export function useBlockSuitePageMeta(
+export function useAllBlockSuitePageMeta(
   blockSuiteWorkspace: Workspace
 ): PageMeta[] {
   if (!weakMap.has(blockSuiteWorkspace)) {
@@ -26,6 +28,18 @@ export function useBlockSuitePageMeta(
   return useAtomValue(weakMap.get(blockSuiteWorkspace) as Atom<PageMeta[]>);
 }
 
+export function useBlockSuitePageMeta(blocksuiteWorkspace: Workspace) {
+  const pageMetas = useAllBlockSuitePageMeta(blocksuiteWorkspace);
+  const { isPageJournal } = useJournalHelper(blocksuiteWorkspace);
+  return useMemo(
+    () =>
+      pageMetas.filter(
+        pageMeta => !isPageJournal(pageMeta.id) || !!pageMeta.updatedDate
+      ),
+    [isPageJournal, pageMetas]
+  );
+}
+
 export function usePageMetaHelper(blockSuiteWorkspace: Workspace) {
   return useMemo(
     () => ({

packages/frontend/core/src/hooks/use-journal.ts

@@ -33,6 +33,10 @@ export const useJournalHelper = (workspace: BlockSuiteWorkspace) => {
     (maybeDate: MaybeDate) => {
       const title = dayjs(maybeDate).format(JOURNAL_DATE_FORMAT);
       const page = bsWorkspaceHelper.createPage();
+      // set created date to match the journal date
+      page.workspace.setPageMeta(page.id, {
+        createDate: dayjs(maybeDate).toDate().getTime(),
+      });
       initEmptyPage(page, title).catch(err =>
         console.error('Failed to load journal page', err)
       );

packages/frontend/core/src/pages/workspace/detail-page/editor-sidebar/extensions/journal.tsx

@@ -6,6 +6,7 @@ import {
 } from '@affine/component';
 import { MoveToTrash } from '@affine/core/components/page-list';
 import { useTrashModalHelper } from '@affine/core/hooks/affine/use-trash-modal-helper';
+import { useBlockSuitePageMeta } from '@affine/core/hooks/use-block-suite-page-meta';
 import { useBlockSuiteWorkspacePageTitle } from '@affine/core/hooks/use-block-suite-workspace-page-title';
 import {
   useJournalHelper,
@@ -13,6 +14,7 @@ import {
   useJournalRouteHelper,
 } from '@affine/core/hooks/use-journal';
 import { useNavigateHelper } from '@affine/core/hooks/use-navigate-helper';
+import type { BlockSuiteWorkspace } from '@affine/core/shared';
 import { useAFFiNEI18N } from '@affine/i18n/hooks';
 import {
   EdgelessIcon,
@@ -20,7 +22,7 @@ import {
   PageIcon,
   TodayIcon,
 } from '@blocksuite/icons';
-import type { Page } from '@blocksuite/store';
+import type { Page, PageMeta } from '@blocksuite/store';
 import { assignInlineVars } from '@vanilla-extract/dynamic';
 import clsx from 'clsx';
 import dayjs from 'dayjs';
@@ -41,21 +43,28 @@ const CountDisplay = ({
   return <span {...attrs}>{count > max ? `${max}+` : count}</span>;
 };
 interface PageItemProps extends HTMLAttributes<HTMLDivElement> {
-  page: Page;
+  pageMeta: PageMeta;
+  workspace: BlockSuiteWorkspace;
   right?: ReactNode;
 }
-const PageItem = ({ page, right, className, ...attrs }: PageItemProps) => {
-  const { isJournal } = useJournalInfoHelper(page.workspace, page.id);
-  const title = useBlockSuiteWorkspacePageTitle(page.workspace, page.id);
+const PageItem = ({
+  pageMeta,
+  workspace,
+  right,
+  className,
+  ...attrs
+}: PageItemProps) => {
+  const { isJournal } = useJournalInfoHelper(workspace, pageMeta.id);
+  const title = useBlockSuiteWorkspacePageTitle(workspace, pageMeta.id);
 
   const Icon = isJournal
     ? TodayIcon
-    : page.meta.mode === 'edgeless'
+    : pageMeta.mode === 'edgeless'
       ? EdgelessIcon
       : PageIcon;
   return (
     <div
-      aria-label={page.meta.title}
+      aria-label={pageMeta.title}
       className={clsx(className, styles.pageItem)}
       {...attrs}
     >
@@ -114,15 +123,12 @@ const EditorJournalPanel = (props: EditorExtensionProps) => {
 };
 
 const sortPagesByDate = (
-  pages: Page[],
+  pages: PageMeta[],
   field: 'updatedDate' | 'createDate',
   order: 'asc' | 'desc' = 'desc'
 ) => {
   return [...pages].sort((a, b) => {
-    return (
-      (order === 'asc' ? 1 : -1) *
-      dayjs(b.meta[field]).diff(dayjs(a.meta[field]))
-    );
+    return (order === 'asc' ? 1 : -1) * dayjs(b[field]).diff(dayjs(a[field]));
   });
 };
 
@@ -141,21 +147,21 @@ const JournalDailyCountBlock = ({ workspace, date }: JournalBlockProps) => {
   const nodeRef = useRef<HTMLDivElement>(null);
   const t = useAFFiNEI18N();
   const [activeItem, setActiveItem] = useState<NavItemName>('createdToday');
+  const pageMetas = useBlockSuitePageMeta(workspace);
 
   const navigateHelper = useNavigateHelper();
 
   const getTodaysPages = useCallback(
     (field: 'createDate' | 'updatedDate') => {
-      const pages: Page[] = [];
-      Array.from(workspace.pages.values()).forEach(page => {
-        if (page.meta.trash) return;
-        if (page.meta[field] && dayjs(page.meta[field]).isSame(date, 'day')) {
-          pages.push(page);
-        }
-      });
-      return sortPagesByDate(pages, field);
+      return sortPagesByDate(
+        pageMetas.filter(pageMeta => {
+          if (pageMeta.trash) return false;
+          return pageMeta[field] && dayjs(pageMeta[field]).isSame(date, 'day');
+        }),
+        field
+      );
     },
-    [date, workspace.pages]
+    [date, pageMetas]
   );
 
   const createdToday = useMemo(
@@ -224,14 +230,15 @@ const JournalDailyCountBlock = ({ workspace, date }: JournalBlockProps) => {
               <Scrollable.Scrollbar />
               <Scrollable.Viewport>
                 <div className={styles.dailyCountContent} ref={nodeRef}>
-                  {renderList.map((page, index) => (
+                  {renderList.map((pageMeta, index) => (
                     <PageItem
                       onClick={() =>
-                        navigateHelper.openPage(workspace.id, page.id)
+                        navigateHelper.openPage(workspace.id, pageMeta.id)
                       }
                       tabIndex={name === activeItem ? 0 : -1}
                       key={index}
-                      page={page}
+                      pageMeta={pageMeta}
+                      workspace={workspace}
                     />
                   ))}
                 </div>
@@ -282,7 +289,8 @@ const ConflictList = ({
           <PageItem
             aria-label={page.meta.title}
             aria-selected={isCurrent}
-            page={page}
+            pageMeta={page.meta}
+            workspace={workspace}
             key={page.id}
             right={
               <Menu

packages/frontend/core/src/pages/workspace/index.tsx

@@ -8,7 +8,13 @@ import {
 } from '@affine/core/modules/workspace';
 import { type Workspace } from '@affine/workspace';
 import { useAtom, useAtomValue } from 'jotai';
-import { type ReactElement, Suspense, useEffect, useMemo } from 'react';
+import {
+  type ReactElement,
+  Suspense,
+  useEffect,
+  useMemo,
+  useState,
+} from 'react';
 import { Outlet, useParams } from 'react-router-dom';
 
 import { AffineErrorBoundary } from '../../components/affine/affine-error-boundary';
@@ -67,12 +73,34 @@ export const Component = (): ReactElement => {
     localStorage.setItem('last_workspace_id', workspace.id);
   }, [setCurrentWorkspace, meta, workspaceManager, workspace]);
 
+  const [workspaceIsLoading, setWorkspaceIsLoading] = useState(true);
+
+  // hotfix: avoid doing operation, before workspace is loaded
+  useEffect(() => {
+    if (!workspace) {
+      setWorkspaceIsLoading(true);
+      return;
+    }
+    const metaYMap = workspace.blockSuiteWorkspace.doc.getMap('meta');
+
+    const handleYMapChanged = () => {
+      setWorkspaceIsLoading(metaYMap.size === 0);
+    };
+
+    handleYMapChanged();
+
+    metaYMap.observe(handleYMapChanged);
+    return () => {
+      metaYMap.unobserve(handleYMapChanged);
+    };
+  }, [workspace]);
+
   // if listLoading is false, we can show 404 page, otherwise we should show loading page.
   if (listLoading === false && meta === undefined) {
     return <PageNotFound />;
   }
 
-  if (!workspace) {
+  if (!workspace || workspaceIsLoading) {
     return <WorkspaceFallback key="workspaceLoading" />;
   }
 

packages/frontend/graphql/src/graphql/create-checkout-link.gql

@@ -1,6 +1,3 @@
-mutation checkout(
-  $recurring: SubscriptionRecurring!
-  $idempotencyKey: String!
-) {
-  checkout(recurring: $recurring, idempotencyKey: $idempotencyKey)
+mutation createCheckoutSession($input: CreateCheckoutSessionInput!) {
+  createCheckoutSession(input: $input)
 }

packages/frontend/graphql/src/graphql/index.ts

@@ -127,14 +127,14 @@ mutation changePassword($token: String!, $newPassword: String!) {
 }`,
 };
 
-export const checkoutMutation = {
-  id: 'checkoutMutation' as const,
-  operationName: 'checkout',
-  definitionName: 'checkout',
+export const createCheckoutSessionMutation = {
+  id: 'createCheckoutSessionMutation' as const,
+  operationName: 'createCheckoutSession',
+  definitionName: 'createCheckoutSession',
   containsFile: false,
   query: `
-mutation checkout($recurring: SubscriptionRecurring!, $idempotencyKey: String!) {
-  checkout(recurring: $recurring, idempotencyKey: $idempotencyKey)
+mutation createCheckoutSession($input: CreateCheckoutSessionInput!) {
+  createCheckoutSession(input: $input)
 }`,
 };
 

packages/frontend/graphql/src/schema.ts

@@ -192,12 +192,14 @@ export type ChangePasswordMutation = {
   };
 };
 
-export type CheckoutMutationVariables = Exact<{
-  recurring: SubscriptionRecurring;
-  idempotencyKey: Scalars['String']['input'];
+export type CreateCheckoutSessionMutationVariables = Exact<{
+  input: CreateCheckoutSessionInput;
 }>;
 
-export type CheckoutMutation = { __typename?: 'Mutation'; checkout: string };
+export type CreateCheckoutSessionMutation = {
+  __typename?: 'Mutation';
+  createCheckoutSession: string;
+};
 
 export type CreateCustomerPortalMutationVariables = Exact<{
   [key: string]: never;
@@ -1041,9 +1043,9 @@ export type Mutations =
       response: ChangePasswordMutation;
     }
   | {
-      name: 'checkoutMutation';
-      variables: CheckoutMutationVariables;
-      response: CheckoutMutation;
+      name: 'createCheckoutSessionMutation';
+      variables: CreateCheckoutSessionMutationVariables;
+      response: CreateCheckoutSessionMutation;
     }
   | {
       name: 'createCustomerPortalMutation';

packages/frontend/workspace-impl/package.json

@@ -23,6 +23,7 @@
     "lodash-es": "^4.17.21",
     "nanoid": "^5.0.3",
     "next-auth": "^4.24.5",
+    "rxjs": "^7.8.1",
     "socket.io-client": "^4.7.2",
     "y-protocols": "^1.0.6",
     "yjs": "^13.6.10"

packages/frontend/workspace-impl/src/cloud/awareness.ts

@@ -18,26 +18,70 @@ export function createCloudAwarenessProvider(
   workspaceId: string,
   awareness: Awareness
 ): AwarenessProvider {
-  const socket = getIoManager().socket('/');
+  return new AffineCloudAwarenessProvider(workspaceId, awareness);
+}
 
-  const awarenessBroadcast = ({
+class AffineCloudAwarenessProvider implements AwarenessProvider {
+  socket = getIoManager().socket('/');
+
+  constructor(
+    private readonly workspaceId: string,
+    private readonly awareness: Awareness
+  ) {}
+
+  connect(): void {
+    this.socket.on('server-awareness-broadcast', this.awarenessBroadcast);
+    this.socket.on(
+      'new-client-awareness-init',
+      this.newClientAwarenessInitHandler
+    );
+    this.awareness.on('update', this.awarenessUpdate);
+
+    window.addEventListener('beforeunload', this.windowBeforeUnloadHandler);
+
+    this.socket.on('connect', () => this.handleConnect());
+
+    if (this.socket.connected) {
+      this.handleConnect();
+    } else {
+      this.socket.connect();
+    }
+  }
+  disconnect(): void {
+    removeAwarenessStates(
+      this.awareness,
+      [this.awareness.clientID],
+      'disconnect'
+    );
+    this.awareness.off('update', this.awarenessUpdate);
+    this.socket.emit('client-leave-awareness', this.workspaceId);
+    this.socket.off('server-awareness-broadcast', this.awarenessBroadcast);
+    this.socket.off(
+      'new-client-awareness-init',
+      this.newClientAwarenessInitHandler
+    );
+    this.socket.off('connect', this.handleConnect);
+    window.removeEventListener('unload', this.windowBeforeUnloadHandler);
+  }
+
+  awarenessBroadcast = ({
     workspaceId: wsId,
     awarenessUpdate,
   }: {
     workspaceId: string;
     awarenessUpdate: string;
   }) => {
-    if (wsId !== workspaceId) {
+    if (wsId !== this.workspaceId) {
       return;
     }
     applyAwarenessUpdate(
-      awareness,
+      this.awareness,
       base64ToUint8Array(awarenessUpdate),
       'remote'
     );
   };
 
-  const awarenessUpdate = (changes: AwarenessChanges, origin: unknown) => {
+  awarenessUpdate = (changes: AwarenessChanges, origin: unknown) => {
     if (origin === 'remote') {
       return;
     }
@@ -46,63 +90,41 @@ export function createCloudAwarenessProvider(
       res.concat(cur)
     );
 
-    const update = encodeAwarenessUpdate(awareness, changedClients);
+    const update = encodeAwarenessUpdate(this.awareness, changedClients);
     uint8ArrayToBase64(update)
       .then(encodedUpdate => {
-        socket.emit('awareness-update', {
-          workspaceId: workspaceId,
+        this.socket.emit('awareness-update', {
+          workspaceId: this.workspaceId,
           awarenessUpdate: encodedUpdate,
         });
       })
       .catch(err => logger.error(err));
   };
 
-  const newClientAwarenessInitHandler = () => {
-    const awarenessUpdate = encodeAwarenessUpdate(awareness, [
-      awareness.clientID,
+  newClientAwarenessInitHandler = () => {
+    const awarenessUpdate = encodeAwarenessUpdate(this.awareness, [
+      this.awareness.clientID,
     ]);
     uint8ArrayToBase64(awarenessUpdate)
       .then(encodedAwarenessUpdate => {
-        socket.emit('awareness-update', {
-          guid: workspaceId,
+        this.socket.emit('awareness-update', {
+          guid: this.workspaceId,
           awarenessUpdate: encodedAwarenessUpdate,
         });
       })
       .catch(err => logger.error(err));
   };
 
-  const windowBeforeUnloadHandler = () => {
-    removeAwarenessStates(awareness, [awareness.clientID], 'window unload');
+  windowBeforeUnloadHandler = () => {
+    removeAwarenessStates(
+      this.awareness,
+      [this.awareness.clientID],
+      'window unload'
+    );
   };
 
-  function handleConnect() {
-    socket.emit('client-handshake-awareness', workspaceId);
-    socket.emit('awareness-init', workspaceId);
-  }
-
-  return {
-    connect: () => {
-      socket.on('server-awareness-broadcast', awarenessBroadcast);
-      socket.on('new-client-awareness-init', newClientAwarenessInitHandler);
-      awareness.on('update', awarenessUpdate);
-
-      window.addEventListener('beforeunload', windowBeforeUnloadHandler);
-
-      socket.connect();
-
-      socket.on('connect', handleConnect);
-
-      socket.emit('client-handshake-awareness', workspaceId);
-      socket.emit('awareness-init', workspaceId);
-    },
-    disconnect: () => {
-      removeAwarenessStates(awareness, [awareness.clientID], 'disconnect');
-      awareness.off('update', awarenessUpdate);
-      socket.emit('client-leave-awareness', workspaceId);
-      socket.off('server-awareness-broadcast', awarenessBroadcast);
-      socket.off('new-client-awareness-init', newClientAwarenessInitHandler);
-      socket.off('connect', handleConnect);
-      window.removeEventListener('unload', windowBeforeUnloadHandler);
-    },
+  handleConnect = () => {
+    this.socket.emit('client-handshake-awareness', this.workspaceId);
+    this.socket.emit('awareness-init', this.workspaceId);
   };
 }

packages/frontend/workspace-impl/src/cloud/sync/batch-sync-sender.ts

@@ -1,107 +0,0 @@
-interface SyncUpdateSender {
-  (
-    guid: string,
-    updates: Uint8Array[]
-  ): Promise<{
-    accepted: boolean;
-    retry: boolean;
-  }>;
-}
-
-/**
- * BatchSyncSender is simple wrapper with vanilla update sync with several advanced features:
- *   - ACK mechanism, send updates sequentially with previous sync request correctly responds with ACK
- *   - batching updates, when waiting for previous ACK, new updates will be buffered and sent in single sync request
- *   - retryable, allow retry when previous sync request failed but with retry flag been set to true
- */
-export class BatchSyncSender {
-  private readonly buffered: Uint8Array[] = [];
-  private job: Promise<void> | null = null;
-  private started = true;
-
-  constructor(
-    private readonly guid: string,
-    private readonly rawSender: SyncUpdateSender
-  ) {}
-
-  send(update: Uint8Array) {
-    this.buffered.push(update);
-    this.next();
-    return Promise.resolve();
-  }
-
-  stop() {
-    this.started = false;
-  }
-
-  start() {
-    this.started = true;
-    this.next();
-  }
-
-  private next() {
-    if (!this.started || this.job || !this.buffered.length) {
-      return;
-    }
-
-    const lastIndex = Math.min(
-      this.buffered.length - 1,
-      99 /* max batch updates size */
-    );
-    const updates = this.buffered.slice(0, lastIndex + 1);
-
-    if (updates.length) {
-      this.job = this.rawSender(this.guid, updates)
-        .then(({ accepted, retry }) => {
-          // remove pending updates if updates are accepted
-          if (accepted) {
-            this.buffered.splice(0, lastIndex + 1);
-          }
-
-          // stop when previous sending failed and non-recoverable
-          if (accepted || retry) {
-            // avoid call stack overflow
-            setTimeout(() => {
-              this.next();
-            }, 0);
-          } else {
-            this.stop();
-          }
-        })
-        .catch(() => {
-          this.stop();
-        })
-        .finally(() => {
-          this.job = null;
-        });
-    }
-  }
-}
-
-export class MultipleBatchSyncSender {
-  private senders: Record<string, BatchSyncSender> = {};
-
-  constructor(private readonly rawSender: SyncUpdateSender) {}
-
-  async send(guid: string, update: Uint8Array) {
-    return this.getSender(guid).send(update);
-  }
-
-  private getSender(guid: string) {
-    let sender = this.senders[guid];
-    if (!sender) {
-      sender = new BatchSyncSender(guid, this.rawSender);
-      this.senders[guid] = sender;
-    }
-
-    return sender;
-  }
-
-  start() {
-    Object.values(this.senders).forEach(sender => sender.start());
-  }
-
-  stop() {
-    Object.values(this.senders).forEach(sender => sender.stop());
-  }
-}

packages/frontend/workspace-impl/src/cloud/sync/index.ts

@@ -1,168 +1,199 @@
 import { DebugLogger } from '@affine/debug';
 import { fetchWithTraceReport } from '@affine/graphql';
-import type { SyncStorage } from '@affine/workspace';
+import { type SyncStorage } from '@affine/workspace';
+import {
+  BehaviorSubject,
+  distinctUntilChanged,
+  filter,
+  lastValueFrom,
+  take,
+  timeout,
+} from 'rxjs';
 
 import { getIoManager } from '../../utils/affine-io';
 import { base64ToUint8Array, uint8ArrayToBase64 } from '../../utils/base64';
-import { MultipleBatchSyncSender } from './batch-sync-sender';
 
 const logger = new DebugLogger('affine:storage:socketio');
 
+export class AffineSyncStorage implements SyncStorage {
+  name = 'affine-cloud';
+
+  SEND_TIMEOUT = 30000;
+
+  socket = getIoManager().socket('/');
+
+  connected = new BehaviorSubject(false);
+  handshook = new BehaviorSubject(false);
+
+  constructor(private readonly workspaceId: string) {
+    const handleConnect = () => {
+      this.connected.next(true);
+    };
+    this.socket.on('connect', handleConnect);
+    this.connected.next(this.socket.connected);
+
+    const handleDisconnect = () => {
+      this.connected.next(false);
+      this.handshook.next(false);
+    };
+    this.socket.on('disconnect', handleDisconnect);
+
+    this.connected.pipe(distinctUntilChanged()).subscribe(connected => {
+      if (connected) {
+        this.socket
+          .timeout(this.SEND_TIMEOUT)
+          .emitWithAck('client-handshake-sync', this.workspaceId)
+          .then(() => {
+            this.handshook.next(true);
+          })
+          .catch(err => {
+            logger.error('client-handshake-sync error', {
+              workspaceId: this.workspaceId,
+              error: err,
+            });
+          });
+      }
+    });
+
+    this.disconnect = () => {
+      this.socket.emit('client-leave-sync', this.workspaceId);
+      this.socket.off('connect', handleConnect);
+      this.socket.off('disconnect', handleDisconnect);
+    };
+  }
+
+  async waitForHandshake() {
+    await lastValueFrom(
+      this.handshook.pipe(
+        filter(v => v),
+        timeout({ first: this.SEND_TIMEOUT }),
+        take(1)
+      )
+    );
+  }
+
+  handleConnect = () => {
+    this.socket.emit('client-handshake-sync', this.workspaceId);
+  };
+
+  async pull(
+    docId: string,
+    state: Uint8Array
+  ): Promise<{ data: Uint8Array; state?: Uint8Array } | null> {
+    await this.waitForHandshake();
+    const stateVector = state ? await uint8ArrayToBase64(state) : undefined;
+
+    logger.debug('doc-load-v2', {
+      workspaceId: this.workspaceId,
+      guid: docId,
+      stateVector,
+    });
+
+    const response:
+      | { error: any }
+      | { data: { missing: string; state: string } } = await this.socket
+      .timeout(this.SEND_TIMEOUT)
+      .emitWithAck('doc-load-v2', {
+        workspaceId: this.workspaceId,
+        guid: docId,
+        stateVector,
+      });
+
+    logger.debug('doc-load callback', {
+      workspaceId: this.workspaceId,
+      guid: docId,
+      stateVector,
+      response,
+    });
+
+    if ('error' in response) {
+      // TODO: result `EventError` with server
+      if (response.error.code === 'DOC_NOT_FOUND') {
+        return null;
+      } else {
+        throw new Error(response.error.message);
+      }
+    } else {
+      return {
+        data: base64ToUint8Array(response.data.missing),
+        state: response.data.state
+          ? base64ToUint8Array(response.data.state)
+          : undefined,
+      };
+    }
+  }
+
+  async push(docId: string, update: Uint8Array) {
+    await this.waitForHandshake();
+
+    logger.debug('client-update-v2', {
+      workspaceId: this.workspaceId,
+      guid: docId,
+      update,
+    });
+
+    const payload = await uint8ArrayToBase64(update);
+
+    const response: {
+      // TODO: reuse `EventError` with server
+      error?: any;
+      data: any;
+    } = await this.socket
+      .timeout(this.SEND_TIMEOUT)
+      .emitWithAck('client-update-v2', {
+        workspaceId: this.workspaceId,
+        guid: docId,
+        updates: [payload],
+      });
+
+    // TODO: raise error with different code to users
+    if (response.error) {
+      logger.error('client-update-v2 error', {
+        workspaceId: this.workspaceId,
+        guid: docId,
+        response,
+      });
+
+      throw new Error(response.error);
+    }
+  }
+
+  async subscribe(
+    cb: (docId: string, data: Uint8Array) => void,
+    disconnect: (reason: string) => void
+  ) {
+    const handleUpdate = async (message: {
+      workspaceId: string;
+      guid: string;
+      updates: string[];
+    }) => {
+      if (message.workspaceId === this.workspaceId) {
+        message.updates.forEach(update => {
+          cb(message.guid, base64ToUint8Array(update));
+        });
+      }
+    };
+    const handleDisconnect = (reason: string) => {
+      this.socket.off('server-updates', handleUpdate);
+      disconnect(reason);
+    };
+    this.socket.on('server-updates', handleUpdate);
+
+    this.socket.on('disconnect', handleDisconnect);
+
+    return () => {
+      this.socket.off('server-updates', handleUpdate);
+      this.socket.off('disconnect', handleDisconnect);
+    };
+  }
+
+  disconnect: () => void;
+}
+
 export function createAffineStorage(
   workspaceId: string
 ): SyncStorage & { disconnect: () => void } {
-  logger.debug('createAffineStorage', workspaceId);
-  const socket = getIoManager().socket('/');
-
-  const syncSender = new MultipleBatchSyncSender(async (guid, updates) => {
-    const payload = await Promise.all(
-      updates.map(update => uint8ArrayToBase64(update))
-    );
-
-    return new Promise(resolve => {
-      socket.emit(
-        'client-update-v2',
-        {
-          workspaceId,
-          guid,
-          updates: payload,
-        },
-        (response: {
-          // TODO: reuse `EventError` with server
-          error?: any;
-          data: any;
-        }) => {
-          // TODO: raise error with different code to users
-          if (response.error) {
-            logger.error('client-update-v2 error', {
-              workspaceId,
-              guid,
-              response,
-            });
-          }
-
-          resolve({
-            accepted: !response.error,
-            // TODO: reuse `EventError` with server
-            retry: response.error?.code === 'INTERNAL',
-          });
-        }
-      );
-    });
-  });
-
-  function handleConnect() {
-    socket.emit(
-      'client-handshake-sync',
-      workspaceId,
-      (response: { error?: any }) => {
-        if (!response.error) {
-          syncSender.start();
-        }
-      }
-    );
-  }
-
-  socket.on('connect', handleConnect);
-
-  socket.connect();
-
-  socket.emit(
-    'client-handshake-sync',
-    workspaceId,
-    (response: { error?: any }) => {
-      if (!response.error) {
-        syncSender.start();
-      }
-    }
-  );
-
-  return {
-    name: 'affine-cloud',
-    async pull(docId, state) {
-      const stateVector = state ? await uint8ArrayToBase64(state) : undefined;
-
-      return new Promise((resolve, reject) => {
-        logger.debug('doc-load-v2', {
-          workspaceId: workspaceId,
-          guid: docId,
-          stateVector,
-        });
-        socket.emit(
-          'doc-load-v2',
-          {
-            workspaceId: workspaceId,
-            guid: docId,
-            stateVector,
-          },
-          (
-            response: // TODO: reuse `EventError` with server
-            { error: any } | { data: { missing: string; state: string } }
-          ) => {
-            logger.debug('doc-load callback', {
-              workspaceId: workspaceId,
-              guid: docId,
-              stateVector,
-              response,
-            });
-
-            if ('error' in response) {
-              // TODO: result `EventError` with server
-              if (response.error.code === 'DOC_NOT_FOUND') {
-                resolve(null);
-              } else {
-                reject(new Error(response.error.message));
-              }
-            } else {
-              resolve({
-                data: base64ToUint8Array(response.data.missing),
-                state: response.data.state
-                  ? base64ToUint8Array(response.data.state)
-                  : undefined,
-              });
-            }
-          }
-        );
-      });
-    },
-    async push(docId, update) {
-      logger.debug('client-update-v2', {
-        workspaceId,
-        guid: docId,
-        update,
-      });
-
-      await syncSender.send(docId, update);
-    },
-    async subscribe(cb, disconnect) {
-      const handleUpdate = async (message: {
-        workspaceId: string;
-        guid: string;
-        updates: string[];
-      }) => {
-        if (message.workspaceId === workspaceId) {
-          message.updates.forEach(update => {
-            cb(message.guid, base64ToUint8Array(update));
-          });
-        }
-      };
-      socket.on('server-updates', handleUpdate);
-
-      socket.on('disconnect', reason => {
-        socket.off('server-updates', handleUpdate);
-        disconnect(reason);
-      });
-
-      return () => {
-        socket.off('server-updates', handleUpdate);
-      };
-    },
-    disconnect() {
-      syncSender.stop();
-      socket.emit('client-leave-sync', workspaceId);
-      socket.off('connect', handleConnect);
-    },
-  };
+  return new AffineSyncStorage(workspaceId);
 }
 
 export function createAffineStaticStorage(workspaceId: string): SyncStorage {

yarn.lock

@@ -812,6 +812,7 @@ __metadata:
     lodash-es: "npm:^4.17.21"
     nanoid: "npm:^5.0.3"
     next-auth: "npm:^4.24.5"
+    rxjs: "npm:^7.8.1"
     socket.io-client: "npm:^4.7.2"
     vitest: "npm:1.1.3"
     ws: "npm:^8.14.2"