feat(core): new worker workspace engine (#9257)

packages/common/nbstore/src/impls/idb/db.ts

@@ -25,22 +25,14 @@ export class IDBConnection extends AutoReconnectConnection<{
   }
 
   override async doConnect() {
+    // indexeddb will responsible for version control, so the db.version always match migrator.version
+    const db = await openDB<DocStorageSchema>(this.dbName, migrator.version, {
+      upgrade: migrator.migrate,
+    });
+    db.addEventListener('versionchange', this.handleVersionChange);
+
     return {
-      db: await openDB<DocStorageSchema>(this.dbName, migrator.version, {
-        upgrade: migrator.migrate,
-        blocking: () => {
-          // if, for example, an tab with newer version is opened, this function will be called.
-          // we should close current connection to allow the new version to upgrade the db.
-          this.setStatus(
-            'closed',
-            new Error('Blocking a new version. Closing the connection.')
-          );
-        },
-        blocked: () => {
-          // fallback to retry auto retry
-          this.setStatus('error', new Error('Blocked by other tabs.'));
-        },
-      }),
+      db,
       channel: new BroadcastChannel('idb:' + this.dbName),
     };
   }
@@ -49,7 +41,19 @@ export class IDBConnection extends AutoReconnectConnection<{
     db: IDBPDatabase<DocStorageSchema>;
     channel: BroadcastChannel;
   }) {
+    db.db.removeEventListener('versionchange', this.handleVersionChange);
     db.channel.close();
     db.db.close();
   }
+
+  handleVersionChange = (e: IDBVersionChangeEvent) => {
+    if (e.newVersion !== migrator.version) {
+      this.error = new Error(
+        'Database version mismatch, expected ' +
+          migrator.version +
+          ' but got ' +
+          e.newVersion
+      );
+    }
+  };
 }

packages/common/nbstore/src/impls/idb/doc.ts

