@dxos/echo-pipeline 0.8.3 → 0.8.4-main.1068cf700f

package/src/automerge/automerge-host.ts

@@ -3,57 +3,55 @@
 //
 
 import {
+  type Doc,
+  type Heads,
   getBackend,
   getHeads,
-  isAutomerge,
   equals as headsEquals,
+  isAutomerge,
   save,
-  type Doc,
-  type Heads,
 } from '@automerge/automerge';
 import {
-  type DocHandleChangePayload,
-  Repo,
   type AnyDocumentId,
   type DocHandle,
+  type DocHandleChangePayload,
   type DocumentId,
+  type HandleState,
   type PeerCandidatePayload,
   type PeerDisconnectedPayload,
   type PeerId,
+  Repo,
   type StorageAdapterInterface,
   type StorageKey,
   interpretAsDocumentId,
-  type HandleState,
 } from '@automerge/automerge-repo';
+import type * as Record from 'effect/Record';
 
-import { Event, asyncTimeout } from '@dxos/async';
-import { Context, Resource, cancelWithContext, type Lifecycle } from '@dxos/context';
-import { DatabaseDirectory, type CollectionId } from '@dxos/echo-protocol';
-import { type IndexMetadataStore } from '@dxos/indexing';
+import { DeferredTask, Event, asyncTimeout } from '@dxos/async';
+import { Context, type Lifecycle, Resource, cancelWithContext } from '@dxos/context';
+import { type CollectionId, DatabaseDirectory } from '@dxos/echo-protocol';
 import { invariant } from '@dxos/invariant';
 import { PublicKey } from '@dxos/keys';
 import { type LevelDB } from '@dxos/kv-store';
 import { log } from '@dxos/log';
-import { objectPointerCodec } from '@dxos/protocols';
+import { type SpaceSyncState } from '@dxos/protocols/proto/dxos/echo/service';
 import { type DocHeadsList, type FlushRequest } from '@dxos/protocols/proto/dxos/echo/service';
 import { trace } from '@dxos/tracing';
-import { bufferToArray } from '@dxos/util';
+import { ComplexSet, bufferToArray, defaultMap, isNonNullable, range } from '@dxos/util';
 
-import { CollectionSynchronizer, diffCollectionState, type CollectionState } from './collection-synchronizer';
+import { type CollectionState, CollectionSynchronizer, diffCollectionState } from './collection-synchronizer';
 import { type EchoDataMonitor } from './echo-data-monitor';
 import { EchoNetworkAdapter, isEchoPeerMetadata } from './echo-network-adapter';
-import { type EchoReplicator, type RemoteDocumentExistenceCheckParams } from './echo-replicator';
+import { type AutomergeReplicator, type RemoteDocumentExistenceCheckProps } from './echo-replicator';
 import { HeadsStore } from './heads-store';
-import { LevelDBStorageAdapter, type BeforeSaveParams } from './leveldb-storage-adapter';
+import { type BeforeSaveProps, LevelDBStorageAdapter } from './leveldb-storage-adapter';
 
 export type PeerIdProvider = () => string | undefined;
 
 export type RootDocumentSpaceKeyProvider = (documentId: string) => PublicKey | undefined;
 
