@dxos/echo-pipeline 0.6.12 → 0.6.13-main.548ca8d

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

@@ -0,0 +1,208 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { DeferredTask } from '@dxos/async';
+import { getHeads, type Doc } from '@dxos/automerge/automerge';
+import { type DocHandle, type DocumentId } from '@dxos/automerge/automerge-repo';
+import { Stream } from '@dxos/codec-protobuf';
+import { Context, Resource } from '@dxos/context';
+import { type SpaceDoc } from '@dxos/echo-protocol';
+import { type ObjectSnapshot, type Indexer, type IdToHeads } from '@dxos/indexing';
+import { log } from '@dxos/log';
+import { objectPointerCodec } from '@dxos/protocols';
+import { type IndexConfig } from '@dxos/protocols/proto/dxos/echo/indexing';
+import {
+  type QueryRequest,
+  type QueryResponse,
+  type QueryService,
+  type QueryResult,
+} from '@dxos/protocols/proto/dxos/echo/query';
+import { trace } from '@dxos/tracing';
+
+import { QueryState } from './query-state';
+import { type AutomergeHost, getSpaceKeyFromDoc } from '../automerge';
+
+export type QueryServiceParams = {
+  indexer: Indexer;
+  automergeHost: AutomergeHost;
+};
+
+/**
+ * Represents an active query (stream and query state connected to that stream).
+ */
+type ActiveQuery = {
+  state: QueryState;
+  sendResults: (results: QueryResult[]) => void;
+  close: () => Promise<void>;
+};
+
+export class QueryServiceImpl extends Resource implements QueryService {
+  private readonly _queries = new Set<ActiveQuery>();
+
+  private readonly _updateQueries = new DeferredTask(this._ctx, async () => {
+    await Promise.all(
+      Array.from(this._queries).map(async (query) => {
+        try {
+          const { changed } = await query.state.execQuery();
+          if (changed) {
+            query.sendResults(query.state.getResults());
+          }
+        } catch (err) {
+          log.catch(err);
+        }
+      }),
+    );
+  });
+
+  // TODO(burdon): OK for options, but not params. Pass separately and type readonly here.
+  constructor(private readonly _params: QueryServiceParams) {
+    super();
+
+    trace.diagnostic({
+      id: 'active-queries',
+      name: 'Active Queries',
+      fetch: () => {
+        return Array.from(this._queries).map((query) => {
+          return {
+            filter: JSON.stringify(query.state.filter),
+            metrics: query.state.metrics,
+          };
+        });
+      },
+    });
+  }
+
+  override async _open() {
+    this._params.indexer.updated.on(this._ctx, () => this._updateQueries.schedule());
+  }
+
+  override async _close() {
+    await Promise.all(Array.from(this._queries).map((query) => query.close()));
+  }
+
+  async setConfig(config: IndexConfig): Promise<void> {
+    if (this._params.indexer.initialized) {
+      log.warn('Indexer already initialized.');
+      return;
+    }
+    this._params.indexer.setConfig(config);
+  }
+
+  execQuery(request: QueryRequest): Stream<QueryResponse> {
+    return new Stream<QueryResponse>(({ next, close, ctx }) => {
+      const query: ActiveQuery = {
+        state: new QueryState({
+          indexer: this._params.indexer,
+          automergeHost: this._params.automergeHost,
+          request,
+        }),
+        sendResults: (results) => {
+          if (ctx.disposed) {
+            return;
+          }
+          next({ queryId: request.queryId, results });
+        },
+        close: async () => {
+          close();
+          await query.state.close();
+          this._queries.delete(query);
+        },
+      };
+      this._queries.add(query);
+
+      queueMicrotask(async () => {
+        await query.state.open();
+
+        try {
+          const { changed } = await query.state.execQuery();
+          if (changed) {
+            query.sendResults(query.state.getResults());
+          }
+        } catch (error) {
+          log.catch(error);
+        }
+      });
+
+      return query.close;
+    });
+  }
+
+  /**
+   * Re-index all loaded documents.
+   */
+  async reindex() {
+    log.info('Reindexing all documents...');
+    const iterator = createDocumentsIterator(this._params.automergeHost);
+    const ids: IdToHeads = new Map();
+    for await (const documents of iterator()) {
+      for (const { id, heads } of documents) {
+        ids.set(id, heads);
+      }
+      if (ids.size % 100 === 0) {
+        log.info('Collected documents...', { count: ids.size });
+      }
+    }
+
+    log.info('Marking all documents as dirty...', { count: ids.size });
+    await this._params.indexer.reindex(ids);
+  }
+}
+
+/**
+ * Factory for `getAllDocuments` iterator.
+ */
+// TODO(dmaretskyi): Get roots from echo-host.
+const createDocumentsIterator = (automergeHost: AutomergeHost) =>
+  /**
+   * Recursively get all object data blobs from loaded documents from Automerge Repo.
+   */
+  // TODO(mykola): Unload automerge handles after usage.
+  async function* getAllDocuments(): AsyncGenerator<ObjectSnapshot[], void, void> {
+    /** visited automerge handles */
+    const visited = new Set<string>();
+
+    async function* getObjectsFromHandle(handle: DocHandle<any>): AsyncGenerator<ObjectSnapshot[]> {
+      if (visited.has(handle.documentId)) {
+        return;
+      }
+      const doc: Doc<SpaceDoc> = handle.docSync();
+
+      const spaceKey = getSpaceKeyFromDoc(doc) ?? undefined;
+
+      if (doc.objects) {
+        yield Object.entries(doc.objects as { [key: string]: any }).map(([objectId, object]) => {
+          return {
+            id: objectPointerCodec.encode({ documentId: handle.documentId, objectId, spaceKey }),
+            object,
+            heads: getHeads(doc),
+          };
+        });
+      }
+
+      if (doc.links) {
+        for (const id of Object.values(doc.links as { [echoId: string]: string })) {
+          if (visited.has(id)) {
+            continue;
+          }
+          const linkHandle = await automergeHost.loadDoc(Context.default(), id as DocumentId);
+          for await (const result of getObjectsFromHandle(linkHandle)) {
+            yield result;
+          }
+        }
+      }
+
+      visited.add(handle.documentId);
+    }
+
+    // TODO(mykola): Use list of roots instead of iterating over all handles.
+    for (const handle of Object.values(automergeHost.repo.handles)) {
+      if (visited.has(handle.documentId)) {
+        continue;
+      }
+      for await (const result of getObjectsFromHandle(handle)) {
+        yield result;
+      }
+      visited.add(handle.documentId);
+    }
+  };