@@ -29,26 +29,35 @@ export class IndexedDBDocStorage extends DocStorageBase<IDBConnectionOptions> {
 
   override locker = new IndexedDBLocker(this.connection);
 
-  private _lastTimestamp = new Date(0);
-
-  private generateTimestamp() {
-    const timestamp = new Date();
-    if (timestamp.getTime() <= this._lastTimestamp.getTime()) {
-      timestamp.setTime(this._lastTimestamp.getTime() + 1);
-    }
-    this._lastTimestamp = timestamp;
-    return timestamp;
-  }
-
   override async pushDocUpdate(update: DocUpdate, origin?: string) {
-    const trx = this.db.transaction(['updates', 'clocks'], 'readwrite');
-    const timestamp = this.generateTimestamp();
-    await trx.objectStore('updates').add({
-      ...update,
-      createdAt: timestamp,
-    });
+    let timestamp = new Date();
 
-    await trx.objectStore('clocks').put({ docId: update.docId, timestamp });
+    let retry = 0;
+
+    while (true) {
+      try {
+        const trx = this.db.transaction(['updates', 'clocks'], 'readwrite');
+
+        await trx.objectStore('updates').add({
+          ...update,
+          createdAt: timestamp,
+        });
+
+        await trx.objectStore('clocks').put({ docId: update.docId, timestamp });
+
+        trx.commit();
+      } catch (e) {
+        if (e instanceof Error && e.name === 'ConstraintError') {
+          retry++;
+          if (retry < 10) {
+            timestamp = new Date(timestamp.getTime() + 1);
+            continue;
+          }
+        }
+        throw e;
+      }
+      break;
+    }
 
     this.emit(
       'update',
@@ -191,9 +200,9 @@ export class IndexedDBDocStorage extends DocStorageBase<IDBConnectionOptions> {
     };
   }
 
-  handleChannelMessage(event: MessageEvent<ChannelMessage>) {
+  handleChannelMessage = (event: MessageEvent<ChannelMessage>) => {
     if (event.data.type === 'update') {
       this.emit('update', event.data.update, event.data.origin);
     }
-  }
+  };
 }

packages/common/nbstore/src/impls/idb/index.ts

@@ -2,7 +2,6 @@ import type { StorageConstructor } from '..';
 import { IndexedDBBlobStorage } from './blob';
 import { IndexedDBDocStorage } from './doc';
 import { IndexedDBSyncStorage } from './sync';
-import { IndexedDBV1BlobStorage, IndexedDBV1DocStorage } from './v1';
 
 export * from './blob';
 export * from './doc';
@@ -13,8 +12,3 @@ export const idbStorages = [
   IndexedDBBlobStorage,
   IndexedDBSyncStorage,
 ] satisfies StorageConstructor[];
-
-export const idbv1Storages = [
-  IndexedDBV1DocStorage,
-  IndexedDBV1BlobStorage,
-] satisfies StorageConstructor[];

packages/common/nbstore/src/impls/idb/v1/blob.ts

@@ -19,6 +19,9 @@ export class IndexedDBV1BlobStorage extends BlobStorageBase {
   }
 
   override async get(key: string) {
+    if (!this.db) {
+      return null;
+    }
     const trx = this.db.transaction('blob', 'readonly');
     const blob = await trx.store.get(key);
     if (!blob) {
@@ -34,6 +37,9 @@ export class IndexedDBV1BlobStorage extends BlobStorageBase {
   }
 
   override async delete(key: string, permanently: boolean) {
+    if (!this.db) {
+      return;
+    }
     if (permanently) {
       const trx = this.db.transaction('blob', 'readwrite');
       await trx.store.delete(key);
@@ -41,6 +47,9 @@ export class IndexedDBV1BlobStorage extends BlobStorageBase {
   }
 
   override async list() {
+    if (!this.db) {
+      return [];
+    }
     const trx = this.db.transaction('blob', 'readonly');
     const it = trx.store.iterate();
 

packages/common/nbstore/src/impls/idb/v1/db.ts

@@ -15,23 +15,26 @@ export interface DocDBSchema extends DBSchema {
   };
 }
 
-export class DocIDBConnection extends AutoReconnectConnection<
-  IDBPDatabase<DocDBSchema>
-> {
+export class DocIDBConnection extends AutoReconnectConnection<IDBPDatabase<DocDBSchema> | null> {
   override get shareId() {
     return 'idb(old):affine-local';
   }
 
   override async doConnect() {
-    return openDB<DocDBSchema>('affine-local', 1, {
-      upgrade: db => {
-        db.createObjectStore('workspace', { keyPath: 'id' });
-      },
-    });
+    const dbs = await indexedDB.databases();
+    if (dbs.some(d => d.name === 'affine-local')) {
+      return openDB<DocDBSchema>('affine-local', 1, {
+        upgrade: db => {
+          db.createObjectStore('workspace', { keyPath: 'id' });
+        },
+      });
+    } else {
+      return null;
+    }
   }
 
-  override doDisconnect(conn: IDBPDatabase<DocDBSchema>) {
-    conn.close();
+  override doDisconnect(conn: IDBPDatabase<DocDBSchema> | null) {
+    conn?.close();
   }
 }
 
@@ -46,9 +49,7 @@ export interface BlobIDBConnectionOptions {
   id: string;
 }
 
-export class BlobIDBConnection extends AutoReconnectConnection<
-  IDBPDatabase<BlobDBSchema>
-> {
+export class BlobIDBConnection extends AutoReconnectConnection<IDBPDatabase<BlobDBSchema> | null> {
   constructor(private readonly options: BlobIDBConnectionOptions) {
     super();
   }
@@ -58,14 +59,19 @@ export class BlobIDBConnection extends AutoReconnectConnection<
   }
 
   override async doConnect() {
-    return openDB<BlobDBSchema>(`${this.options.id}_blob`, 1, {
-      upgrade: db => {
-        db.createObjectStore('blob');
-      },
-    });
+    const dbs = await indexedDB.databases();
+    if (dbs.some(d => d.name === `${this.options.id}_blob`)) {
+      return openDB<BlobDBSchema>(`${this.options.id}_blob`, 1, {
+        upgrade: db => {
+          db.createObjectStore('blob');
+        },
+      });
+    } else {
+      return null;
+    }
   }
 
-  override doDisconnect(conn: IDBPDatabase<BlobDBSchema>) {
-    conn.close();
+  override doDisconnect(conn: IDBPDatabase<BlobDBSchema> | null) {
+    conn?.close();
   }
 }

packages/common/nbstore/src/impls/idb/v1/doc.ts

@@ -1,9 +1,20 @@
+import { once } from 'lodash-es';
+import {
+  applyUpdate,
+  type Array as YArray,
+  Doc as YDoc,
+  type Map as YMap,
+} from 'yjs';
+
 import { share } from '../../../connection';
 import {
+  type DocClocks,
   type DocRecord,
   DocStorageBase,
+  type DocStorageOptions,
   type DocUpdate,
 } from '../../../storage';
+import { getIdConverter } from '../../../utils/id-converter';
 import { DocIDBConnection } from './db';
 
 /**
@@ -14,6 +25,13 @@ export class IndexedDBV1DocStorage extends DocStorageBase {
 
   readonly connection = share(new DocIDBConnection());
 
+  constructor(opts: DocStorageOptions) {
+    super({
+      ...opts,
+      readonlyMode: true,
+    });
+  }
+
   get db() {
     return this.connection.inner;
   }
@@ -23,26 +41,11 @@ export class IndexedDBV1DocStorage extends DocStorageBase {
   }
 
   override async getDoc(docId: string) {
-    const trx = this.db.transaction('workspace', 'readonly');
-    const record = await trx.store.get(docId);
-
-    if (!record?.updates.length) {
+    if (!this.db) {
       return null;
     }
-
-    if (record.updates.length === 1) {
-      return {
-        docId,
-        bin: record.updates[0].update,
-        timestamp: new Date(record.updates[0].timestamp),
-      };
-    }
-
-    return {
-      docId,
-      bin: await this.mergeUpdates(record.updates.map(update => update.update)),
-      timestamp: new Date(record.updates.at(-1)?.timestamp ?? Date.now()),
-    };
+    const oldId = (await this.getIdConverter()).newIdToOldId(docId);
+    return this.rawGetDoc(oldId);
   }
 
   protected override async getDocSnapshot() {
@@ -55,12 +58,60 @@ export class IndexedDBV1DocStorage extends DocStorageBase {
   }
 
   override async deleteDoc(docId: string) {
+    if (!this.db) {
+      return;
+    }
+    const oldId = (await this.getIdConverter()).newIdToOldId(docId);
     const trx = this.db.transaction('workspace', 'readwrite');
-    await trx.store.delete(docId);
+    await trx.store.delete(oldId);
   }
 
-  override async getDocTimestamps() {
-    return {};
+  override async getDocTimestamps(): Promise<DocClocks> {
+    if (!this.db) {
+      return {};
+    }
+
+    const idConverter = await this.getIdConverter();
+
+    const oldIds: string[] = [this.spaceId];
+
+    const rootDocBuffer = await this.rawGetDoc(this.spaceId);
+    if (rootDocBuffer) {
+      const ydoc = new YDoc({
+        guid: this.spaceId,
+      });
+      applyUpdate(ydoc, rootDocBuffer.bin);
+
+      // get all ids from rootDoc.meta.pages.[*].id, trust this id as normalized id
+      const normalizedDocIds = (
+        (ydoc.getMap('meta') as YMap<any> | undefined)?.get('pages') as
+          | YArray<YMap<any>>
+          | undefined
+      )
+        ?.map(i => i.get('id') as string)
+        .filter(i => !!i);
+
+      const spaces = ydoc.getMap('spaces') as YMap<any> | undefined;
+      for (const pageId of normalizedDocIds ?? []) {
+        const subdoc = spaces?.get(pageId);
+        if (subdoc && subdoc instanceof YDoc) {
+          oldIds.push(subdoc.guid);
+        }
+      }
+    }
+
+    const trx = this.db.transaction('workspace', 'readonly');
+    const allKeys = await trx.store.getAllKeys();
+    oldIds.push(...allKeys.filter(k => k.startsWith(`db$${this.spaceId}$`)));
+    oldIds.push(
+      ...allKeys.filter(k =>
+        k.match(new RegExp(`^userdata\\$[\\w-]+\\$${this.spaceId}$`))
+      )
+    );
+
+    return Object.fromEntries(
+      oldIds.map(id => [idConverter.oldIdToNewId(id), new Date(1)])
+    );
   }
 
   override async getDocTimestamp(_docId: string) {
@@ -78,4 +129,59 @@ export class IndexedDBV1DocStorage extends DocStorageBase {
   protected override async markUpdatesMerged(): Promise<number> {
     return 0;
   }
+
+  private async rawGetDoc(id: string) {
+    if (!this.db) {
+      return null;
+    }
+    const trx = this.db.transaction('workspace', 'readonly');
+    const record = await trx.store.get(id);
+
+    if (!record?.updates.length) {
+      return null;
+    }
+
+    if (record.updates.length === 1) {
+      return {
+        docId: id,
+        bin: record.updates[0].update,
+        timestamp: new Date(record.updates[0].timestamp),
+      };
+    }
+
+    return {
+      docId: id,
+      bin: await this.mergeUpdates(record.updates.map(update => update.update)),
+      timestamp: new Date(record.updates.at(-1)?.timestamp ?? Date.now()),
+    };
+  }
+
+  private readonly getIdConverter = once(async () => {
+    const idConverter = getIdConverter(
+      {
+        getDocBuffer: async id => {
+          if (!this.db) {
+            return null;
+          }
+          const trx = this.db.transaction('workspace', 'readonly');
+          const record = await trx.store.get(id);
+
+          if (!record?.updates.length) {
+            return null;
+          }
+
+          if (record.updates.length === 1) {
+            return record.updates[0].update;
+          }
+
+          return await this.mergeUpdates(
+            record.updates.map(update => update.update)
+          );
+        },
+      },
+      this.spaceId
+    );
+
+    return await idConverter;
+  });
 }

packages/common/nbstore/src/impls/idb/v1/index.ts

@@ -1,2 +1,11 @@
+import type { StorageConstructor } from '../..';
+import { IndexedDBV1BlobStorage } from './blob';
+import { IndexedDBV1DocStorage } from './doc';
+
 export * from './blob';
 export * from './doc';
+
+export const idbV1Storages = [
+  IndexedDBV1DocStorage,
+  IndexedDBV1BlobStorage,
+] satisfies StorageConstructor[];