@dxos/echo-pipeline 0.6.2 → 0.6.3-main.0308ae2

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

@@ -3,16 +3,30 @@
 //
 
 import { synchronized, Trigger } from '@dxos/async';
-import { type Message, NetworkAdapter, type PeerId, type PeerMetadata } from '@dxos/automerge/automerge-repo';
+import { NetworkAdapter, type Message, 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 ShouldAdvertiseParams } from './echo-replicator';
+import { nonNullable } from '@dxos/util';
+
+import {
+  type EchoReplicator,
+  type ReplicatorConnection,
+  type ShouldAdvertiseParams,
+  type ShouldSyncCollectionParams,
+} from './echo-replicator';
+import {
+  isCollectionQueryMessage,
+  isCollectionStateMessage,
+  type CollectionQueryMessage,
+  type CollectionStateMessage,
+} from './network-protocol';
 
 export type EchoNetworkAdapterParams = {
   getContainingSpaceForDocument: (documentId: string) => Promise<PublicKey | null>;
+  onCollectionStateQueried: (collectionId: string, peerId: PeerId) => void;
+  onCollectionStateReceived: (collectionId: string, peerId: PeerId, state: unknown) => void;
 };
 
 /**
@@ -119,6 +133,47 @@ export class EchoNetworkAdapter extends NetworkAdapter {
     return connection.connection.shouldAdvertise(params);
   }
 
+  shouldSyncCollection(peerId: PeerId, params: ShouldSyncCollectionParams): boolean {
+    const connection = this._connections.get(peerId);
+    if (!connection) {
+      return false;
+    }
+
+    return connection.connection.shouldSyncCollection(params);
+  }
+
+  queryCollectionState(collectionId: string, targetId: PeerId): void {
+    const message: CollectionQueryMessage = {
+      type: 'collection-query',
+      senderId: this.peerId as PeerId,
+      targetId,
+      collectionId,
+    };
+    this.send(message);
+  }
+
+  sendCollectionState(collectionId: string, targetId: PeerId, state: unknown): void {
+    const message: CollectionStateMessage = {
+      type: 'collection-state',
+      senderId: this.peerId as PeerId,
+      targetId,
+      collectionId,
+      state,
+    };
+    this.send(message);
+  }
+
+  // TODO(dmaretskyi): Remove.
+  getPeersInterestedInCollection(collectionId: string): PeerId[] {
+    return Array.from(this._connections.values())
+      .map((connection) => {
+        return connection.connection.shouldSyncCollection({ collectionId })
+          ? (connection.connection.peerId as PeerId)
+          : null;
+      })
+      .filter(nonNullable);
+  }
+
   private _onConnectionOpen(connection: ReplicatorConnection) {
     log('Connection opened', { peerId: connection.peerId });
     invariant(!this._connections.has(connection.peerId as PeerId));
@@ -136,7 +191,7 @@ export class EchoNetworkAdapter extends NetworkAdapter {
             break;
           }
 
-          this.emit('message', value);
+          this._onMessage(value);
         }
       } catch (err) {
         if (connectionEntry.isOpen) {
@@ -149,6 +204,16 @@ export class EchoNetworkAdapter extends NetworkAdapter {
     this._emitPeerCandidate(connection);
   }
 
+  private _onMessage(message: Message) {
+    if (isCollectionQueryMessage(message)) {
+      this._params.onCollectionStateQueried(message.collectionId, message.senderId);
+    } else if (isCollectionStateMessage(message)) {
+      this._params.onCollectionStateReceived(message.collectionId, message.senderId, message.state);
+    } else {
+      this.emit('message', message);
+    }
+  }
+
   /**
    * Trigger doc-synchronizer shared documents set recalculation. Happens on peer-candidate.
    * TODO(y): replace with a proper API call when sharePolicy update becomes supported by automerge-repo

package/src/automerge/echo-replicator.ts

@@ -51,8 +51,17 @@ export interface ReplicatorConnection {
    * The remote peer can still request the document by its id bypassing this check.
    */
   shouldAdvertise(params: ShouldAdvertiseParams): Promise<boolean>;
+
+  /**
+   * @returns true if the collection should be synced to this peer.
+   */
+  shouldSyncCollection(params: ShouldSyncCollectionParams): boolean;
 }
 
 export type ShouldAdvertiseParams = {
   documentId: string;
 };
+
+export type ShouldSyncCollectionParams = {
+  collectionId: string;
+};

package/src/automerge/heads-store.ts

@@ -26,14 +26,11 @@ export class HeadsStore {
     });
   }
 