package/src/db-host/query-state.ts

@@ -0,0 +1,200 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { type DocumentId } from '@dxos/automerge/automerge-repo';
+import { Context, LifecycleState, Resource } from '@dxos/context';
+import { createIdFromSpaceKey } from '@dxos/echo-protocol';
+import { type Indexer, type IndexQuery } from '@dxos/indexing';
+import { invariant } from '@dxos/invariant';
+import { PublicKey } from '@dxos/keys';
+import { objectPointerCodec } from '@dxos/protocols';
+import { type Filter as FilterProto } from '@dxos/protocols/proto/dxos/echo/filter';
+import { type QueryRequest, type QueryResult } from '@dxos/protocols/proto/dxos/echo/query';
+import { trace } from '@dxos/tracing';
+import { nonNullable } from '@dxos/util';
+
+import { type AutomergeHost, getSpaceKeyFromDoc } from '../automerge';
+
+type QueryStateParams = {
+  indexer: Indexer;
+  automergeHost: AutomergeHost;
+  request: QueryRequest;
+};
+
+type QueryRunResult = {
+  changed: boolean;
+};
+
+export type QueryMetrics = {
+  objectsReturned: number;
+  objectsReturnedFromIndex: number;
+  documentsLoaded: number;
+  executionTime: number;
+  indexQueryTime: number;
+  documentLoadTime: number;
+};
+
+/**
+ * Manages querying logic on service side.
+ */
+@trace.resource()
+export class QueryState extends Resource {
+  private _results: QueryResult[] = [];
+
+  /**
+   * Metrics are only captured for the first run of the query since that is the most representative.
+   * We plan to change the query logic so that reactive updates do not require a full re-run of the query.
+   */
+  private _firstRun = true;
+
+  @trace.info({ depth: null })
+  public readonly filter: FilterProto;
+
+  @trace.info()
+  public metrics: QueryMetrics = {
+    objectsReturned: 0,
+    objectsReturnedFromIndex: 0,
+    documentsLoaded: 0,
+    executionTime: 0,
+    indexQueryTime: 0,
+    documentLoadTime: 0,
+  };
+
+  @trace.info()
+  get active() {
+    return this._lifecycleState === LifecycleState.OPEN;
+  }
+
+  constructor(private readonly _params: QueryStateParams) {
+    super();
+    this.filter = _params.request.filter;
+  }
+
+  getResults() {
+    return this._results;
+  }
+
+  // https://github.com/open-telemetry/semantic-conventions/blob/main/docs/attributes-registry/db.md#generic-database-attributes
+  @trace.span({ showInBrowserTimeline: true, op: 'db.query', attributes: { 'db.system': 'echo' } })
+  async execQuery(): Promise<QueryRunResult> {
+    const filter = this._params.request.filter;
+
+    const beginQuery = performance.now();
+
+    // For object id filters, we return no results as those are handled by the SpaceQuerySource.
+    const hits =
+      filter.objectIds && filter.objectIds?.length > 0
+        ? []
+        : await this._params.indexer.execQuery(filterToIndexQuery(filter));
+    if (this._firstRun) {
+      this.metrics.indexQueryTime = performance.now() - beginQuery;
+    }
+
+    const beginFilter = performance.now();
+
+    const results: QueryResult[] = (
+      await Promise.all(
+        hits.map(async (result) => {
+          if (this._firstRun) {
+            this.metrics.objectsReturnedFromIndex++;
+          }
+
+          const { objectId, documentId, spaceKey: spaceKeyInIndex } = objectPointerCodec.decode(result.id);
+
+          let spaceKey: string | null;
+          if (spaceKeyInIndex !== undefined) {
+            spaceKey = spaceKeyInIndex;
+          } else {
+            // Indexes created by older versions of the indexer do not have the spaceKey in the index.
+            // If the spaceKey is not in the index, we need to load the document to get it.
+
+            if (this._firstRun) {
+              this.metrics.documentsLoaded++;
+            }
+
+            const handle = await this._params.automergeHost.loadDoc(Context.default(), documentId as DocumentId);
+
+            // `whenReady` creates a timeout so we guard it with an if to skip it if the handle is already ready.
+            if (this._ctx.disposed) {
+              return;
+            }
+            spaceKey = getSpaceKeyFromDoc(handle.docSync());
+          }
+
+          if (!spaceKey) {
+            return;
+          }
+          // TODO(mykola): Remove business logic from here.
+          if (
+            this._params.request.filter.options?.spaces?.length &&
+            !this._params.request.filter.options.spaces.some((key) => key.equals(spaceKey!))
+          ) {
+            return;
+          }
+
+          if (this._firstRun) {
+            this.metrics.objectsReturned++;
+          }
+
+          return {
+            id: objectId,
+            documentId,
+            spaceId: await createIdFromSpaceKey(PublicKey.from(spaceKey)),
+            spaceKey: PublicKey.from(spaceKey),
+            rank: result.rank,
+          } satisfies QueryResult;
+        }),
+      )
+    ).filter(nonNullable);
+
+    if (this._firstRun) {
+      this.metrics.documentLoadTime = performance.now() - beginFilter;
+    }
+
+    if (this._ctx.disposed) {
+      return { changed: false };
+    }
+
+    const areResultsUnchanged =
+      !this._firstRun &&
+      this._results.length === results.length &&
+      this._results.every((oldResult) => results.some((result) => result.id === oldResult.id)) &&
+      results.every((result) => this._results.some((oldResult) => oldResult.id === result.id));
+
+    if (this._firstRun) {
+      this.metrics.executionTime = performance.now() - beginQuery;
+    }
+
+    this._firstRun = false;
+    if (areResultsUnchanged) {
+      return { changed: false };
+    }
+
+    this._results = results;
+    return { changed: true };
+  }
+}
+
+// TODO(burdon): Process Filter DSL.
+const filterToIndexQuery = (filter: FilterProto): IndexQuery => {
+  invariant(!(filter.type && (filter.or ?? []).length > 0), 'Cannot mix type and or filters.');
+  invariant(
+    (filter.or ?? []).every((subFilter) => !(subFilter.type && (subFilter.or ?? []).length > 0)),
+    'Cannot mix type and or filters.',
+  );
+  if (
+    filter.type ||
+    ((filter.or ?? []).length > 0 && (filter.or ?? []).every((subFilter) => !subFilter.not && subFilter.type))
+  ) {
+    return {
+      typenames: filter.type?.objectId
+        ? [filter.type.objectId]
+        : (filter.or ?? []).map((f) => f.type?.objectId).filter(nonNullable),
+      inverted: filter.not,
+    };
+  } else {
+    // Query all objects.
+    return { typenames: [] };
+  }
+};

