@dxos/echo-pipeline 0.6.13 → 0.6.14-main.2b6a0f3

package/src/edge/echo-edge-replicator.ts

@@ -0,0 +1,341 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { Mutex, scheduleTask } from '@dxos/async';
+import * as A from '@dxos/automerge/automerge';
+import { cbor } from '@dxos/automerge/automerge-repo';
+import { Context, Resource } from '@dxos/context';
+import { randomUUID } from '@dxos/crypto';
+import type { CollectionId } from '@dxos/echo-protocol';
+import { type EdgeConnection } from '@dxos/edge-client';
+import { invariant } from '@dxos/invariant';
+import type { SpaceId } from '@dxos/keys';
+import { log } from '@dxos/log';
+import { EdgeService, type AutomergeProtocolMessage, type PeerId } from '@dxos/protocols';
+import { buf } from '@dxos/protocols/buf';
+import {
+  type Message as RouterMessage,
+  MessageSchema as RouterMessageSchema,
+} from '@dxos/protocols/buf/dxos/edge/messenger_pb';
+import { bufferToArray } from '@dxos/util';
+
+import {
+  getSpaceIdFromCollectionId,
+  type EchoReplicator,
+  type EchoReplicatorContext,
+  type ReplicatorConnection,
+  type ShouldAdvertiseParams,
+  type ShouldSyncCollectionParams,
+} from '../automerge';
+
+/**
+ * Delay before restarting the connection after receiving a forbidden error.
+ */
+const INITIAL_RESTART_DELAY = 500;
+const RESTART_DELAY_JITTER = 250;
+const MAX_RESTART_DELAY = 5000;
+
+export type EchoEdgeReplicatorParams = {
+  edgeConnection: EdgeConnection;
+  disableSharePolicy?: boolean;
+};
+
+export class EchoEdgeReplicator implements EchoReplicator {
+  private readonly _edgeConnection: EdgeConnection;
+  private readonly _mutex = new Mutex();
+
+  private _ctx?: Context = undefined;
+  private _context: EchoReplicatorContext | null = null;
+  private _connectedSpaces = new Set<SpaceId>();
+  private _connections = new Map<SpaceId, EdgeReplicatorConnection>();
+  private _sharePolicyEnabled = true;
+
+  constructor({ edgeConnection, disableSharePolicy }: EchoEdgeReplicatorParams) {
+    this._edgeConnection = edgeConnection;
+    this._sharePolicyEnabled = !disableSharePolicy;
+  }
+
+  async connect(context: EchoReplicatorContext): Promise<void> {
+    log.info('connect', { peerId: context.peerId, connectedSpaces: this._connectedSpaces.size });
+    this._context = context;
+
+    this._ctx = Context.default();
+    this._edgeConnection.reconnect.on(this._ctx, async () => {
+      using _guard = await this._mutex.acquire();
+
+      const spaces = [...this._connectedSpaces];
+      for (const connection of this._connections.values()) {
+        await connection.close();
+      }
+      this._connections.clear();
+
+      if (this._context !== null) {
+        for (const spaceId of spaces) {
+          await this._openConnection(spaceId);
+        }
+      }
+    });
+
+    for (const spaceId of this._connectedSpaces) {
+      await this._openConnection(spaceId);
+    }
+  }
+
+  async disconnect(): Promise<void> {
+    using _guard = await this._mutex.acquire();
+    await this._ctx?.dispose();
+
+    for (const connection of this._connections.values()) {
+      await connection.close();
+    }
+    this._connections.clear();
+  }
+
+  async connectToSpace(spaceId: SpaceId) {
+    using _guard = await this._mutex.acquire();
+
+    this._connectedSpaces.add(spaceId);
+
+    // Check if AM-repo requested that we connect to remote peers.
+    if (this._context !== null) {
+      await this._openConnection(spaceId);
+    }
+  }
+
+  async disconnectFromSpace(spaceId: SpaceId) {
+    using _guard = await this._mutex.acquire();
+
+    this._connectedSpaces.delete(spaceId);
+
+    const connection = this._connections.get(spaceId);
+    if (connection) {
+      await connection.close();
+      this._connections.delete(spaceId);
+    }
+  }
+
+  private async _openConnection(spaceId: SpaceId, reconnects: number = 0) {
+    invariant(this._context);
+    invariant(!this._connections.has(spaceId));
+
+    let restartScheduled = false;
+
+    const connection = new EdgeReplicatorConnection({
+      edgeConnection: this._edgeConnection,
+      spaceId,
+      context: this._context,
+      sharedPolicyEnabled: this._sharePolicyEnabled,
+      onRemoteConnected: async () => {
+        this._context?.onConnectionOpen(connection);
+      },
+      onRemoteDisconnected: async () => {
+        this._context?.onConnectionClosed(connection);
+      },
+      onRestartRequested: async () => {
+        if (!this._ctx || restartScheduled) {
+          return;
+        }
+
+        const restartDelay =
+          Math.min(MAX_RESTART_DELAY, INITIAL_RESTART_DELAY * reconnects) + Math.random() * RESTART_DELAY_JITTER;
+
+        log.info('connection restart scheduled', { spaceId, reconnects, restartDelay });
+
+        restartScheduled = true;
+        scheduleTask(
+          this._ctx,
+          async () => {
+            using _guard = await this._mutex.acquire();
+            if (this._connections.get(spaceId) !== connection) {
+              return;
+            }
+
+            const ctx = this._ctx;
+            await connection.close(); // Will call onRemoteDisconnected
+            this._connections.delete(spaceId);
+            if (ctx?.disposed) {
+              return;
+            }
+            await this._openConnection(spaceId, reconnects + 1);
+          },
+          restartDelay,
+        );
+      },
+    });
+    this._connections.set(spaceId, connection);
+
+    await connection.open();
+  }
+}
+
+type EdgeReplicatorConnectionsParams = {
+  edgeConnection: EdgeConnection;
+  spaceId: SpaceId;
+  context: EchoReplicatorContext;
+  sharedPolicyEnabled: boolean;
+  onRemoteConnected: () => Promise<void>;
+  onRemoteDisconnected: () => Promise<void>;
+  onRestartRequested: () => Promise<void>;
+};
+
+class EdgeReplicatorConnection extends Resource implements ReplicatorConnection {
+  private readonly _edgeConnection: EdgeConnection;
+  private _remotePeerId: string | null = null;
+  private readonly _targetServiceId: string;
+  private readonly _spaceId: SpaceId;
+  private readonly _context: EchoReplicatorContext;
+  private readonly _sharedPolicyEnabled: boolean;
+  private readonly _onRemoteConnected: () => Promise<void>;
+  private readonly _onRemoteDisconnected: () => Promise<void>;
+  private readonly _onRestartRequested: () => void;
+
+  private _readableStreamController!: ReadableStreamDefaultController<AutomergeProtocolMessage>;
+
+  public readable: ReadableStream<AutomergeProtocolMessage>;
+  public writable: WritableStream<AutomergeProtocolMessage>;
+
+  constructor({
+    edgeConnection,
+    spaceId,
+    context,
+    sharedPolicyEnabled,
+    onRemoteConnected,
+    onRemoteDisconnected,
+    onRestartRequested,
+  }: EdgeReplicatorConnectionsParams) {
+    super();
+    this._edgeConnection = edgeConnection;
+    this._spaceId = spaceId;
+    this._context = context;
+
+    // Generate a unique peer id for every connection.
+    // This way automerge-repo will have separate sync states for every connection.
+    // This is important because the previous connection might have had some messages that failed to deliver
+    // abd if we don't clear the sync-state, automerge will not attempt to deliver them again.
+    this._remotePeerId = `${EdgeService.AUTOMERGE_REPLICATOR}:${spaceId}-${randomUUID()}`;
+    this._targetServiceId = `${EdgeService.AUTOMERGE_REPLICATOR}:${spaceId}`;
+    this._sharedPolicyEnabled = sharedPolicyEnabled;
+    this._onRemoteConnected = onRemoteConnected;
+    this._onRemoteDisconnected = onRemoteDisconnected;
+    this._onRestartRequested = onRestartRequested;
+
+    this.readable = new ReadableStream<AutomergeProtocolMessage>({
+      start: (controller) => {
+        this._readableStreamController = controller;
+      },
+    });
+
+    this.writable = new WritableStream<AutomergeProtocolMessage>({
+      write: async (message: AutomergeProtocolMessage, controller) => {
+        await this._sendMessage(message);
+      },
+    });
+  }
+
+  protected override async _open(ctx: Context): Promise<void> {
+    log('open');
+    // TODO: handle reconnects
+    this._ctx.onDispose(
+      this._edgeConnection.addListener((msg: RouterMessage) => {
+        this._onMessage(msg);
+      }),
+    );
+
+    await this._onRemoteConnected();
+  }
+
+  protected override async _close(): Promise<void> {
+    log('close');
+    this._readableStreamController.close();
+    await this._onRemoteDisconnected();
+  }
+
+  get peerId(): string {
+    invariant(this._remotePeerId, 'Not connected');
+    return this._remotePeerId;
+  }
+
+  async shouldAdvertise(params: ShouldAdvertiseParams): Promise<boolean> {
+    if (!this._sharedPolicyEnabled) {
+      return true;
+    }
+    const spaceId = await this._context.getContainingSpaceIdForDocument(params.documentId);
+    if (!spaceId) {
+      // There's no spaceId if the document is not present locally. This means the sharePolicy check is being
+      // performed on message reception, so spaceId check was already performed in _onMessage.
+      return true;
+    }
+    return spaceId === this._spaceId;
+  }
+
+  shouldSyncCollection(params: ShouldSyncCollectionParams): boolean {
+    if (!this._sharedPolicyEnabled) {
+      return true;
+    }
+    const spaceId = getSpaceIdFromCollectionId(params.collectionId as CollectionId);
+    return spaceId === this._spaceId;
+  }
+
+  private _onMessage(message: RouterMessage) {
+    if (message.serviceId !== this._targetServiceId) {
+      return;
+    }
+
+    const payload = cbor.decode(message.payload!.value) as AutomergeProtocolMessage;
+    log('recv', () => {
+      const decodedData =
+        payload.type === 'sync' && payload.data
+          ? A.decodeSyncMessage(payload.data)
+          : payload.type === 'collection-state'
+            ? (payload as any).state
+            : payload;
+      return { from: message.serviceId, type: payload.type, decodedData };
+    });
+    // Fix the peer id.
+    payload.senderId = this._remotePeerId! as PeerId;
+    this._processMessage(payload);
+  }
+
+  private _processMessage(message: AutomergeProtocolMessage) {
+    // There's a race between the credentials being replicated that are needed for access control and the data replication.
+    // AutomergeReplicator might return a Forbidden error if the credentials are not yet replicated.
+    // We restart the connection with some delay to account for that.
+    if (isForbiddenErrorMessage(message)) {
+      this._onRestartRequested();
+      return;
+    }
+
+    this._readableStreamController.enqueue(message);
+  }
+
+  private async _sendMessage(message: AutomergeProtocolMessage) {
+    // Fix the peer id.
+    (message as any).targetId = this._targetServiceId as PeerId;
+
+    log('send', {
+      type: message.type,
+      senderId: message.senderId,
+      targetId: (message as any).targetId,
+      documentId: (message as any).documentId,
+    });
+    const encoded = cbor.encode(message);
+
+    await this._edgeConnection.send(
+      buf.create(RouterMessageSchema, {
+        serviceId: this._targetServiceId,
+        source: {
+          identityKey: this._edgeConnection.identityKey,
+          peerKey: this._edgeConnection.peerKey,
+        },
+        payload: { value: bufferToArray(encoded) },
+      }),
+    );
+  }
+}
+
+/**
+ * This message is sent by EDGE AutomergeReplicator when the authorization is denied.
+ */
+const isForbiddenErrorMessage = (message: AutomergeProtocolMessage) =>
+  message.type === 'error' && message.message === 'Forbidden';

