@dxos/echo-pipeline 0.6.0 → 0.6.1-main.09a92b0

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-connection.ts

@@ -0,0 +1,130 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { cbor, type Message } from '@dxos/automerge/automerge-repo';
+import { Resource } from '@dxos/context';
+import { invariant } from '@dxos/invariant';
+import { type PublicKey } from '@dxos/keys';
+import { log } from '@dxos/log';
+import { AutomergeReplicator, type AutomergeReplicatorFactory } from '@dxos/teleport-extension-automerge-replicator';
+
+import type { ReplicatorConnection, ShouldAdvertiseParams } from './echo-replicator';
+
+const DEFAULT_FACTORY: AutomergeReplicatorFactory = (params) => new AutomergeReplicator(...params);
+
+export type MeshReplicatorConnectionParams = {
+  ownPeerId: string;
+  onRemoteConnected: () => void;
+  onRemoteDisconnected: () => void;
+  shouldAdvertise: (params: ShouldAdvertiseParams) => Promise<boolean>;
+  replicatorFactory?: AutomergeReplicatorFactory;
+};
+
+export class MeshReplicatorConnection extends Resource implements ReplicatorConnection {
+  public readable: ReadableStream<Message>;
+  public writable: WritableStream<Message>;
+  public remoteDeviceKey: PublicKey | null = null;
+
+  public readonly replicatorExtension: AutomergeReplicator;
+
+  private _remotePeerId: string | null = null;
+  private _isEnabled = false;
+
+  constructor(private readonly _params: MeshReplicatorConnectionParams) {
+    super();
+
+    let readableStreamController!: ReadableStreamDefaultController<Message>;
+    this.readable = new ReadableStream<Message>({
+      start: (controller) => {
+        readableStreamController = controller;
+        this._ctx.onDispose(() => controller.close());
+      },
+    });
+
+    this.writable = new WritableStream<Message>({
+      write: async (message: Message, controller) => {
+        invariant(this._isEnabled, 'Writing to a disabled connection');
+        try {
+          await this.replicatorExtension.sendSyncMessage({ payload: cbor.encode(message) });
+        } catch (err) {
+          controller.error(err);
+          this._disconnectIfEnabled();
+        }
+      },
+    });
+
+    const createAutomergeReplicator = this._params.replicatorFactory ?? DEFAULT_FACTORY;
+    this.replicatorExtension = createAutomergeReplicator([
+      {
+        peerId: this._params.ownPeerId,
+      },
+      {
+        onStartReplication: async (info, remotePeerId /** Teleport ID */) => {
+          // Note: We store only one extension per peer.
+          //       There can be a case where two connected peers have more than one teleport connection between them
+          //       and each of them uses different teleport connections to send messages.
+          //       It works because we receive messages from all teleport connections and Automerge Repo dedup them.
+          // TODO(mykola): Use only one teleport connection per peer.
+
+          // TODO(dmaretskyi): Critical bug.
+          // - two peers get connected via swarm 1
+          // - they get connected via swarm 2
+          // - swarm 1 gets disconnected
+          // - automerge repo thinks that peer 2 got disconnected even though swarm 2 is still active
+
+          this.remoteDeviceKey = remotePeerId;
+
+          // Set automerge id.
+          this._remotePeerId = info.id;
+
+          log('onStartReplication', { id: info.id, thisPeerId: this.peerId, remotePeerId: remotePeerId.toHex() });
+
+          this._params.onRemoteConnected();
+        },
+        onSyncMessage: async ({ payload }) => {
+          if (!this._isEnabled) {
+            return;
+          }
+          const message = cbor.decode(payload) as Message;
+          // Note: automerge Repo dedup messages.
+          readableStreamController.enqueue(message);
+        },
+        onClose: async () => {
+          this._disconnectIfEnabled();
+        },
+      },
+    ]);
+  }
+
+  private _disconnectIfEnabled() {
+    if (this._isEnabled) {
+      this._params.onRemoteDisconnected();
+    }
+  }
+
+  get peerId(): string {
+    invariant(this._remotePeerId != null, 'Remote peer has not connected yet.');
+    return this._remotePeerId;
+  }
+
+  async shouldAdvertise(params: ShouldAdvertiseParams): Promise<boolean> {
+    return this._params.shouldAdvertise(params);
+  }
+
+  /**
+   * Start exchanging messages with the remote peer.
+   * Call after the remote peer has connected.
+   */
+  enable() {
+    invariant(this._remotePeerId != null, 'Remote peer has not connected yet.');
+    this._isEnabled = true;
+  }
+
+  /**
+   * Stop exchanging messages with the remote peer.
+   */
+  disable() {
+    this._isEnabled = false;
+  }
+}

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

