perf(electron): add index for updates (#6951)

![image.png](https://graphite-user-uploaded-assets-prod.s3.amazonaws.com/T2klNLEk0wxLh4NRDzhk/cd2e982a-f78a-4cc3-b090-ee4c0090e19d.png)

Above image shows the performance on querying a 20k rows of updates table, which is super slow at 150+ms. After adding index for doc_id the performance should be greatly improved.

After:
![image.png](https://graphite-user-uploaded-assets-prod.s3.amazonaws.com/T2klNLEk0wxLh4NRDzhk/45ea4389-1833-4dc5-bd64-84d8c99cd647.png)

fix TOV-866

packages/frontend/electron/src/helper/db/db-adapter.ts

@@ -2,16 +2,14 @@ import type { InsertRow } from '@affine/native';
 import { SqliteConnection, ValidationResult } from '@affine/native';
 import { WorkspaceVersion } from '@toeverything/infra/blocksuite';
 
-import { applyGuidCompatibilityFix, migrateToLatest } from '../db/migration';
 import { logger } from '../logger';
+import { applyGuidCompatibilityFix, migrateToLatest } from './migration';
 
 /**
  * A base class for SQLite DB adapter that provides basic methods around updates & blobs
  */
-export abstract class BaseSQLiteAdapter {
+export class SQLiteAdapter {
   db: SqliteConnection | null = null;
-  abstract role: string;
-
   constructor(public readonly path: string) {}
 
   async connectIfNeeded() {
@@ -27,7 +25,7 @@ export abstract class BaseSQLiteAdapter {
         await migrateToLatest(this.path, WorkspaceVersion.Surface);
       }
       await applyGuidCompatibilityFix(this.db);
-      logger.info(`[SQLiteAdapter:${this.role}]`, 'connected:', this.path);
+      logger.info(`[SQLiteAdapter]`, 'connected:', this.path);
     }
     return this.db;
   }
@@ -36,7 +34,7 @@ export abstract class BaseSQLiteAdapter {
     const { db } = this;
     this.db = null;
     // log after close will sometimes crash the app when quitting
-    logger.info(`[SQLiteAdapter:${this.role}]`, 'destroyed:', this.path);
+    logger.info(`[SQLiteAdapter]`, 'destroyed:', this.path);
     await db?.close();
   }
 
@@ -128,7 +126,7 @@ export abstract class BaseSQLiteAdapter {
       const start = performance.now();
       await this.db.insertUpdates(updates);
       logger.debug(
-        `[SQLiteAdapter][${this.role}] addUpdateToSQLite`,
+        `[SQLiteAdapter] addUpdateToSQLite`,
         'length:',
         updates.length,
         'docids',
@@ -140,4 +138,41 @@ export abstract class BaseSQLiteAdapter {
       logger.error('addUpdateToSQLite', this.path, error);
     }
   }
+
+  async deleteUpdates(docId?: string) {
+    try {
+      if (!this.db) {
+        logger.warn(`${this.path} is not connected`);
+        return;
+      }
+      await this.db.deleteUpdates(docId);
+    } catch (error) {
+      logger.error('deleteUpdates', error);
+    }
+  }
+
+  async getUpdatesCount(docId?: string) {
+    try {
+      if (!this.db) {
+        logger.warn(`${this.path} is not connected`);
+        return 0;
+      }
+      return await this.db.getUpdatesCount(docId);
+    } catch (error) {
+      logger.error('getUpdatesCount', error);
+      return 0;
+    }
+  }
+
+  async replaceUpdates(docId: string | null | undefined, updates: InsertRow[]) {
+    try {
+      if (!this.db) {
+        logger.warn(`${this.path} is not connected`);
+        return;
+      }
+      await this.db.replaceUpdates(docId, updates);
+    } catch (error) {
+      logger.error('replaceUpdates', error);
+    }
+  }
 }

packages/frontend/electron/src/helper/db/index.ts

@@ -5,22 +5,21 @@ import { ensureSQLiteDB } from './ensure-db';
 export * from './ensure-db';
 
 export const dbHandlers = {
-  getDocAsUpdates: async (workspaceId: string, subdocId?: string) => {
+  getDocAsUpdates: async (workspaceId: string, subdocId: string) => {
     const workspaceDB = await ensureSQLiteDB(workspaceId);
     return workspaceDB.getDocAsUpdates(subdocId);
   },
   applyDocUpdate: async (
     workspaceId: string,
     update: Uint8Array,
-    subdocId?: string
+    subdocId: string
   ) => {
     const workspaceDB = await ensureSQLiteDB(workspaceId);
-    return workspaceDB.addUpdateToSQLite([
-      {
-        data: update,
-        docId: subdocId,
-      },
-    ]);
+    return workspaceDB.addUpdateToSQLite(update, subdocId);
+  },
+  deleteDoc: async (workspaceId: string, subdocId: string) => {
+    const workspaceDB = await ensureSQLiteDB(workspaceId);
+    return workspaceDB.deleteUpdate(subdocId);
   },
   addBlob: async (workspaceId: string, key: string, data: Uint8Array) => {
     const workspaceDB = await ensureSQLiteDB(workspaceId);

packages/frontend/electron/src/helper/db/workspace-db-adapter.ts

@@ -1,36 +1,43 @@
-import type { InsertRow } from '@affine/native';
+import { AsyncLock } from '@toeverything/infra';
 import { Subject } from 'rxjs';
 import { applyUpdate, Doc as YDoc } from 'yjs';
 
 import { logger } from '../logger';
 import { getWorkspaceMeta } from '../workspace/meta';
-import { BaseSQLiteAdapter } from './base-db-adapter';
+import { SQLiteAdapter } from './db-adapter';
 import { mergeUpdate } from './merge-update';
 
 const TRIM_SIZE = 500;
 
-export class WorkspaceSQLiteDB extends BaseSQLiteAdapter {
-  role = 'primary';
-
+export class WorkspaceSQLiteDB {
+  lock = new AsyncLock();
   update$ = new Subject<void>();
+  adapter = new SQLiteAdapter(this.path);
 
   constructor(
-    public override path: string,
+    public path: string,
     public workspaceId: string
-  ) {
-    super(path);
+  ) {}
+
+  async transaction<T>(cb: () => Promise<T>): Promise<T> {
+    using _lock = await this.lock.acquire();
+    return await cb();
   }
 
-  override async destroy() {
-    await super.destroy();
+  async destroy() {
+    await this.adapter.destroy();
 
     // when db is closed, we can safely remove it from ensure-db list
     this.update$.complete();
   }
 
+  toDBDocId = (docId: string) => {
+    return this.workspaceId === docId ? undefined : docId;
+  };
+
   getWorkspaceName = async () => {
     const ydoc = new YDoc();
-    const updates = await this.getUpdates();
+    const updates = await this.adapter.getUpdates();
     updates.forEach(update => {
       applyUpdate(ydoc, update.data);
     });
@@ -38,44 +45,75 @@ export class WorkspaceSQLiteDB extends BaseSQLiteAdapter {
   };
 
   async init() {
-    const db = await super.connectIfNeeded();
+    const db = await this.adapter.connectIfNeeded();
     await this.tryTrim();
     return db;
   }
 
+  async get(docId: string) {
+    return this.adapter.getUpdates(docId);
+  }
+
   // getUpdates then encode
-  getDocAsUpdates = async (docId?: string) => {
-    const updates = await this.getUpdates(docId);
-    return mergeUpdate(updates.map(row => row.data));
+  getDocAsUpdates = async (docId: string) => {
+    const dbID = this.toDBDocId(docId);
+    const update = await this.tryTrim(dbID);
+    if (update) {
+      return update;
+    } else {
+      const updates = await this.adapter.getUpdates(dbID);
+      return mergeUpdate(updates.map(row => row.data));
+    }
   };
 
-  override async addBlob(key: string, value: Uint8Array) {
+  async addBlob(key: string, value: Uint8Array) {
     this.update$.next();
-    const res = await super.addBlob(key, value);
+    const res = await this.adapter.addBlob(key, value);
     return res;
   }
 
-  override async deleteBlob(key: string) {
-    this.update$.next();
-    await super.deleteBlob(key);
+  async getBlob(key: string) {
+    return this.adapter.getBlob(key);
   }
 
-  override async addUpdateToSQLite(data: InsertRow[]) {
-    this.update$.next();
-    await super.addUpdateToSQLite(data);
+  async getBlobKeys() {
+    return this.adapter.getBlobKeys();
   }
 
-  private readonly tryTrim = async (docId?: string) => {
-    const count = (await this.db?.getUpdatesCount(docId)) ?? 0;
+  async deleteBlob(key: string) {
+    this.update$.next();
+    await this.adapter.deleteBlob(key);
+  }
+
+  async addUpdateToSQLite(update: Uint8Array, subdocId: string) {
+    this.update$.next();
+    await this.adapter.addUpdateToSQLite([
+      {
+        data: update,
+        docId: this.toDBDocId(subdocId),
+      },
+    ]);
+  }
+
+  async deleteUpdate(subdocId: string) {
+    this.update$.next();
+    await this.adapter.deleteUpdates(this.toDBDocId(subdocId));
+  }
+
+  private readonly tryTrim = async (dbID?: string) => {
+    const count = (await this.adapter?.getUpdatesCount(dbID)) ?? 0;
     if (count > TRIM_SIZE) {
-      logger.debug(`trim ${this.workspaceId}:${docId} ${count}`);
-      const update = await this.getDocAsUpdates(docId);
-      if (update) {
-        const insertRows = [{ data: update, docId }];
-        await this.db?.replaceUpdates(docId, insertRows);
-        logger.debug(`trim ${this.workspaceId}:${docId} successfully`);
-      }
+      return await this.transaction(async () => {
+        logger.debug(`trim ${this.workspaceId}:${dbID} ${count}`);
+        const updates = await this.adapter.getUpdates(dbID);
+        const update = mergeUpdate(updates.map(row => row.data));
+        const insertRows = [{ data: update, dbID }];
+        await this.adapter?.replaceUpdates(dbID, insertRows);
+        logger.debug(`trim ${this.workspaceId}:${dbID} successfully`);
+        return update;
+      });
     }
+    return null;
   };
 }