-  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;
-    }
+  // TODO(dmaretskyi): Make batched.
+  async getHeads(documentIds: DocumentId[]): Promise<Array<Heads | undefined>> {
+    return this._db.getMany<DocumentId, Heads>(documentIds, {
+      keyEncoding: 'utf8',
+      valueEncoding: headsEncoding,
+    });
   }
 }

package/src/automerge/index.ts

@@ -6,3 +6,5 @@ export * from './automerge-host';
 export * from './leveldb-storage-adapter';
 export * from './mesh-echo-replicator';
 export * from './echo-replicator';
+export { diffCollectionState } from './collection-synchronizer';
+export * from './space-collection';

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

@@ -9,7 +9,7 @@ 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';
+import type { ReplicatorConnection, ShouldAdvertiseParams, ShouldSyncCollectionParams } from './echo-replicator';
 
 const DEFAULT_FACTORY: AutomergeReplicatorFactory = (params) => new AutomergeReplicator(...params);
 
@@ -18,6 +18,7 @@ export type MeshReplicatorConnectionParams = {
   onRemoteConnected: () => void;
   onRemoteDisconnected: () => void;
   shouldAdvertise: (params: ShouldAdvertiseParams) => Promise<boolean>;
+  shouldSyncCollection: (params: ShouldSyncCollectionParams) => boolean;
   replicatorFactory?: AutomergeReplicatorFactory;
 };
 
@@ -112,6 +113,10 @@ export class MeshReplicatorConnection extends Resource implements ReplicatorConn
     return this._params.shouldAdvertise(params);
   }
 
+  shouldSyncCollection(params: ShouldSyncCollectionParams): boolean {
+    return this._params.shouldSyncCollection(params);
+  }
+
   /**
    * Start exchanging messages with the remote peer.
    * Call after the remote peer has connected.

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

@@ -3,16 +3,18 @@
 //
 
 import { invariant } from '@dxos/invariant';
-import { PublicKey } from '@dxos/keys';
+import { PublicKey, type SpaceId } from '@dxos/keys';
 import { log } from '@dxos/log';
 import {
   type AutomergeReplicator,
   type AutomergeReplicatorFactory,
 } from '@dxos/teleport-extension-automerge-replicator';
-import { ComplexMap, ComplexSet, defaultMap } from '@dxos/util';
+import { ComplexSet, defaultMap } from '@dxos/util';
 
 import { type EchoReplicator, type EchoReplicatorContext, type ShouldAdvertiseParams } from './echo-replicator';
 import { MeshReplicatorConnection } from './mesh-echo-replicator-connection';
+import { getSpaceIdFromCollectionId } from './space-collection';
+import { createIdFromSpaceKey } from '../space';
 
 // TODO(dmaretskyi): Move out of @dxos/echo-pipeline.
 
@@ -27,9 +29,9 @@ export class MeshEchoReplicator implements EchoReplicator {
   private readonly _connectionsPerPeer = new Map<string, MeshReplicatorConnection>();
 
   /**
-   * spaceKey -> deviceKey[]
+   * spaceId -> deviceKey[]
    */
-  private readonly _authorizedDevices = new ComplexMap<PublicKey, ComplexSet<PublicKey>>(PublicKey.hash);
+  private readonly _authorizedDevices = new Map<SpaceId, ComplexSet<PublicKey>>();
 
   private _context: EchoReplicatorContext | null = null;
 
@@ -88,7 +90,9 @@ export class MeshEchoReplicator implements EchoReplicator {
             return false;
           }
 
-          const authorizedDevices = this._authorizedDevices.get(spaceKey);
+          const spaceId = await createIdFromSpaceKey(spaceKey);
+
+          const authorizedDevices = this._authorizedDevices.get(spaceId);
 
           if (!connection.remoteDeviceKey) {
             log('device key not found for share policy check', {
@@ -113,15 +117,32 @@ export class MeshEchoReplicator implements EchoReplicator {
           return false;
         }
       },
+      shouldSyncCollection: ({ collectionId }) => {
+        const spaceId = getSpaceIdFromCollectionId(collectionId);
+
+        const authorizedDevices = this._authorizedDevices.get(spaceId);
+
+        if (!connection.remoteDeviceKey) {
+          log('device key not found for collection sync check', {
+            peerId: connection.peerId,
+            collectionId,
+          });
+          return false;
+        }
+
+        const isAuthorized = authorizedDevices?.has(connection.remoteDeviceKey) ?? false;
+        return isAuthorized;
+      },
     });
     this._connections.add(connection);
 
     return connection.replicatorExtension;
   }
 