package/src/edge/index.ts

@@ -0,0 +1,5 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+export * from './echo-edge-replicator';

package/src/index.ts

@@ -8,3 +8,4 @@ export * from './metadata';
 export * from './pipeline';
 export * from './space';
 export * from './automerge';
+export * from './edge';

package/src/metadata/metadata-store.ts

@@ -20,6 +20,7 @@ import {
   type IdentityRecord,
   type SpaceCache,
   type LargeSpaceMetadata,
+  type EdgeReplicationSetting,
 } from '@dxos/protocols/proto/dxos/echo/metadata';
 import { type Directory, type File } from '@dxos/random-access-storage';
 import { type Timeframe } from '@dxos/timeframe';
@@ -229,11 +230,20 @@ export class MetadataStore {
       return this._metadata.identity.haloSpace;
     }
 
-    const space = this.spaces.find((space) => space.key === spaceKey);
+    const space = this.spaces.find((space) => space.key.equals(spaceKey));
     invariant(space, 'Space not found');
     return space;
   }
 
+  hasSpace(spaceKey: PublicKey): boolean {
+    if (this._metadata.identity?.haloSpace.key.equals(spaceKey)) {
+      // Check if the space is the identity space.
+      return true;
+    }
+
+    return !!this.spaces.find((space) => space.key.equals(spaceKey));
+  }
+
   private _getLargeSpaceMetadata(key: PublicKey): LargeSpaceMetadata {
     let entry = this._spaceLargeMetadata.get(key);
     if (entry) {
@@ -288,7 +298,7 @@ export class MetadataStore {
 
   async addSpace(record: SpaceMetadata) {
     invariant(
-      !(this._metadata.spaces ?? []).find((space) => space.key === record.key),
+      !(this._metadata.spaces ?? []).find((space) => space.key.equals(record.key)),
       'Cannot overwrite existing space in metadata',
     );
 
@@ -336,6 +346,16 @@ export class MetadataStore {
     await this._saveSpaceLargeMetadata(spaceKey);
     await this.flush();
   }
+
+  getSpaceEdgeReplicationSetting(spaceKey: PublicKey): EdgeReplicationSetting | undefined {
+    return this.hasSpace(spaceKey) ? this._getSpace(spaceKey).edgeReplication : undefined;
+  }
+
+  async setSpaceEdgeReplicationSetting(spaceKey: PublicKey, setting: EdgeReplicationSetting) {
+    this._getSpace(spaceKey).edgeReplication = setting;
+    await this._save();
+    await this.flush();
+  }
 }
 
 const fromBytesInt32 = (buf: Buffer) => buf.readInt32LE(0);

package/src/pipeline/pipeline-stress.test.ts

@@ -2,9 +2,9 @@
 // Copyright 2022 DXOS.org
 //
 
-import expect from 'expect';
 import * as fc from 'fast-check';
-import { inspect } from 'util';
+import { inspect } from 'node:util';
+import { describe, expect, test } from 'vitest';
 
 import { asyncTimeout } from '@dxos/async';
 import { type FeedStore, type FeedWrapper } from '@dxos/feed-store';
@@ -12,7 +12,6 @@ import { PublicKey } from '@dxos/keys';
 import { log } from '@dxos/log';
 import { type FeedMessageBlock } from '@dxos/protocols';
 import { type FeedMessage } from '@dxos/protocols/proto/dxos/echo/feed';
-import { describe, test } from '@dxos/test';
 import { Timeframe } from '@dxos/timeframe';
 import { range } from '@dxos/util';
 
@@ -24,53 +23,51 @@ const NUM_MESSAGES = 10;
 
 // TODO(burdon): Describe test.
 describe('pipeline/stress test', () => {
-  test
-    .skip('stress', async () => {
-      const builder = new TestFeedBuilder();
-
-      const agentIds = range(NUM_AGENTS).map(() => PublicKey.random().toHex().slice(0, 8));
-      const anAgentId = fc.constantFrom(...agentIds);
-
-      const commands = fc.commands(
-        [
-          fc.tuple(anAgentId, fc.integer({ min: 1, max: 10 })).map(([agent, count]) => new WriteCommand(agent, count)),
-          fc.constant(new SyncCommand()),
-          anAgentId.map((agent) => new RestartCommand(agent)),
-        ],
-        { size: 'large' },
-      );
-
-      const model = fc.asyncProperty(commands, async (commands) => {
-        const feedStore = builder.createFeedStore();
-
-        const agents = new Map(agentIds.map((id) => [id, new Agent(builder, feedStore, id)]));
-        await Promise.all(Array.from(agents.values()).map((agent) => agent.open()));
-        await Promise.all(Array.from(agents.values()).map((agent) => agent.start()));
-
-        const setup: fc.ModelRunSetup<Model, Real> = () => ({
-          model: {},
-          real: {
-            feedStore,
-            agents,
-          },
-        });
-
-        try {
-          await fc.asyncModelRun(setup, [...commands, new SyncCommand()]);
-        } finally {
-          await Promise.all(Array.from(agents.values()).map((agent) => agent.stop()));
-          await Promise.all(Array.from(agents.values()).map((agent) => agent.close()));
-        }
+  test.skip('stress', { timeout: 60_000 }, async () => {
+    const builder = new TestFeedBuilder();
+
+    const agentIds = range(NUM_AGENTS).map(() => PublicKey.random().toHex().slice(0, 8));
+    const anAgentId = fc.constantFrom(...agentIds);
+
+    const commands = fc.commands(
+      [
+        fc.tuple(anAgentId, fc.integer({ min: 1, max: 10 })).map(([agent, count]) => new WriteCommand(agent, count)),
+        fc.constant(new SyncCommand()),
+        anAgentId.map((agent) => new RestartCommand(agent)),
+      ],
+      { size: 'large' },
+    );
+
+    const model = fc.asyncProperty(commands, async (commands) => {
+      const feedStore = builder.createFeedStore();
+
+      const agents = new Map(agentIds.map((id) => [id, new Agent(builder, feedStore, id)]));
+      await Promise.all(Array.from(agents.values()).map((agent) => agent.open()));
+      await Promise.all(Array.from(agents.values()).map((agent) => agent.start()));
+
+      const setup: fc.ModelRunSetup<Model, Real> = () => ({
+        model: {},
+        real: {
+          feedStore,
+          agents,
+        },
       });
 
-      const examples: [commands: Iterable<fc.AsyncCommand<Model, Real, boolean>>][] = [
-        [[new WriteCommand(agentIds[0], 10), new WriteCommand(agentIds[1], 10), new SyncCommand()]],
-        [[new WriteCommand(agentIds[0], 4), new RestartCommand(agentIds[0]), new SyncCommand()]],
-      ];
+      try {
+        await fc.asyncModelRun(setup, [...commands, new SyncCommand()]);
+      } finally {
+        await Promise.all(Array.from(agents.values()).map((agent) => agent.stop()));
+        await Promise.all(Array.from(agents.values()).map((agent) => agent.close()));
+      }
+    });
+
+    const examples: [commands: Iterable<fc.AsyncCommand<Model, Real, boolean>>][] = [
+      [[new WriteCommand(agentIds[0], 10), new WriteCommand(agentIds[1], 10), new SyncCommand()]],
+      [[new WriteCommand(agentIds[0], 4), new RestartCommand(agentIds[0]), new SyncCommand()]],
+    ];
 
-      await fc.assert(model, { examples });
-    })
-    .timeout(60_000);
+    await fc.assert(model, { examples });
+  });
 });
 
 class Agent {

package/src/pipeline/pipeline.test.ts

@@ -2,11 +2,10 @@
 // Copyright 2022 DXOS.org
 //
 
-import expect from 'expect';
+import { describe, expect, test, onTestFinished } from 'vitest';
 
 import { Event, sleep } from '@dxos/async';
 import { type FeedMessage } from '@dxos/protocols/proto/dxos/echo/feed';
-import { describe, test, afterTest } from '@dxos/test';
 import { Timeframe } from '@dxos/timeframe';
 import { range } from '@dxos/util';
 
@@ -62,7 +61,7 @@ describe('pipeline/Pipeline', () => {
 
   test('reading and writing with cursor changes', async () => {
     const pipeline = new Pipeline();
-    afterTest(() => pipeline.stop());
+    onTestFinished(() => pipeline.stop());
 
     const builder = new TestFeedBuilder();
     const feedStore = builder.createFeedStore();
@@ -105,7 +104,7 @@ describe('pipeline/Pipeline', () => {
 
   test('cursor change while polling', async () => {
     const pipeline = new Pipeline();
-    afterTest(() => pipeline.stop());
+    onTestFinished(() => pipeline.stop());
 
     const builder = new TestFeedBuilder();
     const feedStore = builder.createFeedStore();

package/src/space/control-pipeline.test.ts

@@ -2,7 +2,7 @@
 // Copyright 2022 DXOS.org
 //
 
-import expect from 'expect';
+import { describe, expect, test, onTestFinished } from 'vitest';
 
 import { CredentialGenerator, createCredential } from '@dxos/credentials';
 import { FeedFactory, FeedStore } from '@dxos/feed-store';
@@ -12,7 +12,6 @@ import { log } from '@dxos/log';
 import type { FeedMessage } from '@dxos/protocols/proto/dxos/echo/feed';
 import { AdmittedFeed } from '@dxos/protocols/proto/dxos/halo/credentials';
 import { createStorage, StorageType } from '@dxos/random-access-storage';
-import { describe, test, afterTest } from '@dxos/test';
 import { Timeframe } from '@dxos/timeframe';
 
 import { ControlPipeline } from './control-pipeline';
@@ -62,7 +61,7 @@ describe('space/control-pipeline', () => {
     await controlPipeline.setWriteFeed(genesisFeed);
     await controlPipeline.start();
 
-    afterTest(() => controlPipeline.stop());
+    onTestFinished(() => controlPipeline.stop());
 
     //
     // Genesis

package/src/space/control-pipeline.ts

@@ -160,7 +160,7 @@ export class ControlPipeline {
     };
     await this._pipeline.unpause();
 
-    log('save snapshot', { key: this._spaceKey, snapshot });
+    log('save snapshot', { key: this._spaceKey, snapshot: getSnapshotLoggerContext(snapshot) });
     await this._metadata.setSpaceControlPipelineSnapshot(this._spaceKey, snapshot);
   }
 
@@ -229,3 +229,12 @@ export class ControlPipeline {
     }
   }
 }
+
+const getSnapshotLoggerContext = (snapshot: ControlPipelineSnapshot) => {
+  return snapshot.messages?.map((msg) => {
+    const issuer = msg.credential.issuer;
+    const subject = msg.credential.subject.id;
+    const type = msg.credential.subject.assertion['@type'];
+    return { issuer, subject, type };
+  });
+};

package/src/space/replication.browser.test.ts

@@ -2,16 +2,12 @@
 // Copyright 2021 DXOS.org
 //
 
-// @dxos/test platform=browser
-
-import expect from 'expect';
-import waitForExpect from 'wait-for-expect';
+import { describe, expect, test } from 'vitest';
 
 import { FeedFactory, FeedStore } from '@dxos/feed-store';
 import { Keyring } from '@dxos/keyring';
 import type { FeedMessage } from '@dxos/protocols/proto/dxos/echo/feed';
 import { createStorage } from '@dxos/random-access-storage';
-import { describe, test } from '@dxos/test';
 import { Timeframe } from '@dxos/timeframe';
 
 import { valueEncoding } from '../common';
@@ -57,8 +53,6 @@ describe('replication', () => {
       timeframe: new Timeframe([[feed1.key, 123]]),
     });
 
-    await waitForExpect(() => {
-      expect(feed2.properties.length).toEqual(1);
-    });
+    await expect.poll(() => feed2.properties.length).toEqual(1);
   });
 });

package/src/space/space-manager.browser.test.ts

@@ -2,10 +2,9 @@
 // Copyright 2022 DXOS.org
 //
 
-// @dxos/test platform=browser
+import { describe, test, onTestFinished } from 'vitest';
 
 import { createStorage } from '@dxos/random-access-storage';
-import { describe, test, afterTest } from '@dxos/test';
 
 import { TestAgentBuilder, WebsocketNetworkManagerProvider } from '../testing';
 
@@ -20,7 +19,9 @@ describe('space-manager', () => {
       storage: createStorage(),
       networkManagerProvider: WebsocketNetworkManagerProvider(SIGNAL_URL),
     });
-    afterTest(async () => await builder.close());
+    onTestFinished(async () => {
+      await builder.close();
+    });
 
     const peer1 = await builder.createPeer();
     const spaceManager1 = peer1.spaceManager;
@@ -30,8 +31,8 @@ describe('space-manager', () => {
     const spaceManager2 = peer2.spaceManager;
     await spaceManager2.open();
 
-    afterTest(() => spaceManager1.close());
-    afterTest(() => spaceManager2.close());
+    onTestFinished(() => spaceManager1.close());
+    onTestFinished(() => spaceManager2.close());
 
     // const space1 = await spaceManager1.createSpace();
     // expect(space1.isOpen).to.be.true;