@@ -2,20 +2,17 @@
 // Copyright 2024 DXOS.org
 //
 
-import { type Message, cbor } from '@dxos/automerge/automerge-repo';
-import { Resource } from '@dxos/context';
 import { invariant } from '@dxos/invariant';
 import { PublicKey } from '@dxos/keys';
 import { log } from '@dxos/log';
-import { AutomergeReplicator } from '@dxos/teleport-extension-automerge-replicator';
+import {
+  type AutomergeReplicator,
+  type AutomergeReplicatorFactory,
+} from '@dxos/teleport-extension-automerge-replicator';
 import { ComplexMap, ComplexSet, defaultMap } from '@dxos/util';
 
-import {
-  type EchoReplicator,
-  type EchoReplicatorContext,
-  type ReplicatorConnection,
-  type ShouldAdvertizeParams,
-} from './echo-replicator';
+import { type EchoReplicator, type EchoReplicatorContext, type ShouldAdvertiseParams } from './echo-replicator';
+import { MeshReplicatorConnection } from './mesh-echo-replicator-connection';
 
 // TODO(dmaretskyi): Move out of @dxos/echo-pipeline.
 
@@ -50,11 +47,12 @@ export class MeshEchoReplicator implements EchoReplicator {
     this._context = null;
   }
 