package/src/db-host/space-state-manager.ts

@@ -0,0 +1,90 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import isEqual from 'lodash.isequal';
+
+import { Event, UpdateScheduler } from '@dxos/async';
+import { interpretAsDocumentId, type DocHandle, type DocumentId } from '@dxos/automerge/automerge-repo';
+import { Resource, Context } from '@dxos/context';
+import { type SpaceDoc } from '@dxos/echo-protocol';
+import { type SpaceId } from '@dxos/keys';
+
+import { DatabaseRoot } from './database-root';
+
+export class SpaceStateManager extends Resource {
+  private readonly _roots = new Map<DocumentId, DatabaseRoot>();
+  private readonly _rootBySpace = new Map<SpaceId, DocumentId>();
+  private readonly _perRootContext = new Map<DocumentId, Context>();
+  private readonly _lastSpaceDocumentList = new Map<SpaceId, DocumentId[]>();
+
+  public readonly spaceDocumentListUpdated = new Event<SpaceDocumentListUpdatedEvent>();
+
+  protected override async _close(ctx: Context): Promise<void> {
+    for (const [_, rootCtx] of this._perRootContext) {
+      await rootCtx.dispose();
+    }
+    this._roots.clear();
+  }
+
+  get roots(): ReadonlyMap<DocumentId, DatabaseRoot> {
+    return this._roots;
+  }
+
+  getRootByDocumentId(documentId: DocumentId): DatabaseRoot | undefined {
+    return this._roots.get(documentId);
+  }
+
+  async assignRootToSpace(spaceId: SpaceId, handle: DocHandle<SpaceDoc>): Promise<DatabaseRoot> {
+    let root: DatabaseRoot;
+    if (this._roots.has(handle.documentId)) {
+      root = this._roots.get(handle.documentId)!;
+    } else {
+      root = new DatabaseRoot(handle);
+      this._roots.set(handle.documentId, root);
+    }
+
+    if (this._rootBySpace.get(spaceId) === root.handle.documentId) {
+      return root;
+    }
+
+    const prevRootId = this._rootBySpace.get(spaceId);
+    if (prevRootId) {
+      void this._perRootContext.get(prevRootId)?.dispose();
+      this._perRootContext.delete(prevRootId);
+    }
+
+    this._rootBySpace.set(spaceId, root.handle.documentId);
+    const ctx = new Context();
+
+    this._perRootContext.set(root.handle.documentId, ctx);
+
+    await root.handle.whenReady();
+
+    const documentListCheckScheduler = new UpdateScheduler(
+      ctx,
+      async () => {
+        const documentIds = [root.documentId, ...root.getAllLinkedDocuments().map((url) => interpretAsDocumentId(url))];
+        if (!isEqual(documentIds, this._lastSpaceDocumentList.get(spaceId))) {
+          this._lastSpaceDocumentList.set(spaceId, documentIds);
+          this.spaceDocumentListUpdated.emit(new SpaceDocumentListUpdatedEvent(spaceId, documentIds));
+        }
+      },
+      { maxFrequency: 50 },
+    );
+    const triggerCheckOnChange = () => documentListCheckScheduler.trigger();
+    root.handle.addListener('change', triggerCheckOnChange);
+    ctx.onDispose(() => root.handle.removeListener('change', triggerCheckOnChange));
+
+    documentListCheckScheduler.trigger();
+
+    return root;
+  }
+}
+
+export class SpaceDocumentListUpdatedEvent {
+  constructor(
+    public readonly spaceId: SpaceId,
+    public readonly documentIds: DocumentId[],
+  ) {}
+}

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