-  authorizeDevice(spaceKey: PublicKey, deviceKey: PublicKey) {
+  async authorizeDevice(spaceKey: PublicKey, deviceKey: PublicKey) {
     log('authorizeDevice', { spaceKey, deviceKey });
-    defaultMap(this._authorizedDevices, spaceKey, () => new ComplexSet(PublicKey.hash)).add(deviceKey);
+    const spaceId = await createIdFromSpaceKey(spaceKey);
+    defaultMap(this._authorizedDevices, spaceId, () => new ComplexSet(PublicKey.hash)).add(deviceKey);
     for (const connection of this._connections) {
       if (connection.remoteDeviceKey && connection.remoteDeviceKey.equals(deviceKey)) {
         if (this._connectionsPerPeer.has(connection.peerId)) {

package/src/automerge/network-protocol.ts

@@ -0,0 +1,45 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import type { Message, PeerId } from '@dxos/automerge/automerge-repo';
+
+export const MESSAGE_TYPE_COLLECTION_QUERY = 'collection-query';
+
+export type CollectionQueryMessage = {
+  type: typeof MESSAGE_TYPE_COLLECTION_QUERY;
+  senderId: PeerId;
+  targetId: PeerId;
+  collectionId: string;
+
+  /**
+   * Identifier of the current state.
+   * Remote peer will skip sending the state if it has the same tag.
+   */
+  stateTag?: string;
+};
+
+export const isCollectionQueryMessage = (message: Message): message is CollectionQueryMessage =>
+  message.type === MESSAGE_TYPE_COLLECTION_QUERY;
+
+export const MESSAGE_TYPE_COLLECTION_STATE = 'collection-state';
+
+export type CollectionStateMessage = {
+  type: typeof MESSAGE_TYPE_COLLECTION_STATE;
+  senderId: PeerId;
+  targetId: PeerId;
+  collectionId: string;
+
+  /**
+   *  State representation is implementation-defined.
+   */
+  state?: unknown;
+
+  /**
+   * Identifier of the current state.
+   */
+  stateTag?: string;
+};
+
+export const isCollectionStateMessage = (message: Message): message is CollectionStateMessage =>
+  message.type === MESSAGE_TYPE_COLLECTION_STATE;

package/src/automerge/space-collection.ts

@@ -0,0 +1,14 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { invariant } from '@dxos/invariant';
+import { SpaceId } from '@dxos/keys';
+
+export const deriveCollectionIdFromSpaceId = (spaceId: SpaceId): string => `space:${spaceId}`;
+
+export const getSpaceIdFromCollectionId = (collectionId: string): SpaceId => {
+  const spaceId = collectionId.replace(/^space:/, '');
+  invariant(SpaceId.isValid(spaceId));
+  return spaceId;
+};

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

@@ -5,10 +5,10 @@
 import { type DocumentId } from '@dxos/automerge/automerge-repo';
 import { type RequestOptions, Stream } from '@dxos/codec-protobuf';
 import { invariant } from '@dxos/invariant';
+import { SpaceId } from '@dxos/keys';
 import { log } from '@dxos/log';
 import {
   type DataService,
-  type DocHeadsList,
   type FlushRequest,
   type SubscribeRequest,
   type BatchedDocumentUpdates,
@@ -18,10 +18,12 @@ import {
   type ReIndexHeadsRequest,
   type WaitUntilHeadsReplicatedRequest,
   type UpdateRequest,
+  type GetSpaceSyncStateRequest,
+  type SpaceSyncState,
 } from '@dxos/protocols/proto/dxos/echo/service';
 
 import { DocumentsSynchronizer } from './documents-synchronizer';
-import { type AutomergeHost } from '../automerge';
+import { deriveCollectionIdFromSpaceId, type AutomergeHost } from '../automerge';
 
 export type DataServiceParams = {
   automergeHost: AutomergeHost;
@@ -91,18 +93,14 @@ export class DataServiceImpl implements DataService {
   }
 
   async getDocumentHeads(request: GetDocumentHeadsRequest): Promise<GetDocumentHeadsResponse> {
-    const entries = await Promise.all(
-      request.documentIds?.map(async (documentId): Promise<DocHeadsList.Entry> => {
-        const heads = await this._automergeHost.getHeads(documentId as DocumentId);
-        return {
-          documentId,
-          heads,
-        };
-      }) ?? [],
-    );
+    const documentIds = request.documentIds;
+    if (!documentIds) {
+      return { heads: { entries: [] } };
+    }
+    const heads = await this._automergeHost.getHeads(documentIds as DocumentId[]);
     return {
       heads: {
-        entries,
+        entries: heads.map((heads, idx) => ({ documentId: documentIds[idx], heads })),
       },
     };
   }
@@ -121,4 +119,20 @@ export class DataServiceImpl implements DataService {
   async updateIndexes() {
     await this._updateIndexes();
   }
+
+  async getSpaceSyncState(
+    request: GetSpaceSyncStateRequest,
+    options?: RequestOptions | undefined,
+  ): Promise<SpaceSyncState> {
+    invariant(SpaceId.isValid(request.spaceId));
+    const collectionId = deriveCollectionIdFromSpaceId(request.spaceId);
+    const state = await this._automergeHost.getCollectionSyncState(collectionId);
+
+    return {
+      peers: state.peers.map((peer) => ({
+        peerId: peer.peerId,
+        documentsToReconcile: peer.differentDocuments,
+      })),
+    };
+  }
 }

package/src/db-host/documents-synchronizer.ts

@@ -44,14 +44,26 @@ export class DocumentsSynchronizer extends Resource {
     super();
   }
 
-  async addDocuments(documentIds: DocumentId[]) {
+  addDocuments(documentIds: DocumentId[], retryCounter = 0) {
+    if (retryCounter > 3) {
+      log.warn('Failed to load document, retry limit reached', { documentIds });
+      return;
+    }
+
     for (const documentId of documentIds) {
       const doc = this._params.repo.find(documentId as DocumentId);
-      await doc.whenReady();
-      this._startSync(doc);
-      this._pendingUpdates.add(doc.documentId);
+      doc
+        .whenReady()
+        .then(() => {
+          this._startSync(doc);
+          this._pendingUpdates.add(doc.documentId);
+          this._sendUpdatesJob!.trigger();
+        })
+        .catch((error) => {
+          log.warn('Failed to load document, wraparound', { documentId, error });
+          this.addDocuments([documentId], retryCounter + 1);
+        });
     }
-    this._sendUpdatesJob!.trigger();
   }
 
   removeDocuments(documentIds: DocumentId[]) {

package/src/metadata/metadata-store.ts

@@ -150,7 +150,7 @@ export class MetadataStore {
 
       // post-processing
       this._metadata.spaces?.forEach((space) => {
-        space.state ??= SpaceState.ACTIVE;
+        space.state ??= SpaceState.SPACE_ACTIVE;
       });
     } catch (err: any) {
       log.error('failed to load metadata', { err });

package/src/testing/index.ts

@@ -6,3 +6,4 @@ export * from './change-metadata';
 export * from './test-agent-builder';
 export * from './test-feed-builder';
 export * from './test-network-adapter';
+export * from './test-replicator';

package/src/testing/test-replicator.ts

@@ -0,0 +1,194 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { sleep, synchronized } from '@dxos/async';
+import { type Message } from '@dxos/automerge/automerge-repo';
+import { type Context, LifecycleState, Resource } from '@dxos/context';
+import { invariant } from '@dxos/invariant';
+import { log } from '@dxos/log';
+import { AutomergeReplicator, type AutomergeReplicatorFactory } from '@dxos/teleport-extension-automerge-replicator';
+
+import type {
+  EchoReplicator,
+  EchoReplicatorContext,
+  ReplicatorConnection,
+  ShouldAdvertiseParams,
+  ShouldSyncCollectionParams,
+} from '../automerge';
+
+export type TestReplicatorNetworkOptions = {
+  latency?: number;
+};
+
+export class TestReplicationNetwork extends Resource {
+  private readonly _replicators = new Set<TestReplicator>();
+  private readonly _latency?: number = undefined;
+
+  constructor(options: TestReplicatorNetworkOptions = {}) {
+    super();
+    this._latency = options.latency;
+  }
+
+  protected override async _close(ctx: Context): Promise<void> {
+    for (const replicator of this._replicators) {
+      for (const connection of replicator.connections) {
+        void connection.writable.abort();
+        void connection.readable.cancel();
+      }
+    }
+  }
+
+  async createReplicator(): Promise<TestReplicator> {
+    const replicator = new TestReplicator({
+      onConnect: async () => {
+        invariant(this._lifecycleState === LifecycleState.OPEN);
+        await this._connectReplicator(replicator);
+      },
+      onDisconnect: async () => {
+        invariant(this._lifecycleState === LifecycleState.OPEN);
+        await this._disconnectReplicator(replicator);
+      },
+    });
+    this._replicators.add(replicator);
+    return replicator;
+  }
+
+  @synchronized
+  private async _connectReplicator(replicator: TestReplicator) {
+    for (const otherReplicator of this._replicators.values()) {
+      if (otherReplicator === replicator || !otherReplicator.connected) {
+        continue;
+      }
+      log('create connection', { from: replicator.context!.peerId, to: otherReplicator.context!.peerId });
+      const [connection1, connection2] = this._createConnectionPair(
+        replicator.context!.peerId,
+        otherReplicator.context!.peerId,
+      );
+      await replicator.context!.onConnectionOpen(connection1);
+      await otherReplicator.context!.onConnectionOpen(connection2);
+    }
+  }
+
+  private async _disconnectReplicator(replicator: TestReplicator) {
+    for (const connection of replicator.connections) {
+      await replicator.context!.onConnectionClosed(connection);
+      await connection.otherSide!.owningReplicator!.removeConnection(connection.otherSide!);
+    }
+  }
+
+  private _createConnectionPair(peer1: string, peer2: string): [TestReplicatorConnection, TestReplicatorConnection] {
+    const LOG = false;
+
+    const forward = new TransformStream({
+      transform: async (message, controller) => {
+        if (LOG) {
+          log.info('replicate', { from: peer1, to: peer2, message });
+        }
+
+        if (this._latency !== undefined) {
+          await sleep(this._latency);
+        }
+
+        controller.enqueue(message);
+      },
+    });
+    const backwards = new TransformStream({
+      transform: async (message, controller) => {
+        if (LOG) {
+          log.info('replicate', { from: peer2, to: peer1, message });
+        }
+
+        if (this._latency !== undefined) {
+          await sleep(this._latency);
+        }
+
+        controller.enqueue(message);
+      },
+    });
+
+    const connection1 = new TestReplicatorConnection(peer2, backwards.readable, forward.writable);
+    const connection2 = new TestReplicatorConnection(peer1, forward.readable, backwards.writable);
+    connection1.otherSide = connection2;
+    connection2.otherSide = connection1;
+    return [connection1, connection2];
+  }
+}
+
+type TestReplicatorParams = {
+  onConnect: () => Promise<void>;
+  onDisconnect: () => Promise<void>;
+};
+
+export class TestReplicator implements EchoReplicator {
+  constructor(private readonly _params: TestReplicatorParams) {}
+
+  public connected = false;
+  public context: EchoReplicatorContext | undefined = undefined;
+  public connections = new Set<TestReplicatorConnection>();
+
+  async connect(context: EchoReplicatorContext): Promise<void> {
+    log('connect', { peerId: context.peerId });
+    this.context = context;
+    this.connected = true;
+    await this._params.onConnect();
+  }
+
+  async disconnect(): Promise<void> {
+    log('disconnect', { peerId: this.context!.peerId });
+    this.connected = false;
+    await this._params.onDisconnect();
+  }
+
+  async addConnection(connection: TestReplicatorConnection): Promise<void> {
+    connection.owningReplicator = this;
+    this.connections.add(connection);
+    this.context!.onConnectionOpen(connection);
+  }
+
+  async removeConnection(connection: TestReplicatorConnection): Promise<void> {
+    connection.owningReplicator = undefined;
+    this.context!.onConnectionClosed(connection);
+    this.connections.delete(connection);
+  }
+}
+
+export class TestReplicatorConnection implements ReplicatorConnection {
+  public otherSide: TestReplicatorConnection | undefined = undefined;
+  public owningReplicator: TestReplicator | undefined = undefined;
+
+  constructor(
+    public readonly peerId: string,
+    public readonly readable: ReadableStream<Message>,
+    public readonly writable: WritableStream<Message>,
+  ) {}
+
+  async shouldAdvertise(params: ShouldAdvertiseParams): Promise<boolean> {
+    return true;
+  }
+
+  shouldSyncCollection(params: ShouldSyncCollectionParams): boolean {
+    return true;
+  }
+}
+
+export const testAutomergeReplicatorFactory: AutomergeReplicatorFactory = (params) => {
+  return new AutomergeReplicator(
+    {
+      ...params[0],
+      sendSyncRetryPolicy: {
+        retryBackoff: 20,
+        retriesBeforeBackoff: 2,
+        maxRetries: 3,
+      },
+    },
+    params[1],
+  );
+};
+
+export const brokenAutomergeReplicatorFactory: AutomergeReplicatorFactory = (params) => {
+  params[1]!.onSyncMessage = () => {
+    throw new Error();
+  };
+  return testAutomergeReplicatorFactory(params);
+};