@dxos/echo-pipeline 0.6.0 → 0.6.1-main.04e8aa0

package/src/automerge/automerge-doc-loader.ts

@@ -225,7 +225,7 @@ export class AutomergeDocumentLoaderImpl implements AutomergeDocumentLoader {
 
   private async _createObjectOnDocumentLoad(handle: DocHandle<SpaceDoc>, objectId: string) {
     try {
-      await handle.doc(['ready']);
+      await handle.whenReady();
       const logMeta = { objectId, docUrl: handle.url };
       if (this.onObjectDocumentLoaded.listenerCount() === 0) {
         log.info('document loaded after all listeners were removed', logMeta);

package/src/automerge/automerge-host.test.ts

@@ -4,7 +4,9 @@
 
 import expect from 'expect';
 
+import { getHeads } from '@dxos/automerge/automerge';
 import { IndexMetadataStore } from '@dxos/indexing';
+import type { LevelDB } from '@dxos/kv-store';
 import { createTestLevel } from '@dxos/kv-store/testing';
 import { afterTest, describe, test } from '@dxos/test';
 
@@ -15,13 +17,8 @@ describe('AutomergeHost', () => {
     const level = createTestLevel();
     await level.open();
     afterTest(() => level.close());
-    const host = new AutomergeHost({
-      db: level.sublevel('automerge'),
-      indexMetadataStore: new IndexMetadataStore({ db: level.sublevel('index-metadata') }),
-    });
-    await host.open();
-    afterTest(() => host.close());
 
+    const host = await setupAutomergeHost({ level });
     const handle = host.repo.create();
     handle.change((doc: any) => {
       doc.text = 'Hello world';
@@ -35,11 +32,7 @@ describe('AutomergeHost', () => {
     await level.open();
     afterTest(() => level.close());
 
-    const host = new AutomergeHost({
-      db: level.sublevel('automerge'),
-      indexMetadataStore: new IndexMetadataStore({ db: level.sublevel('index-metadata') }),
-    });
-    await host.open();
+    const host = await setupAutomergeHost({ level });
     const handle = host.repo.create();
     handle.change((doc: any) => {
       doc.text = 'Hello world';
@@ -49,15 +42,39 @@ describe('AutomergeHost', () => {
     await host.repo.flush();
     await host.close();
 
-    const host2 = new AutomergeHost({
-      db: level.sublevel('automerge'),
-      indexMetadataStore: new IndexMetadataStore({ db: level.sublevel('index-metadata') }),
-    });
-    await host2.open();
-    afterTest(() => host2.close());
+    const host2 = await setupAutomergeHost({ level });
     const handle2 = host2.repo.find(url);
     await handle2.whenReady();
     expect(handle2.docSync().text).toEqual('Hello world');
     await host2.repo.flush();
   });
+
+  test('query document heads', async () => {
+    const level = createTestLevel();
+    await level.open();
+    afterTest(() => level.close());
+
+    const host = await setupAutomergeHost({ level });
+    const handle = host.createDoc({ text: 'Hello world' });
+    const expectedHeads = getHeads(handle.docSync());
+    await host.flush();
+
+    expect(await host.getHeads(handle.documentId)).toEqual(expectedHeads);
+
+    // Simulate a restart.
+    {
+      const host = await setupAutomergeHost({ level });
+      expect(await host.getHeads(handle.documentId)).toEqual(expectedHeads);
+    }
+  });
 });
+
+const setupAutomergeHost = async ({ level }: { level: LevelDB }) => {
+  const host = new AutomergeHost({
+    db: level,
+    indexMetadataStore: new IndexMetadataStore({ db: level.sublevel('index-metadata') }),
+  });
+  await host.open();
+  afterTest(() => host.close());
+  return host;
+};

package/src/automerge/automerge-host.ts

@@ -22,11 +22,13 @@ import {
   type StorageAdapterInterface,
 } from '@dxos/automerge/automerge-repo';
 import { type Stream } from '@dxos/codec-protobuf';
-import { Context, cancelWithContext, type Lifecycle } from '@dxos/context';
+import { type Context, Resource, cancelWithContext, type Lifecycle } from '@dxos/context';
 import { type SpaceDoc } from '@dxos/echo-protocol';
 import { type IndexMetadataStore } from '@dxos/indexing';
+import { invariant } from '@dxos/invariant';
 import { PublicKey } from '@dxos/keys';
-import { type SublevelDB } from '@dxos/kv-store';
+import { type LevelDB } from '@dxos/kv-store';
+import { log } from '@dxos/log';
 import { objectPointerCodec } from '@dxos/protocols';
 import {
   type FlushRequest,
@@ -39,6 +41,7 @@ import { mapValues } from '@dxos/util';
 
 import { EchoNetworkAdapter, isEchoPeerMetadata } from './echo-network-adapter';
 import { type EchoReplicator } from './echo-replicator';
+import { HeadsStore } from './heads-store';
 import { LevelDBStorageAdapter, type BeforeSaveParams } from './leveldb-storage-adapter';
 import { LocalHostNetworkAdapter } from './local-host-network-adapter';
 
@@ -46,7 +49,7 @@ import { LocalHostNetworkAdapter } from './local-host-network-adapter';
 export type { DocumentId };
 
 export type AutomergeHostParams = {
-  db: SublevelDB;
+  db: LevelDB;
 
   indexMetadataStore: IndexMetadataStore;
 };
@@ -66,9 +69,9 @@ export type CreateDocOptions = {
  * Abstracts over the AutomergeRepo.
  */
 @trace.resource()
-export class AutomergeHost {
+export class AutomergeHost extends Resource {
+  private readonly _db: LevelDB;
   private readonly _indexMetadataStore: IndexMetadataStore;
-  private readonly _ctx = new Context();
   private readonly _echoNetworkAdapter = new EchoNetworkAdapter({
     getContainingSpaceForDocument: this._getContainingSpaceForDocument.bind(this),
   });
@@ -76,22 +79,26 @@ export class AutomergeHost {
   private _repo!: Repo;
   private _clientNetwork!: LocalHostNetworkAdapter;
   private _storage!: StorageAdapterInterface & Lifecycle;
+  private readonly _headsStore: HeadsStore;
 
   @trace.info()
   private _peerId!: string;
 
   constructor({ db, indexMetadataStore }: AutomergeHostParams) {
+    super();
+    this._db = db;
     this._storage = new LevelDBStorageAdapter({
-      db,
+      db: db.sublevel('automerge'),
       callbacks: {
         beforeSave: async (params) => this._beforeSave(params),
         afterSave: async () => this._afterSave(),
       },
     });
+    this._headsStore = new HeadsStore({ db: db.sublevel('heads') });
     this._indexMetadataStore = indexMetadataStore;
   }
 
-  async open() {
+  protected override async _open() {
     // TODO(burdon): Should this be stable?
     this._peerId = `host-${PublicKey.random().toHex()}` as PeerId;
 
@@ -117,7 +124,7 @@ export class AutomergeHost {
     await this._echoNetworkAdapter.whenConnected();
   }
 
-  async close() {
+  protected override async _close() {
     await this._storage.close?.();
     await this._clientNetwork.close();
     await this._echoNetworkAdapter.close();
@@ -179,14 +186,32 @@ export class AutomergeHost {
     }
   }
 
+  async reIndexHeads(documentIds: DocumentId[]) {
+    for (const documentId of documentIds) {
+      log.info('reindexing heads for document', { documentId });
+      const handle = this._repo.find(documentId);
+      await handle.whenReady(['ready', 'requesting']);
+      if (handle.inState(['requesting'])) {
+        log.warn('document is not available locally, skipping', { documentId });
+        continue; // Handle not available locally.
+      }
+
+      const doc = handle.docSync();
+      invariant(doc);
+
+      const heads = getHeads(doc);
+      const batch = this._db.batch();
+      this._headsStore.setHeads(documentId, heads, batch);
+      await batch.write();
+    }
+    log.info('done reindexing heads');
+  }
+
   // TODO(dmaretskyi): Share based on HALO permissions and space affinity.
   // Hosts, running in the worker, don't share documents unless requested by other peers.
   // NOTE: If both peers return sharePolicy=false the replication will not happen
   // https://github.com/automerge/automerge-repo/pull/292
-  private async _sharePolicy(
-    peerId: PeerId /* device key */,
-    documentId?: DocumentId /* space key */,
-  ): Promise<boolean> {
+  private async _sharePolicy(peerId: PeerId, documentId?: DocumentId): Promise<boolean> {
     if (peerId.startsWith('client-')) {
       return false; // Only send docs to clients if they are requested.
     }
@@ -215,13 +240,15 @@ export class AutomergeHost {
 
     const spaceKey = getSpaceKeyFromDoc(doc) ?? undefined;
 
-    const lastAvailableHash = getHeads(doc);
+    const heads = getHeads(doc);
+
+    this._headsStore.setHeads(handle.documentId, heads, batch);
 
     const objectIds = Object.keys(doc.objects ?? {});
     const encodedIds = objectIds.map((objectId) =>
       objectPointerCodec.encode({ documentId: handle.documentId, objectId, spaceKey }),
     );
-    const idToLastHash = new Map(encodedIds.map((id) => [id, lastAvailableHash]));
+    const idToLastHash = new Map(encodedIds.map((id) => [id, heads]));
     this._indexMetadataStore.markDirty(idToLastHash, batch);
   }
 
@@ -281,7 +308,7 @@ export class AutomergeHost {
    * Flush documents to disk.
    */
   @trace.span({ showInBrowserTimeline: true })
-  async flush({ states }: FlushRequest): Promise<void> {
+  async flush({ states }: FlushRequest = {}): Promise<void> {
     // Note: Wait for all requested documents to be loaded/synced from thin-client.
     if (states) {
       await Promise.all(
@@ -289,7 +316,7 @@ export class AutomergeHost {
           if (!heads) {
             return;
           }
-          const handle = this.repo.handles[documentId as DocumentId] ?? this._repo.find(documentId as DocumentId);
+          const handle = this._repo.handles[documentId as DocumentId] ?? this._repo.find(documentId as DocumentId);
           await waitForHeads(handle, heads);
         }) ?? [],
       );
@@ -298,6 +325,19 @@ export class AutomergeHost {
     await this._repo.flush(states?.map(({ documentId }) => documentId as DocumentId));
   }
 
+  async getHeads(documentId: DocumentId): Promise<Heads | undefined> {
+    const handle = this._repo.handles[documentId];
+    if (handle) {
+      const doc = handle.docSync();
+      if (!doc) {
+        return undefined;
+      }
+      return getHeads(doc);
+    } else {
+      return this._headsStore.getHeads(documentId);
+    }
+  }
+
   /**
    * Host <-> Client sync.
    */

package/src/automerge/automerge-repo.test.ts

@@ -7,7 +7,14 @@ import waitForExpect from 'wait-for-expect';
 
 import { asyncTimeout, sleep } from '@dxos/async';
 import { type Heads, change, clone, equals, from, getBackend, getHeads } from '@dxos/automerge/automerge';
-import { type Message, Repo, type PeerId, type DocumentId, type HandleState } from '@dxos/automerge/automerge-repo';
+import {
+  type Message,
+  Repo,
+  type PeerId,
+  type DocumentId,
+  type HandleState,
+  type AutomergeUrl,
+} from '@dxos/automerge/automerge-repo';
 import { randomBytes } from '@dxos/crypto';
 import { PublicKey } from '@dxos/keys';
 import { createTestLevel } from '@dxos/kv-store/testing';
@@ -92,6 +99,74 @@ describe('AutomergeRepo', () => {
     expect(equals(a, c)).to.be.true;
   });
 
+  test('documents missing from local storage go to requesting state', async () => {
+    const hostAdapter: TestAdapter = new TestAdapter({
+      send: (message: Message) => {
+        console.log('hostAdapter.send', message);
+        clientAdapter.receive(message);
+      },
+    });
+    const clientAdapter: TestAdapter = new TestAdapter({
+      send: (message: Message) => {
+        console.log('clientAdapter.send', message);
+        if (message.type !== 'doc-unavailable' && message.type !== 'sync') {
+          hostAdapter.receive(message);
+        }
+      },
+    });
+
+    const host = new Repo({
+      network: [hostAdapter],
+    });
+    const _client = new Repo({
+      network: [clientAdapter],
+    });
+    hostAdapter.ready();
+    clientAdapter.ready();
+    await hostAdapter.onConnect.wait();
+    await clientAdapter.onConnect.wait();
+    hostAdapter.peerCandidate(clientAdapter.peerId!);
+    clientAdapter.peerCandidate(hostAdapter.peerId!);
+
+    const url = 'automerge:3JN8F3Z4dUWEEKKFN7WE9gEGvVUT';
+    const handle = host.find(url as AutomergeUrl);
+    await handle.whenReady(['requesting']);
+  });
+
+  test('documents on disk go to ready state', async () => {
+    const level = createTestLevel();
+    const storage = new LevelDBStorageAdapter({ db: level.sublevel('automerge') });
+    await openAndClose(level, storage);
+
+    let url: AutomergeUrl | undefined;
+    {
+      const repo = new Repo({
+        network: [],
+        storage,
+      });
+      const handle = repo.create<{ field?: string }>();
+      url = handle.url;
+      await repo.flush();
+    }
+
+    {
+      const repo = new Repo({
+        network: [],
+        storage,
+      });
+      const handle = repo.find(url as AutomergeUrl);
+
+      let requestingState = false;
+      queueMicrotask(async () => {
+        await handle.whenReady(['requesting']);
+        requestingState = true;
+      });
+
+      await handle.whenReady(['ready']);
+      expect(requestingState).to.be.false;
+    }
+  });
+
   describe('network', () => {
     test('basic networking', async () => {
       const hostAdapter: TestAdapter = new TestAdapter({

package/src/automerge/echo-network-adapter.ts

@@ -2,14 +2,14 @@
 // Copyright 2024 DXOS.org
 //
 
-import { Trigger, synchronized } from '@dxos/async';
+import { synchronized, Trigger } from '@dxos/async';
 import { type Message, NetworkAdapter, type PeerId, type PeerMetadata } from '@dxos/automerge/automerge-repo';
 import { LifecycleState } from '@dxos/context';
 import { invariant } from '@dxos/invariant';
 import { type PublicKey } from '@dxos/keys';
 import { log } from '@dxos/log';
 
-import { type EchoReplicator, type ReplicatorConnection, type ShouldAdvertizeParams } from './echo-replicator';
+import { type EchoReplicator, type ReplicatorConnection, type ShouldAdvertiseParams } from './echo-replicator';
 
 export type EchoNetworkAdapterParams = {
   getContainingSpaceForDocument: (documentId: string) => Promise<PublicKey | null>;
@@ -57,7 +57,9 @@ export class EchoNetworkAdapter extends NetworkAdapter {
 
   @synchronized
   async open() {
-    invariant(this._lifecycleState === LifecycleState.CLOSED);
+    if (this._lifecycleState === LifecycleState.OPEN) {
+      return;
+    }
     this._lifecycleState = LifecycleState.OPEN;
 
     log('emit ready');
@@ -68,7 +70,9 @@ export class EchoNetworkAdapter extends NetworkAdapter {
 
   @synchronized
   async close() {
-    invariant(this._lifecycleState === LifecycleState.OPEN);
+    if (this._lifecycleState === LifecycleState.CLOSED) {
+      return;
+    }
 
     for (const replicator of this._replicators) {
       await replicator.disconnect();
@@ -106,13 +110,13 @@ export class EchoNetworkAdapter extends NetworkAdapter {
     this._replicators.delete(replicator);
   }
 
-  async shouldAdvertize(peerId: PeerId, params: ShouldAdvertizeParams): Promise<boolean> {
+  async shouldAdvertize(peerId: PeerId, params: ShouldAdvertiseParams): Promise<boolean> {
     const connection = this._connections.get(peerId);
     if (!connection) {
       return false;
     }
 
-    return connection.connection.shouldAdvertize(params);
+    return connection.connection.shouldAdvertise(params);
   }
 
   private _onConnectionOpen(connection: ReplicatorConnection) {

package/src/automerge/echo-replicator.ts

@@ -23,13 +23,11 @@ export interface EchoReplicatorContext {
    */
   get peerId(): string;
 
-  onConnectionOpen(connection: ReplicatorConnection): void;
+  getContainingSpaceForDocument(documentId: string): Promise<PublicKey | null>;
 
+  onConnectionOpen(connection: ReplicatorConnection): void;
   onConnectionClosed(connection: ReplicatorConnection): void;
-
   onConnectionAuthScopeChanged(connection: ReplicatorConnection): void;
-
-  getContainingSpaceForDocument(documentId: string): Promise<PublicKey | null>;
 }
 
 export interface ReplicatorConnection {
@@ -49,13 +47,12 @@ export interface ReplicatorConnection {
   writable: WritableStream<Message>;
 
   /**
-   * @returns true if the document should be advertized to this peer.
-   *
-   * The remote peer can still request the document by it's id bypassing this check.
+   * @returns true if the document should be advertised to this peer.
+   * The remote peer can still request the document by its id bypassing this check.
    */
-  shouldAdvertize(params: ShouldAdvertizeParams): Promise<boolean>;
+  shouldAdvertise(params: ShouldAdvertiseParams): Promise<boolean>;
 }
 
-export type ShouldAdvertizeParams = {
+export type ShouldAdvertiseParams = {
   documentId: string;
 };

package/src/automerge/heads-store.ts

@@ -0,0 +1,39 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import type { Heads } from '@dxos/automerge/automerge';
+import type { DocumentId } from '@dxos/automerge/automerge-repo';
+import { headsEncoding } from '@dxos/indexing';
+import type { BatchLevel, SublevelDB } from '@dxos/kv-store';
+
+export type HeadsStoreParams = {
+  db: SublevelDB;
+};
+
+export class HeadsStore {
+  private readonly _db: SublevelDB;
+
+  constructor({ db }: HeadsStoreParams) {
+    this._db = db;
+  }
+
+  setHeads(documentId: DocumentId, heads: Heads, batch: BatchLevel) {
+    batch.put<DocumentId, Heads>(documentId, heads, {
+      sublevel: this._db,
+      keyEncoding: 'utf8',
+      valueEncoding: headsEncoding,
+    });
+  }
+
+  async getHeads(documentId: DocumentId): Promise<Heads | undefined> {
+    try {
+      return await this._db.get<DocumentId, Heads>(documentId, { keyEncoding: 'utf8', valueEncoding: headsEncoding });
+    } catch (err: any) {
+      if (err.notFound) {
+        return undefined;
+      }
+      throw err;
+    }
+  }
+}

package/src/automerge/mesh-echo-replicator.ts

@@ -2,7 +2,7 @@
 // Copyright 2024 DXOS.org
 //
 
-import { type Message, cbor } from '@dxos/automerge/automerge-repo';
+import { cbor, type Message } from '@dxos/automerge/automerge-repo';
 import { Resource } from '@dxos/context';
 import { invariant } from '@dxos/invariant';
 import { PublicKey } from '@dxos/keys';
@@ -14,7 +14,7 @@ import {
   type EchoReplicator,
   type EchoReplicatorContext,
   type ReplicatorConnection,
-  type ShouldAdvertizeParams,
+  type ShouldAdvertiseParams,
 } from './echo-replicator';
 
 // TODO(dmaretskyi): Move out of @dxos/echo-pipeline.
@@ -74,7 +74,7 @@ export class MeshEchoReplicator implements EchoReplicator {
         await connection.disable();
         this._connections.delete(connection);
       },
-      shouldAdvertize: async (params: ShouldAdvertizeParams) => {
+      shouldAdvertize: async (params: ShouldAdvertiseParams) => {
         log('shouldAdvertize', { peerId: connection.peerId, documentId: params.documentId });
         invariant(this._context);
         try {
@@ -133,7 +133,7 @@ type MeshReplicatorConnectionParams = {
   ownPeerId: string;
   onRemoteConnected: () => Promise<void>;
   onRemoteDisconnected: () => Promise<void>;
-  shouldAdvertize: (params: ShouldAdvertizeParams) => Promise<boolean>;
+  shouldAdvertize: (params: ShouldAdvertiseParams) => Promise<boolean>;
 };
 
 class MeshReplicatorConnection extends Resource implements ReplicatorConnection {
@@ -216,7 +216,7 @@ class MeshReplicatorConnection extends Resource implements ReplicatorConnection
     return this._remotePeerId;
   }
 
-  async shouldAdvertize(params: ShouldAdvertizeParams): Promise<boolean> {
+  async shouldAdvertise(params: ShouldAdvertiseParams): Promise<boolean> {
     return this._params.shouldAdvertize(params);
   }
 

package/src/db-host/data-service.ts

@@ -2,20 +2,23 @@
 // Copyright 2021 DXOS.org
 //
 
-import { type Stream } from '@dxos/codec-protobuf';
+import { type RequestOptions, type Stream } from '@dxos/codec-protobuf';
 import {
   type DataService,
   type EchoEvent,
   type FlushRequest,
+  type GetDocumentHeadsRequest,
+  type GetDocumentHeadsResponse,
   type HostInfo,
   type MutationReceipt,
+  type ReIndexHeadsRequest,
   type SubscribeRequest,
   type SyncRepoRequest,
   type SyncRepoResponse,
   type WriteRequest,
 } from '@dxos/protocols/proto/dxos/echo/service';
 
-import { type AutomergeHost } from '../automerge';
+import { type AutomergeHost, type DocumentId } from '../automerge';
 
 /**
  * Data sync between client and services.
@@ -49,4 +52,21 @@ export class DataServiceImpl implements DataService {
   sendSyncMessage(request: SyncRepoRequest): Promise<void> {
     return this._automergeHost.sendSyncMessage(request);
   }
+
+  async getDocumentHeads(request: GetDocumentHeadsRequest): Promise<GetDocumentHeadsResponse> {
+    const states = await Promise.all(
+      request.documentIds?.map(async (documentId): Promise<GetDocumentHeadsResponse.DocState> => {
+        const heads = await this._automergeHost.getHeads(documentId as DocumentId);
+        return {
+          documentId,
+          heads,
+        };
+      }) ?? [],
+    );
+    return { states };
+  }
+
+  async reIndexHeads(request: ReIndexHeadsRequest, options?: RequestOptions): Promise<void> {
+    await this._automergeHost.reIndexHeads((request.documentIds ?? []) as DocumentId[]);
+  }
 }