@@ -0,0 +1,96 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { describe, test } from 'vitest';
+
+import { Event } from '@dxos/async';
+import { cbor } from '@dxos/automerge/automerge-repo';
+import { createEphemeralEdgeIdentity, EdgeClient, MessageSchema } from '@dxos/edge-client';
+import { createTestEdgeWsServer } from '@dxos/edge-client/testing';
+import { PublicKey, SpaceId } from '@dxos/keys';
+import { EdgeService } from '@dxos/protocols';
+import type { AutomergeProtocolMessage } from '@dxos/protocols';
+import { createBuf } from '@dxos/protocols/buf';
+import type { Peer } from '@dxos/protocols/proto/dxos/edge/messenger';
+import { openAndClose } from '@dxos/test-utils';
+
+import { EchoEdgeReplicator } from './echo-edge-replicator';
+import type { EchoReplicatorContext } from '../automerge';
+
+describe('EchoEdgeReplicator', () => {
+  test('reconnects', async ({ onTestFinished }) => {
+    const { endpoint, cleanup, sendMessage } = await createTestEdgeWsServer(8001);
+    onTestFinished(cleanup);
+    const client = new EdgeClient(await createEphemeralEdgeIdentity(), { socketEndpoint: endpoint });
+    await openAndClose(client);
+
+    const spaceId = SpaceId.random();
+
+    const replicator = new EchoEdgeReplicator({ edgeConnection: client });
+    const { context, connectionOpen } = createMockContext();
+    await replicator.connect(context);
+    await replicator.connectToSpace(spaceId);
+
+    client.setIdentity(await createEphemeralEdgeIdentity());
+    await connectionOpen.waitForCount(1);
+
+    sendMessage(
+      createForbiddenMessage(
+        {
+          identityKey: client.identityKey,
+          peerKey: client.peerKey,
+        },
+        spaceId,
+      ),
+    );
+    await connectionOpen.waitForCount(1);
+
+    // Double restart to check for race conditions.
+    client.setIdentity(await createEphemeralEdgeIdentity());
+    sendMessage(
+      createForbiddenMessage(
+        {
+          identityKey: client.identityKey,
+          peerKey: client.peerKey,
+        },
+        spaceId,
+      ),
+    );
+    await connectionOpen.waitForCount(1);
+
+    await replicator.disconnect();
+  });
+});
+
+const createMockContext = () => {
+  const connectionOpen = new Event();
+  return {
+    context: {
+      getContainingSpaceIdForDocument: async (documentId) => null,
+      getContainingSpaceForDocument: async (documentId) => null,
+      onConnectionAuthScopeChanged: (connection) => {},
+      isDocumentInRemoteCollection: async (params) => false,
+      onConnectionClosed: (connection) => {},
+      onConnectionOpen: (connection) => {
+        connectionOpen.emit();
+      },
+
+      peerId: PublicKey.random().toHex(),
+    } satisfies EchoReplicatorContext,
+
+    connectionOpen,
+  };
+};
+
+const createForbiddenMessage = (target: Peer, spaceId: SpaceId) =>
+  createBuf(MessageSchema, {
+    target: [target],
+    serviceId: `${EdgeService.AUTOMERGE_REPLICATOR}:${spaceId}`,
+    payload: {
+      value: cbor.encode({
+        type: 'error',
+        message: 'Forbidden',
+      } satisfies AutomergeProtocolMessage),
+    },
+  });