-export type AutomergeHostParams = {
+export type AutomergeHostProps = {
   db: LevelDB;
-
-  indexMetadataStore: IndexMetadataStore;
   dataMonitor?: EchoDataMonitor;
 
   /**
@@ -65,6 +63,12 @@ export type AutomergeHostParams = {
 
 export type LoadDocOptions = {
   timeout?: number;
+
+  /**
+   * If true, the document will be fetched from the network if it is not found in the local storage.
+   * Setting this to false does not guarantee that the document will not be fetched from the network.
+   */
+  fetchFromNetwork?: boolean;
 };
 
 export type CreateDocOptions = {
@@ -72,19 +76,38 @@ export type CreateDocOptions = {
    * Import the document together with its history.
    */
   preserveHistory?: boolean;
+
+  documentId?: DocumentId;
 };
 
 export const FIND_PARAMS = {
   allowableStates: ['ready', 'requesting'] satisfies HandleState[],
 };
 
+/**
+ * Maximum amount of documents to sync in a single bundle.
+ */
+const BUNDLE_SIZE = 100;
+/**
+ * Maximum amount of concurrent tasks to run when pushing or pulling bundles.
+ */
+const BUNDLE_SYNC_CONCURRENCY = 2;
+/**
+ * If the number of documents to sync is greater than this threshold, we will use bundles.
+ */
+const BUNDLE_SYNC_THRESHOLD = 50;
+
+/**
+ * Only announce documents that are known to require sync.
+ */
+const OPTIMIZED_SHARE_POLICY = true;
+
 /**
  * Abstracts over the AutomergeRepo.
  */
 @trace.resource()
 export class AutomergeHost extends Resource {
   private readonly _db: LevelDB;
-  private readonly _indexMetadataStore: IndexMetadataStore;
   private readonly _echoNetworkAdapter: EchoNetworkAdapter;
 
   private readonly _collectionSynchronizer = new CollectionSynchronizer({
@@ -97,6 +120,14 @@ export class AutomergeHost extends Resource {
   private _storage!: StorageAdapterInterface & Lifecycle;
   private readonly _headsStore: HeadsStore;
 
+  private _syncTask: DeferredTask | undefined = undefined;
+  /**
+   * Cache of collections that would be synced on next sync task run.
+   */
+  private readonly _collectionsToSync = new ComplexSet<{ collectionId: string; peerId: PeerId }>(
+    ({ collectionId, peerId }) => `${collectionId}|${peerId}`,
+  );
+
   @trace.info()
   private _peerId!: PeerId;
 
@@ -110,13 +141,32 @@ export class AutomergeHost extends Resource {
    */
   public readonly documentsSaved = new Event();
 
-  constructor({
-    db,
-    indexMetadataStore,
-    dataMonitor,
-    peerIdProvider,
-    getSpaceKeyByRootDocumentId,
-  }: AutomergeHostParams) {
+  private readonly _headsUpdates = new Map<DocumentId, Heads>();
+  private _onHeadsChangedTask?: DeferredTask;
+
+  /**
+   * Documents created in this session.
+   */
+  private _createdDocuments = new Set<DocumentId>();
+
+  /**
+   * Documents that need to be synced based on the result of collection-sync.
+   */
+  private _documentsToSync = new Set<DocumentId>();
+
+  /**
+   * Documents that are not avaiale locally that should be requested.
+   */
+  private _documentsToRequest = new Set<DocumentId>();
+
+  /**
+   * Documents requested by remote peers.
+   */
+  private _documentsRequested = new Map<PeerId, Set<DocumentId>>();
+
+  private _sharePolicyChangedTask?: DeferredTask;
+
+  constructor({ db, dataMonitor, peerIdProvider, getSpaceKeyByRootDocumentId }: AutomergeHostProps) {
     super();
     this._db = db;
     this._storage = new LevelDBStorageAdapter({
@@ -134,8 +184,11 @@ export class AutomergeHost extends Resource {
       onCollectionStateReceived: this._onCollectionStateReceived.bind(this),
       monitor: dataMonitor,
     });
+    this._echoNetworkAdapter.documentRequested.on(({ peerId, documentId }) => {
+      defaultMap(this._documentsRequested, peerId, () => new Set()).add(documentId);
+      this._sharePolicyChangedTask!.schedule();
+    });
     this._headsStore = new HeadsStore({ db: db.sublevel('heads') });
-    this._indexMetadataStore = indexMetadataStore;
     this._peerIdProvider = peerIdProvider;
     this._getSpaceKeyByRootDocumentId = getSpaceKeyByRootDocumentId;
   }
@@ -143,12 +196,18 @@ export class AutomergeHost extends Resource {
   protected override async _open(): Promise<void> {
     this._peerId = `host-${this._peerIdProvider?.() ?? PublicKey.random().toHex()}` as PeerId;
 
+    this._onHeadsChangedTask = new DeferredTask(this._ctx, async () => {
+      const docHeads = Array.from(this._headsUpdates.entries());
+      this._headsUpdates.clear();
+      this._onHeadsChanged(docHeads);
+    });
+
     await this._storage.open?.();
 
     // Construct the automerge repo.
     this._repo = new Repo({
       peerId: this._peerId as PeerId,
-      sharePolicy: this._sharePolicy.bind(this),
+      shareConfig: this._shareConfig,
       storage: this._storage,
       network: [
         // Upstream swarm.
@@ -180,6 +239,27 @@ export class AutomergeHost extends Resource {
       }
     });
 
+    this._syncTask = new DeferredTask(this._ctx, async () => {
+      const collectionToSync = Array.from(this._collectionsToSync.values());
+      if (collectionToSync.length === 0) {
+        return;
+      }
+      await Promise.all(
+        collectionToSync.map(async ({ collectionId, peerId }) => {
+          try {
+            await this._handleCollectionSync(collectionId, peerId);
+          } catch (err) {
+            log.error('failed to sync collection', { collectionId, peerId, err });
+          }
+        }),
+      );
+    });
+
+    this._sharePolicyChangedTask = new DeferredTask(this._ctx, async () => {
+      log('share policy changed');
+      this._repo.shareConfigChanged();
+    });
+
     await this._echoNetworkAdapter.open();
     await this._collectionSynchronizer.open();
     await this._echoNetworkAdapter.open();
@@ -190,14 +270,9 @@ export class AutomergeHost extends Resource {
     await this._collectionSynchronizer.close();
     await this._storage.close?.();
     await this._echoNetworkAdapter.close();
-    await this._ctx.dispose();
-  }
-
-  /**
-   * @deprecated To be abstracted away.
-   */
-  get repo(): Repo {
-    return this._repo;
+    this._syncTask = undefined;
+    this._onHeadsChangedTask = undefined;
+    this._sharePolicyChangedTask = undefined;
   }
 
   get peerId(): PeerId {
@@ -208,11 +283,17 @@ export class AutomergeHost extends Resource {
     return Object.keys(this._repo.handles).length;
   }
 
-  async addReplicator(replicator: EchoReplicator): Promise<void> {
+  get handles(): Readonly<Record<DocumentId, DocHandle<any>>> {
+    return this._repo.handles;
+  }
+
+  async addReplicator(replicator: AutomergeReplicator): Promise<void> {
+    invariant(this.isOpen, 'AutomergeHost is not open');
     await this._echoNetworkAdapter.addReplicator(replicator);
   }
 
-  async removeReplicator(replicator: EchoReplicator): Promise<void> {
+  async removeReplicator(replicator: AutomergeReplicator): Promise<void> {
+    invariant(this.isOpen, 'AutomergeHost is not open');
     await this._echoNetworkAdapter.removeReplicator(replicator);
   }
 
@@ -220,13 +301,29 @@ export class AutomergeHost extends Resource {
    * Loads the document handle from the repo and waits for it to be ready.
    */
   async loadDoc<T>(ctx: Context, documentId: AnyDocumentId, opts?: LoadDocOptions): Promise<DocHandle<T>> {
+    invariant(this.isOpen, 'AutomergeHost is not open');
     let handle: DocHandle<T> | undefined;
     if (typeof documentId === 'string') {
       // NOTE: documentId might also be a URL, in which case this lookup will fail.
       handle = this._repo.handles[documentId as DocumentId];
     }
     if (!handle) {
-      handle = await this._repo.find(documentId as DocumentId, FIND_PARAMS);
+      if (!opts?.fetchFromNetwork) {
+        handle = await this._repo.find(documentId as DocumentId, FIND_PARAMS);
+      } else {
+        try {
+          handle = await this._repo.find(documentId as DocumentId, {
+            allowableStates: ['ready', 'requesting', 'unavailable'],
+          });
+          if (handle.state === 'requesting' || handle.state === 'unavailable') {
+            this._documentsToRequest.add(handle.documentId);
+            this._sharePolicyChangedTask!.schedule();
+          }
+        } catch (err) {
+          log.error('failed to load document', { documentId, err });
+          throw err;
+        }
+      }
     }
 
     // `whenReady` creates a timeout so we guard it with an if to skip it if the handle is already ready.
@@ -242,6 +339,7 @@ export class AutomergeHost extends Resource {
   }
 
   async exportDoc(ctx: Context, id: AnyDocumentId): Promise<Uint8Array> {
+    invariant(this.isOpen, 'AutomergeHost is not open');
     const documentId = interpretAsDocumentId(id);
 
     const chunks = await this._storage.loadRange([documentId]);
@@ -251,10 +349,11 @@ export class AutomergeHost extends Resource {
   /**
    * Create new persisted document.
    */
-  createDoc<T>(initialValue?: T | Doc<T> | Uint8Array, opts?: CreateDocOptions): DocHandle<T> {
+  async createDoc<T>(initialValue?: T | Doc<T> | Uint8Array, opts?: CreateDocOptions): Promise<DocHandle<T>> {
+    invariant(this.isOpen, 'AutomergeHost is not open');
     if (opts?.preserveHistory) {
       if (initialValue instanceof Uint8Array) {
-        return this._repo.import(initialValue);
+        return this._repo.import(initialValue, { docId: opts?.documentId });
       }
 
       if (!isAutomerge(initialValue)) {
@@ -262,17 +361,27 @@ export class AutomergeHost extends Resource {
       }
 
       // TODO(dmaretskyi): There's a more efficient way.
-      return this._repo.import(save(initialValue as Doc<T>));
+      const handle = this._repo.import(save(initialValue as Doc<T>), { docId: opts?.documentId });
+      this._createdDocuments.add(handle.documentId);
+      this._sharePolicyChangedTask!.schedule();
+      return handle as DocHandle<T>;
     } else {
       if (initialValue instanceof Uint8Array) {
         throw new Error('Cannot create document from Uint8Array without preserving history');
       }
 
-      return this._repo.create(initialValue);
+      if (opts?.documentId) {
+        throw new Error('Cannot prefil document id when not importing an existing doc');
+      }
+      const handle = await this._repo.create2<T>(initialValue);
+      this._createdDocuments.add(handle.documentId);
+      this._sharePolicyChangedTask!.schedule();
+      return handle;
     }
   }
 
   async waitUntilHeadsReplicated(heads: DocHeadsList): Promise<void> {
+    invariant(this.isOpen, 'AutomergeHost is not open');
     const entries = heads.entries;
     if (!entries?.length) {
       return;
@@ -289,7 +398,7 @@ export class AutomergeHost extends Resource {
     });
     if (headsToWait.length > 0) {
       await Promise.all(
-        headsToWait.map(async (entry, index) => {
+        headsToWait.map(async (entry) => {
           const handle = await this.loadDoc<DatabaseDirectory>(Context.default(), entry.documentId as DocumentId);
           await waitForHeads(handle, entry.heads!);
         }),
@@ -303,6 +412,7 @@ export class AutomergeHost extends Resource {
   }
 
   async reIndexHeads(documentIds: DocumentId[]): Promise<void> {
+    invariant(this.isOpen, 'AutomergeHost is not open');
     for (const documentId of documentIds) {
       log('re-indexing heads for document', { documentId });
       const handle = await this._repo.find(documentId, FIND_PARAMS);
@@ -319,28 +429,42 @@ export class AutomergeHost extends Resource {
     log('done re-indexing 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, documentId?: DocumentId): Promise<boolean> {
-    if (peerId.startsWith('client-')) {
-      return false; // Only send docs to clients if they are requested.
-    }
+  private readonly _shareConfig = {
+    access: async (peerId: PeerId, documentId: DocumentId): Promise<boolean> => {
+      return true;
+    },
 
-    if (!documentId) {
-      return false;
-    }
+    // 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
+    // Called for all loaded documents so they could be advertised to the sync server.
+    announce: async (peerId: PeerId, documentId?: DocumentId): Promise<boolean> => {
+      if (!documentId) {
+        return false;
+      }
 
-    const peerMetadata = this.repo.peerMetadataByPeerId[peerId];
-    if (isEchoPeerMetadata(peerMetadata)) {
-      return this._echoNetworkAdapter.shouldAdvertise(peerId, { documentId });
-    }
+      if (OPTIMIZED_SHARE_POLICY) {
+        if (
+          !this._createdDocuments.has(documentId) &&
+          !this._documentsToSync.has(documentId) &&
+          !this._documentsToRequest.has(documentId)
+        ) {
+          // Skip advertising documents that don't need to be synced.
+          return false;
+        }
+      }
 
-    return false;
-  }
+      const peerMetadata = this._repo.peerMetadataByPeerId[peerId];
+      if (isEchoPeerMetadata(peerMetadata)) {
+        return this._echoNetworkAdapter.shouldAdvertise(peerId, { documentId });
+      }
+
+      return false;
+    },
+  };
 
-  private async _beforeSave({ path, batch }: BeforeSaveParams): Promise<void> {
+  private async _beforeSave({ path, batch }: BeforeSaveProps): Promise<void> {
     const handle = this._repo.handles[path[0] as DocumentId];
     if (!handle || !handle.isReady()) {
       return;
@@ -352,14 +476,6 @@ export class AutomergeHost extends Resource {
 
     const heads = getHeads(doc);
     this._headsStore.setHeads(handle.documentId, heads, batch);
-
-    const spaceKey = DatabaseDirectory.getSpaceKey(doc) ?? undefined;
-    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, heads]));
-    this._indexMetadataStore.markDirty(idToLastHash, batch);
   }
 
   private _shouldSyncCollection(collectionId: string, peerId: PeerId): boolean {
@@ -375,14 +491,24 @@ export class AutomergeHost extends Resource {
    * Called by AutomergeStorageAdapter after levelDB batch commit.
    */
   private async _afterSave(path: StorageKey): Promise<void> {
-    this._indexMetadataStore.notifyMarkedDirty();
+    if (!this.isOpen) {
+      return undefined;
+    }
 
     const documentId = path[0] as DocumentId;
-    const document = this._repo.handles[documentId]?.doc();
-    if (document) {
-      const heads = getHeads(document);
-      this._onHeadsChanged(documentId, heads);
+    const handle = this._repo.handles[documentId];
+    if (!handle || !handle.isReady()) {
+      return;
     }
+    const document = handle.doc();
+    if (!document) {
+      return;
+    }
+
+    const heads = getHeads(document);
+    this._headsUpdates.set(documentId, heads);
+    invariant(this._onHeadsChangedTask, 'onHeadsChangedTask is not initialized');
+    this._onHeadsChangedTask.schedule();
     this.documentsSaved.emit();
   }
 
@@ -391,7 +517,7 @@ export class AutomergeHost extends Resource {
     return this._repo.peers;
   }
 
-  private async _isDocumentInRemoteCollection(params: RemoteDocumentExistenceCheckParams): Promise<boolean> {
+  private async _isDocumentInRemoteCollection(params: RemoteDocumentExistenceCheckProps): Promise<boolean> {
     for (const collectionId of this._collectionSynchronizer.getRegisteredCollectionIds()) {
       const remoteCollections = this._collectionSynchronizer.getRemoteCollectionStates(collectionId);
       const remotePeerDocs = remoteCollections.get(params.peerId as PeerId)?.documents;
@@ -433,11 +559,17 @@ export class AutomergeHost extends Resource {
   async flush({ documentIds }: FlushRequest = {}): Promise<void> {
     // Note: Sync protocol for client and services ensures that all handles should have all changes.
 
-    const loadedDocuments = documentIds?.filter((documentId): documentId is DocumentId => {
-      const handle = this._repo.handles[documentId as DocumentId];
-      return handle && handle.isReady();
-    });
+    const loadedDocuments = (documentIds ?? Object.keys(this._repo.handles)).filter(
+      (documentId): documentId is DocumentId => {
+        const handle = this._repo.handles[documentId as DocumentId];
+        return handle && handle.isReady();
+      },
+    );
     await this._repo.flush(loadedDocuments);
+
+    // Ensure that document verions have propagated accross the system.
+    // This is important for the case where we are doing flush and then waiting for sync to happen.
+    await this._onHeadsChangedTask?.runBlocking();
   }
 
   async getHeads(documentIds: DocumentId[]): Promise<(Heads | undefined)[]> {
@@ -463,6 +595,13 @@ export class AutomergeHost extends Resource {
     return result;
   }
 
+  /**
+   * Iterate over all document heads stored on disk.
+   */
+  listDocumentHeads(): AsyncGenerator<{ documentId: DocumentId; heads: Heads }> {
+    return this._headsStore.iterateAll();
+  }
+
   //
   // Collection sync.
   //
@@ -479,8 +618,8 @@ export class AutomergeHost extends Resource {
     this._collectionSynchronizer.refreshCollection(collectionId);
   }
 
-  async getCollectionSyncState(collectionId: string): Promise<CollectionSyncState> {
-    const result: CollectionSyncState = {
+  async getCollectionSyncState(collectionId: string): Promise<SpaceSyncState> {
+    const result: SpaceSyncState = {
       peers: [],
     };
 
@@ -493,13 +632,16 @@ export class AutomergeHost extends Resource {
 
     for (const [peerId, state] of remoteState) {
       const diff = diffCollectionState(localState, state);
-      result.peers.push({
+      result.peers!.push({
         peerId,
         missingOnRemote: diff.missingOnRemote.length,
         missingOnLocal: diff.missingOnLocal.length,
         differentDocuments: diff.different.length,
-        localDocumentCount: Object.keys(localState.documents).length,
-        remoteDocumentCount: Object.keys(state.documents).length,
+        localDocumentCount: Object.entries(localState.documents).filter(([_, heads]) => heads.length > 0).length,
+        remoteDocumentCount: Object.entries(state.documents).filter(([_, heads]) => heads.length > 0).length,
+
+        totalDocumentCount: new Set([...Object.keys(localState.documents), ...Object.keys(state.documents)]).size,
+        unsyncedDocumentCount: new Set([...diff.missingOnLocal, ...diff.missingOnRemote, ...diff.different]).size,
       });
     }
 
@@ -515,6 +657,16 @@ export class AutomergeHost extends Resource {
       heads.map((heads, index) => [documentIds[index], heads ?? []]),
     );
     this._collectionSynchronizer.setLocalCollectionState(collectionId, { documents });
+
+    // Proactively push our updated local state to peers that are interested in this collection.
+    // This reduces reliance on the next periodic query and prevents replication stalls in fast paths
+    // where the remote queries before our local state is ready.
+    const interestedPeers = this._echoNetworkAdapter.getPeersInterestedInCollection(collectionId);
+    if (interestedPeers.length > 0) {
+      for (const peerId of interestedPeers) {
+        this._sendCollectionState(collectionId, peerId, { documents });
+      }
+    }
   }
 
   async clearLocalCollectionState(collectionId: string): Promise<void> {
@@ -546,6 +698,11 @@ export class AutomergeHost extends Resource {
   }
 
   private _onRemoteCollectionStateUpdated(collectionId: string, peerId: PeerId): void {
+    this._collectionsToSync.add({ collectionId, peerId });
+    this._syncTask?.schedule();
+  }
+
+  private async _handleCollectionSync(collectionId: string, peerId: PeerId) {
     const localState = this._collectionSynchronizer.getLocalCollectionState(collectionId);
     const remoteState = this._collectionSynchronizer.getRemoteCollectionStates(collectionId).get(peerId);
 
@@ -554,36 +711,184 @@ export class AutomergeHost extends Resource {
     }
 
     const { different, missingOnLocal, missingOnRemote } = diffCollectionState(localState, remoteState);
-    const toReplicate = [...missingOnLocal, ...missingOnRemote, ...different];
 
-    if (toReplicate.length === 0) {
+    if (different.length === 0 && missingOnLocal.length === 0 && missingOnRemote.length === 0) {
+      return;
+    }
+
+    const toReplicateWithoutBatching = [...different];
+    const bundleSyncEnabled = this._echoNetworkAdapter.bundleSyncEnabledForPeer(peerId);
+    if (bundleSyncEnabled && missingOnRemote.length >= BUNDLE_SYNC_THRESHOLD) {
+      log('pushing bundle', { amount: missingOnRemote.length });
+      const { syncInteractively } = await this._pushInBundles(peerId, missingOnRemote);
+      toReplicateWithoutBatching.push(...syncInteractively);
+    } else {
+      log.verbose('failed to push bundle, replicating interactively', {
+        collectionId,
+        peerId,
+        amount: missingOnRemote.length,
+      });
+      toReplicateWithoutBatching.push(...missingOnRemote);
+    }
+    if (bundleSyncEnabled && missingOnLocal.length >= BUNDLE_SYNC_THRESHOLD) {
+      log('pulling bundle', { amount: missingOnLocal.length });
+      const { syncInteractively } = await this._pullInBundles(peerId, missingOnLocal);
+      toReplicateWithoutBatching.push(...syncInteractively);
+    } else {
+      log.verbose('failed to pull bundle, replicating interactively', {
+        collectionId,
+        peerId,
+        amount: missingOnLocal.length,
+      });
+      toReplicateWithoutBatching.push(...missingOnLocal);
+    }
+
+    if (toReplicateWithoutBatching.length === 0) {
       return;
     }
 
     log('replicating documents after collection sync', {
       collectionId,
       peerId,
-      toReplicate,
-      count: toReplicate.length,
+      toReplicateWithoutBatching,
+      count: toReplicateWithoutBatching.length,
     });
 
     // Load the documents so they will start syncing.
-    for (const documentId of toReplicate) {
+    for (const documentId of toReplicateWithoutBatching) {
+      // Unless we track the document in "to sync" list, it will not be advertised.
+      this._documentsToSync!.add(documentId);
       this._repo.findWithProgress(documentId);
     }
+    this._sharePolicyChangedTask!.schedule();
   }
 
-  private _onHeadsChanged(documentId: DocumentId, heads: Heads): void {
+  // TODO(mykola): Add retries of batches https://gist.github.com/mykola-vrmchk/fde270259e9209fcbf1331e5abbf12cf
+  private async _pushInBundles(
+    peerId: PeerId,
+    documentIds: DocumentId[],
+  ): Promise<{ syncInteractively: DocumentId[] }> {
+    const documentsToPush = [...documentIds];
+    const syncInteractively: DocumentId[] = [];
+
+    // Push bundles in parallel with BUNDLE_SYNC_CONCURRENCY max concurrent tasks.
+    while (documentsToPush.length > 0) {
+      await Promise.all(
+        range(BUNDLE_SYNC_CONCURRENCY).map(async () => {
+          const bundle = documentsToPush.splice(0, BUNDLE_SIZE);
+          if (bundle.length === 0) {
+            return;
+          }
+          await this._pushBundle(peerId, bundle).catch((err) => {
+            log.warn('failed to push bundle, replicating interactively', { peerId, bundle, err });
+            syncInteractively.push(...bundle);
+          });
+        }),
+      );
+    }
+
+    return { syncInteractively };
+  }
+
+  private async _pushBundle(peerId: PeerId, documentIds: DocumentId[]): Promise<void> {
+    if (this._ctx.disposed) {
+      return;
+    }
+
+    const docs = documentIds.map((documentId) => {
+      const handle = this._repo.handles[documentId];
+      if (!handle || !handle.isReady()) {
+        log.warn('document not ready, skipping', { documentId });
+        return;
+      }
+      const doc = handle.doc();
+      if (!doc) {
+        log.warn('document not available, skipping', { documentId });
+        return;
+      }
+      return {
+        documentId,
+        data: save(doc),
+        heads: getHeads(doc),
+      };
+    });
+
+    await this._echoNetworkAdapter.pushBundle(peerId, docs.filter(isNonNullable));
+  }
+
+  private async _pullInBundles(
+    peerId: PeerId,
+    documentIds: DocumentId[],
+  ): Promise<{ syncInteractively: DocumentId[] }> {
+    const documentsToPull = [...documentIds];
+    const syncInteractively: DocumentId[] = [];
+    const docsToImport: Record<DocumentId, Uint8Array> = {};
+
+    // Pull bundles in parallel with BUNDLE_SYNC_CONCURRENCY max concurrent tasks.
+    while (documentsToPull.length > 0) {
+      await Promise.all(
+        range(BUNDLE_SYNC_CONCURRENCY).map(async () => {
+          const bundle = documentsToPull.splice(0, BUNDLE_SIZE);
+          if (bundle.length === 0) {
+            return;
+          }
+          const result = await this._pullBundle(peerId, bundle).catch((err) => {
+            log.warn('failed to pull bundle, replicating interactively', { peerId, bundle, err });
+            syncInteractively.push(...bundle);
+          });
+          if (result) {
+            Object.assign(docsToImport, result.docsToImport);
+          }
+        }),
+      );
+    }
+
+    for (const [documentId, data] of Object.entries(docsToImport)) {
+      this._repo.import(data, { docId: documentId as DocumentId });
+    }
+    await this._repo.flush(Object.keys(docsToImport) as DocumentId[]);
+
+    return { syncInteractively };
+  }
+
+  private async _pullBundle(
+    peerId: PeerId,
+    documentIds: DocumentId[],
+  ): Promise<{ docsToImport: Record<DocumentId, Uint8Array> } | undefined> {
+    if (this._ctx.disposed) {
+      return;
+    }
+    // NOTE: We are expecting that documents that are being pulled are not present locally, so we are pulling all changes.
+    const docHeads = Object.fromEntries(documentIds.map((documentId) => [documentId, []]));
+    const bundle = await this._echoNetworkAdapter.pullBundle(peerId, docHeads);
+    return { docsToImport: bundle };
+  }
+
+  private _onHeadsChanged(docHeads: [DocumentId, Heads][]): void {
     const collectionsChanged = new Set<CollectionId>();
+
     for (const collectionId of this._collectionSynchronizer.getRegisteredCollectionIds()) {
       const state = this._collectionSynchronizer.getLocalCollectionState(collectionId);
-      if (state?.documents[documentId]) {
-        const newState = structuredClone(state);
-        newState.documents[documentId] = heads;
+      if (!state) {
+        continue;
+      }
+      let newState: CollectionState | undefined;
+
+      for (const [documentId, heads] of docHeads) {
+        if (documentId in state.documents) {
+          if (!newState) {
+            newState = structuredClone(state);
+          }
+          newState.documents[documentId] = heads;
+        }
+      }
+
+      if (newState) {
         this._collectionSynchronizer.setLocalCollectionState(collectionId, newState);
         collectionsChanged.add(collectionId as CollectionId);
       }
     }
+
     for (const collectionId of collectionsChanged) {
       this.collectionStateUpdated.emit({ collectionId });
     }
@@ -619,35 +924,3 @@ const decodeCollectionState = (state: unknown): CollectionState => {
 const encodeCollectionState = (state: CollectionState): unknown => {
   return state;
 };
-
-export type CollectionSyncState = {
-  peers: PeerSyncState[];
-};
-
-export type PeerSyncState = {
-  peerId: PeerId;
-  /**
-   * Documents that are present locally but not on the remote peer.
-   */
-  missingOnRemote: number;
-
-  /**
-   * Documents that are present on the remote peer but not locally.
-   */
-  missingOnLocal: number;
-
-  /**
-   * Documents that are present on both peers but have different heads.
-   */
-  differentDocuments: number;
-
-  /**
-   * Total number of documents locally.
-   */
-  localDocumentCount: number;
-
-  /**
-   * Total number of documents on the remote peer.
-   */
-  remoteDocumentCount: number;
-};