feat(server): implement doc service (#9961)

close CLOUD-94

packages/backend/server/src/core/doc/__tests__/reader-from-database.spec.ts

@@ -0,0 +1,72 @@
+import { randomUUID } from 'node:crypto';
+
+import { User, Workspace } from '@prisma/client';
+import ava, { TestFn } from 'ava';
+
+import { createTestingApp, type TestingApp } from '../../../__tests__/utils';
+import { AppModule } from '../../../app.module';
+import { ConfigModule } from '../../../base/config';
+import { Models } from '../../../models';
+import { DocReader } from '..';
+import { DatabaseDocReader } from '../reader';
+
+const test = ava as TestFn<{
+  models: Models;
+  app: TestingApp;
+  docReader: DocReader;
+}>;
+
+test.before(async t => {
+  const { app } = await createTestingApp({
+    imports: [ConfigModule.forRoot(), AppModule],
+  });
+
+  t.context.models = app.get(Models);
+  t.context.docReader = app.get(DocReader);
+  t.context.app = app;
+});
+
+let user: User;
+let workspace: Workspace;
+
+test.beforeEach(async t => {
+  await t.context.app.initTestingDB();
+  user = await t.context.models.user.create({
+    email: 'test@affine.pro',
+  });
+  workspace = await t.context.models.workspace.create(user.id);
+});
+
+test.after.always(async t => {
+  await t.context.app.close();
+});
+
+test('should return null when doc not found', async t => {
+  const { docReader } = t.context;
+  const docId = randomUUID();
+  const doc = await docReader.getDoc(workspace.id, docId);
+  t.is(doc, null);
+});
+
+test('should return doc when found', async t => {
+  const { docReader } = t.context;
+  t.true(docReader instanceof DatabaseDocReader);
+
+  const docId = randomUUID();
+  const timestamp = Date.now();
+  await t.context.models.doc.createUpdates([
+    {
+      spaceId: workspace.id,
+      docId,
+      blob: Buffer.from('blob1 data'),
+      timestamp,
+      editorId: user.id,
+    },
+  ]);
+
+  const doc = await docReader.getDoc(workspace.id, docId);
+  t.truthy(doc);
+  t.is(doc!.bin.toString(), 'blob1 data');
+  t.is(doc!.timestamp, timestamp);
+  t.is(doc!.editor, user.id);
+});

packages/backend/server/src/core/doc/__tests__/reader-from-rpc.spec.ts

@@ -0,0 +1,124 @@
+import { randomUUID } from 'node:crypto';
+import { mock } from 'node:test';
+
+import { User, Workspace } from '@prisma/client';
+import ava, { TestFn } from 'ava';
+
+import { createTestingApp, type TestingApp } from '../../../__tests__/utils';
+import { AppModule } from '../../../app.module';
+import { Config, InternalServerError } from '../../../base';
+import { ConfigModule } from '../../../base/config';
+import { Models } from '../../../models';
+import { DocReader } from '..';
+import { RpcDocReader } from '../reader';
+
+const test = ava as TestFn<{
+  models: Models;
+  app: TestingApp;
+  docReader: DocReader;
+  config: Config;
+}>;
+
+test.before(async t => {
+  const { app } = await createTestingApp({
+    imports: [
+      ConfigModule.forRoot({
+        flavor: {
+          doc: false,
+        },
+        docService: {
+          endpoint: '',
+        },
+      }),
+      AppModule,
+    ],
+  });
+
+  t.context.models = app.get(Models);
+  t.context.docReader = app.get(DocReader);
+  t.context.config = app.get(Config);
+  t.context.app = app;
+});
+
+let user: User;
+let workspace: Workspace;
+
+test.beforeEach(async t => {
+  t.context.config.docService.endpoint = t.context.app.getHttpServerUrl();
+  await t.context.app.initTestingDB();
+  user = await t.context.models.user.create({
+    email: 'test@affine.pro',
+  });
+  workspace = await t.context.models.workspace.create(user.id);
+});
+
+test.afterEach.always(() => {
+  mock.reset();
+});
+
+test.after.always(async t => {
+  await t.context.app.close();
+});
+
+test('should return null when doc not found', async t => {
+  const { docReader } = t.context;
+  const docId = randomUUID();
+  const doc = await docReader.getDoc(workspace.id, docId);
+  t.is(doc, null);
+});
+
+test('should throw error when doc service internal error', async t => {
+  const { docReader } = t.context;
+  const docId = randomUUID();
+  mock.method(docReader, 'getDoc', async () => {
+    throw new InternalServerError('mock doc service internal error');
+  });
+  await t.throwsAsync(docReader.getDoc(workspace.id, docId), {
+    instanceOf: InternalServerError,
+  });
+});
+
+test('should fallback to database doc service when endpoint network error', async t => {
+  const { docReader } = t.context;
+  t.context.config.docService.endpoint = 'http://localhost:13010';
+  const docId = randomUUID();
+  const timestamp = Date.now();
+  await t.context.models.doc.createUpdates([
+    {
+      spaceId: workspace.id,
+      docId,
+      blob: Buffer.from('blob1 data'),
+      timestamp,
+      editorId: user.id,
+    },
+  ]);
+
+  const doc = await docReader.getDoc(workspace.id, docId);
+  t.truthy(doc);
+  t.is(doc!.bin.toString(), 'blob1 data');
+  t.is(doc!.timestamp, timestamp);
+  t.is(doc!.editor, user.id);
+});
+
+test('should return doc when found', async t => {
+  const { docReader } = t.context;
+  t.true(docReader instanceof RpcDocReader);
+
+  const docId = randomUUID();
+  const timestamp = Date.now();
+  await t.context.models.doc.createUpdates([
+    {
+      spaceId: workspace.id,
+      docId,
+      blob: Buffer.from('blob1 data'),
+      timestamp,
+      editorId: user.id,
+    },
+  ]);
+
+  const doc = await docReader.getDoc(workspace.id, docId);
+  t.truthy(doc);
+  t.is(doc!.bin.toString(), 'blob1 data');
+  t.is(doc!.timestamp, timestamp);
+  t.is(doc!.editor, user.id);
+});

packages/backend/server/src/core/doc/adapters/workspace.ts

@@ -41,7 +41,9 @@ declare global {
 }
 @Injectable()
 export class PgWorkspaceDocStorageAdapter extends DocStorageAdapter {
-  private readonly logger = new Logger(PgWorkspaceDocStorageAdapter.name);
+  protected override readonly logger = new Logger(
+    PgWorkspaceDocStorageAdapter.name
+  );
 
   constructor(
     private readonly models: Models,

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

@@ -8,6 +8,7 @@ import { PgUserspaceDocStorageAdapter } from './adapters/userspace';
 import { PgWorkspaceDocStorageAdapter } from './adapters/workspace';
 import { DocStorageCronJob } from './job';
 import { DocStorageOptions } from './options';
+import { DocReader, DocReaderProvider } from './reader';
 
 @Module({
   imports: [QuotaModule, PermissionModule],
@@ -16,10 +17,15 @@ import { DocStorageOptions } from './options';
     PgWorkspaceDocStorageAdapter,
     PgUserspaceDocStorageAdapter,
     DocStorageCronJob,
+    DocReaderProvider,
   ],
   exports: [PgWorkspaceDocStorageAdapter, PgUserspaceDocStorageAdapter],
 })
 export class DocStorageModule {}
-export { PgUserspaceDocStorageAdapter, PgWorkspaceDocStorageAdapter };
+export {
+  DocReader,
+  PgUserspaceDocStorageAdapter,
+  PgWorkspaceDocStorageAdapter,
+};
 
 export { DocStorageAdapter, type Editor } from './storage';

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

@@ -0,0 +1,93 @@
+import { FactoryProvider, Injectable, Logger } from '@nestjs/common';
+import { ModuleRef } from '@nestjs/core';
+import { ClsService } from 'nestjs-cls';
+
+import { Config, CryptoHelper, UserFriendlyError } from '../../base';
+import { PgWorkspaceDocStorageAdapter } from './adapters/workspace';
+import { type DocRecord } from './storage';
+
+export abstract class DocReader {
+  abstract getDoc(
+    workspaceId: string,
+    docId: string
+  ): Promise<DocRecord | null>;
+}
+
+@Injectable()
+export class DatabaseDocReader extends DocReader {
+  constructor(protected readonly workspace: PgWorkspaceDocStorageAdapter) {
+    super();
+  }
+
+  async getDoc(workspaceId: string, docId: string): Promise<DocRecord | null> {
+    return await this.workspace.getDoc(workspaceId, docId);
+  }
+}
+
+@Injectable()
+export class RpcDocReader extends DatabaseDocReader {
+  private readonly logger = new Logger(DocReader.name);
+
+  constructor(
+    private readonly config: Config,
+    private readonly crypto: CryptoHelper,
+    private readonly cls: ClsService,
+    protected override readonly workspace: PgWorkspaceDocStorageAdapter
+  ) {
+    super(workspace);
+  }
+
+  override async getDoc(
+    workspaceId: string,
+    docId: string
+  ): Promise<DocRecord | null> {
+    const url = `${this.config.docService.endpoint}/rpc/workspaces/${workspaceId}/docs/${docId}`;
+    try {
+      const res = await fetch(url, {
+        headers: {
+          'x-access-token': this.crypto.sign(docId),
+          'x-rpc-trace-id': this.cls.getId(),
+        },
+      });
+      if (!res.ok) {
+        if (res.status === 404) {
+          return null;
+        }
+        const body = (await res.json()) as UserFriendlyError;
+        throw UserFriendlyError.fromUserFriendlyErrorJSON(body);
+      }
+      const timestamp = res.headers.get('x-doc-timestamp') as string;
+      const editor = res.headers.get('x-doc-editor-id') as string;
+      const bin = await res.arrayBuffer();
+      return {
+        spaceId: workspaceId,
+        docId,
+        bin: Buffer.from(bin),
+        timestamp: parseInt(timestamp),
+        editor,
+      };
+    } catch (err) {
+      if (err instanceof UserFriendlyError) {
+        throw err;
+      }
+      // other error
+      this.logger.error(
+        `Failed to fetch doc ${url}, error: ${err}`,
+        (err as Error).stack
+      );
+      // fallback to database doc service if the error is not user friendly, like network error
+      return await super.getDoc(workspaceId, docId);
+    }
+  }
+}
+
+export const DocReaderProvider: FactoryProvider = {
+  provide: DocReader,
+  useFactory: (config: Config, ref: ModuleRef) => {
+    if (config.flavor.doc) {
+      return ref.create(DatabaseDocReader);
+    }
+    return ref.create(RpcDocReader);
+  },
+  inject: [Config, ModuleRef],
+};

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

@@ -1,3 +1,4 @@
+import { Logger } from '@nestjs/common';
 import {
   applyUpdate,
   diffUpdate,
@@ -49,6 +50,7 @@ export interface DocStorageOptions {
 
 export abstract class DocStorageAdapter extends Connection {
   private readonly locker = new SingletonLocker();
+  protected readonly logger = new Logger(DocStorageAdapter.name);
 
   constructor(
     protected readonly options: DocStorageOptions = {
@@ -76,6 +78,9 @@ export abstract class DocStorageAdapter extends Connection {
     const updates = await this.getDocUpdates(spaceId, docId);
 
     if (updates.length) {
+      this.logger.log(
+        `Squashing updates, spaceId: ${spaceId}, docId: ${docId}, updates: ${updates.length}`
+      );
       const { timestamp, bin, editor } = await this.squash(
         snapshot ? [snapshot, ...updates] : updates
       );
@@ -96,7 +101,12 @@ export abstract class DocStorageAdapter extends Connection {
       }
 
       // always mark updates as merged unless throws
-      await this.markUpdatesMerged(spaceId, docId, updates);
+      const count = await this.markUpdatesMerged(spaceId, docId, updates);
+      if (count > 0) {
+        this.logger.log(
+          `Marked ${count} updates as merged, spaceId: ${spaceId}, docId: ${docId}`
+        );
+      }
 
       return newSnapshot;
     }