-  createExtension(): AutomergeReplicator {
+  createExtension(extensionFactory?: AutomergeReplicatorFactory): AutomergeReplicator {
     invariant(this._context);
 
     const connection: MeshReplicatorConnection = new MeshReplicatorConnection({
       ownPeerId: this._context.peerId,
+      replicatorFactory: extensionFactory,
       onRemoteConnected: async () => {
         log('onRemoteConnected', { peerId: connection.peerId });
         invariant(this._context);
@@ -64,18 +62,18 @@ export class MeshEchoReplicator implements EchoReplicator {
         } else {
           this._connectionsPerPeer.set(connection.peerId, connection);
           this._context.onConnectionOpen(connection);
-          await connection.enable();
+          connection.enable();
         }
       },
       onRemoteDisconnected: async () => {
         log('onRemoteDisconnected', { peerId: connection.peerId });
         this._context?.onConnectionClosed(connection);
         this._connectionsPerPeer.delete(connection.peerId);
-        await connection.disable();
+        connection.disable();
         this._connections.delete(connection);
       },
-      shouldAdvertize: async (params: ShouldAdvertizeParams) => {
-        log('shouldAdvertize', { peerId: connection.peerId, documentId: params.documentId });
+      shouldAdvertise: async (params: ShouldAdvertiseParams) => {
+        log('shouldAdvertise', { peerId: connection.peerId, documentId: params.documentId });
         invariant(this._context);
         try {
           const spaceKey = await this._context.getContainingSpaceForDocument(params.documentId);
@@ -123,116 +121,10 @@ export class MeshEchoReplicator implements EchoReplicator {
     defaultMap(this._authorizedDevices, spaceKey, () => new ComplexSet(PublicKey.hash)).add(deviceKey);
     for (const connection of this._connections) {
       if (connection.remoteDeviceKey && connection.remoteDeviceKey.equals(deviceKey)) {
-        this._context?.onConnectionAuthScopeChanged(connection);
+        if (this._connectionsPerPeer.has(connection.peerId)) {
+          this._context?.onConnectionAuthScopeChanged(connection);
+        }
       }
     }
   }
 }
-
-type MeshReplicatorConnectionParams = {
-  ownPeerId: string;
-  onRemoteConnected: () => Promise<void>;
-  onRemoteDisconnected: () => Promise<void>;
-  shouldAdvertize: (params: ShouldAdvertizeParams) => Promise<boolean>;
-};
-
-class MeshReplicatorConnection extends Resource implements ReplicatorConnection {
-  public readable: ReadableStream<Message>;
-  public writable: WritableStream<Message>;
-  public remoteDeviceKey: PublicKey | null = null;
-
-  public readonly replicatorExtension: AutomergeReplicator;
-
-  private _remotePeerId: string | null = null;
-  private _isEnabled = false;
-
-  constructor(private readonly _params: MeshReplicatorConnectionParams) {
-    super();
-
-    let readableStreamController!: ReadableStreamDefaultController<Message>;
-    this.readable = new ReadableStream<Message>({
-      start: (controller) => {
-        readableStreamController = controller;
-        this._ctx.onDispose(() => controller.close());
-      },
-    });
-
-    this.writable = new WritableStream<Message>({
-      write: async (message: Message, controller) => {
-        // TODO(dmaretskyi): Show we block on RPC completing here?
-        this.replicatorExtension.sendSyncMessage({ payload: cbor.encode(message) }).catch((err) => {
-          controller.error(err);
-        });
-      },
-    });
-
-    this.replicatorExtension = new AutomergeReplicator(
-      {
-        peerId: this._params.ownPeerId,
-      },
-      {
-        onStartReplication: async (info, remotePeerId /** Teleport ID */) => {
-          // Note: We store only one extension per peer.
-          //       There can be a case where two connected peers have more than one teleport connection between them
-          //       and each of them uses different teleport connections to send messages.
-          //       It works because we receive messages from all teleport connections and Automerge Repo dedup them.
-          // TODO(mykola): Use only one teleport connection per peer.
-
-          // TODO(dmaretskyi): Critical bug.
-          // - two peers get connected via swarm 1
-          // - they get connected via swarm 2
-          // - swarm 1 gets disconnected
-          // - automerge repo thinks that peer 2 got disconnected even though swarm 2 is still active
-
-          this.remoteDeviceKey = remotePeerId;
-
-          // Set automerge id.
-          this._remotePeerId = info.id;
-
-          log('onStartReplication', { id: info.id, thisPeerId: this.peerId, remotePeerId: remotePeerId.toHex() });
-
-          await this._params.onRemoteConnected();
-        },
-        onSyncMessage: async ({ payload }) => {
-          if (!this._isEnabled) {
-            return;
-          }
-          const message = cbor.decode(payload) as Message;
-          // Note: automerge Repo dedup messages.
-          readableStreamController.enqueue(message);
-        },
-        onClose: async () => {
-          if (!this._isEnabled) {
-            return;
-          }
-          await this._params.onRemoteDisconnected();
-        },
-      },
-    );
-  }
-
-  get peerId(): string {
-    invariant(this._remotePeerId != null, 'Remote peer has not connected yet.');
-    return this._remotePeerId;
-  }
-
-  async shouldAdvertize(params: ShouldAdvertizeParams): Promise<boolean> {
-    return this._params.shouldAdvertize(params);
-  }
-
-  /**
-   * Start exchanging messages with the remote peer.
-   * Call after the remote peer has connected.
-   */
-  async enable() {
-    invariant(this._remotePeerId != null, 'Remote peer has not connected yet.');
-    this._isEnabled = true;
-  }
-
-  /**
-   * Stop exchanging messages with the remote peer.
-   */
-  async disable() {
-    this._isEnabled = false;
-  }
-}

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[]);
+  }
 }