diff --git a/packages/y-indexeddb/src/__tests__/index.spec.ts b/packages/y-indexeddb/src/__tests__/index.spec.ts index 91c6f94628..d92ad54841 100644 --- a/packages/y-indexeddb/src/__tests__/index.spec.ts +++ b/packages/y-indexeddb/src/__tests__/index.spec.ts @@ -7,9 +7,17 @@ import { __unstableSchemas, AffineSchemas } from '@blocksuite/blocks/models'; import { assertExists, uuidv4, Workspace } from '@blocksuite/store'; import { openDB } from 'idb'; import { beforeEach, describe, expect, test, vi } from 'vitest'; +import { applyUpdate, Doc, encodeStateAsUpdate } from 'yjs'; import type { WorkspacePersist } from '../index'; -import { createIndexedDBProvider, dbVersion, setMergeCount } from '../index'; +import { + createIndexedDBProvider, + dbVersion, + getMilestones, + markMilestone, + revertUpdate, + setMergeCount, +} from '../index'; async function getUpdates(id: string): Promise { const db = await openDB('affine-local', dbVersion); @@ -146,3 +154,66 @@ describe('indexeddb provider', () => { } }); }); + +describe('milestone', () => { + test('milestone', async () => { + const doc = new Doc(); + const map = doc.getMap('map'); + const array = doc.getArray('array'); + map.set('1', 1); + array.push([1]); + await markMilestone('1', doc, 'test1'); + const milestones = await getMilestones('1'); + assertExists(milestones); + expect(milestones).toBeDefined(); + expect(Object.keys(milestones).length).toBe(1); + expect(milestones.test1).toBeInstanceOf(Uint8Array); + const snapshot = new Doc(); + applyUpdate(snapshot, milestones.test1); + { + const map = snapshot.getMap('map'); + expect(map.get('1')).toBe(1); + } + map.set('1', 2); + { + const map = snapshot.getMap('map'); + expect(map.get('1')).toBe(1); + } + revertUpdate(doc, milestones.test1, { + map: 'Map', + array: 'Array', + }); + { + const map = doc.getMap('map'); + expect(map.get('1')).toBe(1); + } + + const fn = vi.fn(() => true); + doc.gcFilter = fn; + expect(fn).toBeCalledTimes(0); + + for (let i = 0; i < 1e5; i++) { + map.set(`${i}`, i + 1); + } + for (let i = 0; i < 1e5; i++) { + map.delete(`${i}`); + } + for (let i = 0; i < 1e5; i++) { + map.set(`${i}`, i - 1); + } + + expect(fn).toBeCalled(); + + const doc2 = new Doc(); + applyUpdate(doc2, encodeStateAsUpdate(doc)); + + revertUpdate(doc2, milestones.test1, { + map: 'Map', + array: 'Array', + }); + { + const map = doc2.getMap('map'); + expect(map.get('1')).toBe(1); + } + }); +}); diff --git a/packages/y-indexeddb/src/index.ts b/packages/y-indexeddb/src/index.ts index 57bd5fa69d..078e129d34 100644 --- a/packages/y-indexeddb/src/index.ts +++ b/packages/y-indexeddb/src/index.ts @@ -5,13 +5,57 @@ import { diffUpdate, Doc, encodeStateAsUpdate, + encodeStateVector, mergeUpdates, + UndoManager, } from 'yjs'; const indexeddbOrigin = Symbol('indexeddb-provider-origin'); +const snapshotOrigin = Symbol('snapshot-origin'); let mergeCount = 500; +type Metadata = Record; + +export function revertUpdate( + doc: Doc, + snapshotUpdate: Uint8Array, + metadata: Metadata +) { + const snapshotDoc = new Doc(); + applyUpdate(snapshotDoc, snapshotUpdate, snapshotOrigin); + + const currentStateVector = encodeStateVector(doc); + const snapshotStateVector = encodeStateVector(snapshotDoc); + + const changesSinceSnapshotUpdate = encodeStateAsUpdate( + doc, + snapshotStateVector + ); + const undoManager = new UndoManager( + [...snapshotDoc.share.keys()].map(key => { + if (metadata[key] === 'Text') { + return snapshotDoc.getText(key); + } else if (metadata[key] === 'Map') { + return snapshotDoc.getMap(key); + } else if (metadata[key] === 'Array') { + return snapshotDoc.getArray(key); + } + throw new Error('Unknown type'); + }), + { + trackedOrigins: new Set([snapshotOrigin]), + } + ); + applyUpdate(snapshotDoc, changesSinceSnapshotUpdate, snapshotOrigin); + undoManager.undo(); + const revertChangesSinceSnapshotUpdate = encodeStateAsUpdate( + snapshotDoc, + currentStateVector + ); + applyUpdate(doc, revertChangesSinceSnapshotUpdate, snapshotOrigin); +} + export class EarlyDisconnectError extends Error { constructor() { super('Early disconnect'); @@ -69,6 +113,52 @@ export interface OldYjsDB extends DBSchema { }; } +export const markMilestone = async ( + id: string, + doc: Doc, + name: string, + dbName = 'affine-local' +): Promise => { + const dbPromise = openDB(dbName, dbVersion, { + upgrade: upgradeDB, + }); + const db = await dbPromise; + const store = db + .transaction('milestone', 'readwrite') + .objectStore('milestone'); + const milestone = await store.get('id'); + const binary = encodeStateAsUpdate(doc); + if (!milestone) { + await store.put({ + id, + milestone: { + [name]: binary, + }, + }); + } else { + milestone.milestone[name] = binary; + await store.put(milestone); + } +}; + +export const getMilestones = async ( + id: string, + dbName = 'affine-local' +): Promise => { + const dbPromise = openDB(dbName, dbVersion, { + upgrade: upgradeDB, + }); + const db = await dbPromise; + const store = db + .transaction('milestone', 'readonly') + .objectStore('milestone'); + const milestone = await store.get(id); + if (!milestone) { + return null; + } + return milestone.milestone; +}; + export const createIndexedDBProvider = ( id: string, doc: Doc,