cojson 0.20.9 → 0.20.11

package/src/sync.ts

@@ -1,9 +1,13 @@
+import { base58 } from "@scure/base";
 import { md5 } from "@noble/hashes/legacy";
 import { Histogram, ValueType, metrics } from "@opentelemetry/api";
 import { PeerState } from "./PeerState.js";
 import { SyncStateManager } from "./SyncStateManager.js";
 import { UnsyncedCoValuesTracker } from "./UnsyncedCoValuesTracker.js";
-import { SYNC_SCHEDULER_CONFIG } from "./config.js";
+import {
+  STORAGE_RECONCILIATION_CONFIG,
+  SYNC_SCHEDULER_CONFIG,
+} from "./config.js";
 import {
   getContenDebugInfo,
   getNewTransactionsFromContentMessage,
@@ -21,6 +25,8 @@ import { CoValuePriority } from "./priority.js";
 import { IncomingMessagesQueue } from "./queue/IncomingMessagesQueue.js";
 import { LocalTransactionsSyncQueue } from "./queue/LocalTransactionsSyncQueue.js";
 import type { StorageStreamingQueue } from "./queue/StorageStreamingQueue.js";
+import { OngoingStorageReconciliationTracker } from "./OngoingStorageReconciliationTracker.js";
+import { StorageReconciliationServerAckTracker } from "./StorageReconciliationAckTracker.js";
 import {
   CoValueKnownState,
   knownStateFrom,
@@ -33,7 +39,9 @@ export type SyncMessage =
   | LoadMessage
   | KnownStateMessage
   | NewContentMessage
-  | DoneMessage;
+  | DoneMessage
+  | ReconcileMessage
+  | ReconcileAckMessage;
 
 export type LoadMessage = {
   action: "load";
@@ -68,6 +76,19 @@ export type DoneMessage = {
   id: RawCoID;
 };
 
+export type ReconcileBatchID = string;
+
+export type ReconcileMessage = {
+  action: "reconcile";
+  id: ReconcileBatchID;
+  values: [coValue: RawCoID, sessionsHash: string][];
+};
+
+export type ReconcileAckMessage = {
+  action: "reconcile-ack";
+  id: ReconcileBatchID;
+};
+
 /**
  * Determines when network sync is enabled.
  * - "always": sync is enabled for both Anonymous Authentication and Authenticated Account
@@ -111,9 +132,29 @@ export type ServerPeerSelector = (
   serverPeers: PeerState[],
 ) => PeerState[];
 
+/**
+ * Manages the sync of coValues between peers.
+ * It is responsible for sending, receiving and processing sync messages.
+ * For more details on how the sync protocol works, see the sync protocol documentation:
+ * {@link docs/sync-protocol.md}
+ */
 export class SyncManager {
   peers: { [key: PeerID]: PeerState } = {};
   local: LocalNode;
+  /**
+   * Tracks pending reconcile acks from the server.
+   */
+  private reconciliationAckTracker =
+    new StorageReconciliationServerAckTracker();
+  /**
+   * Tracks ongoing storage reconciliation batches in a server.
+   */
+  private ongoingStorageReconciliationTracker =
+    new OngoingStorageReconciliationTracker();
+
+  get pendingReconciliationAck(): Map<string, number> {
+    return this.reconciliationAckTracker.pendingReconciliationAck;
+  }
 
   // When true, transactions will not be verified.
   // This is useful when syncing only for storage purposes, with the expectation that
@@ -127,6 +168,8 @@ export class SyncManager {
     this._ignoreUnknownCoValuesFromServers = true;
   }
 
+  fullStorageReconciliationEnabled = false;
+
   peersCounter = metrics.getMeter("cojson").createUpDownCounter("jazz.peers", {
     description: "Amount of connected peers",
     valueType: ValueType.INT,
@@ -179,6 +222,15 @@ export class SyncManager {
   }
 
   handleSyncMessage(msg: SyncMessage, peer: PeerState) {
+    if (msg.action === "reconcile") {
+      this.handleReconcile(msg, peer);
+      return;
+    }
+    if (msg.action === "reconcile-ack") {
+      this.handleReconcileAck(msg, peer);
+      return;
+    }
+
     if (!isRawCoID(msg.id)) {
       const errorType = msg.id ? "invalid" : "undefined";
       logger.warn(`Received sync message with ${errorType} id`, {
@@ -233,7 +285,20 @@ export class SyncManager {
     }
   }
 
-  sendNewContent(id: RawCoID, peer: PeerState, seen: Set<RawCoID> = new Set()) {
+  sendNewContent(
+    id: RawCoID,
+    peer: PeerState,
+    forceKnownReplyOnNoDelta: boolean = false,
+  ) {
+    this.#sendNewContent(id, peer, new Set(), forceKnownReplyOnNoDelta);
+  }
+
+  #sendNewContent(
+    id: RawCoID,
+    peer: PeerState,
+    seen: Set<RawCoID>,
+    forceKnownReplyOnNoDelta: boolean,
+  ) {
     if (seen.has(id)) {
       return;
     }
@@ -249,7 +314,7 @@ export class SyncManager {
     const includeDependencies = peer.role !== "server";
     if (includeDependencies) {
       for (const dependency of coValue.getDependedOnCoValues()) {
-        this.sendNewContent(dependency, peer, seen);
+        this.#sendNewContent(dependency, peer, seen, false);
       }
     }
 
@@ -263,7 +328,7 @@ export class SyncManager {
       }
 
       peer.combineOptimisticWith(id, coValue.knownState());
-    } else if (!peer.toldKnownState.has(id)) {
+    } else if (forceKnownReplyOnNoDelta || !peer.toldKnownState.has(id)) {
       if (coValue.isDeleted) {
         // This way we make the peer believe that we've always ingested all the content they sent, even though we skipped it because the coValue is deleted
         this.trySendToPeer(
@@ -281,15 +346,172 @@ export class SyncManager {
     peer.trackToldKnownState(id);
   }
 
+  /**
+   * Reconciles all in-memory CoValues with all persistent server peers
+   */
   reconcileServerPeers() {
     const serverPeers = Object.values(this.peers).filter(
-      (peer) => peer.role === "server",
+      isPersistentServerPeer,
     );
     for (const peer of serverPeers) {
       this.startPeerReconciliation(peer);
     }
   }
 
+  /**
+   * Ensures all CoValues in storage are synced to the given server peer.
+   * Sends "reconcile" message(s) with [coValueId, sessionsHash] for each CoValue.
+   * Server responds with "known" only where it is missing the CoValue or has different sessions,
+   * so that client can send missing content.
+   * Processes CoValues in batches of RECONCILIATION_BATCH_SIZE.
+   * @param peer - The server peer to reconcile with.
+   * @param initialOffset - Offset to start from (for resuming after interrupt). Default 0.
+   * @param onComplete - Called when reconciliation is fully complete (all batches sent and acked).
+   */
+  startStorageReconciliation(
+    peer: PeerState,
+    initialOffset?: number,
+    onComplete?: () => void,
+  ): void {
+    if (!this.local.storage) return;
+    if (!isPersistentServerPeer(peer)) return;
+
+    const startOffset = initialOffset ?? 0;
+    const batchSize = STORAGE_RECONCILIATION_CONFIG.BATCH_SIZE;
+
+    const storage = this.local.storage;
+
+    storage.getCoValueCount((totalCoValueCount) => {
+      const sendReconcileMessage = (
+        batchId: string,
+        entries: [RawCoID, string][],
+        offset: number,
+      ) => {
+        if (entries.length === 0) return;
+
+        this.reconciliationAckTracker.trackBatch(
+          batchId,
+          peer.id,
+          offset + batchSize,
+        );
+
+        this.trySendToPeer(peer, {
+          action: "reconcile",
+          id: batchId,
+          values: entries,
+        });
+      };
+
+      const triggerNextBatch = (lastBatchLength: number, offset: number) => {
+        // This value becomes false when the last covalueid batch picked from the storage
+        // is smaller than the batch size.
+        if (lastBatchLength === batchSize) {
+          logger.info("Reconciled CoValues in storage", {
+            peerId: peer.id,
+            completed: offset + batchSize,
+            total: totalCoValueCount,
+          });
+          processStorageBatch(offset + batchSize);
+        } else {
+          // Note: `completed` can be higher than `total` if CoValues were added
+          // after the reconciliation started
+          logger.info("Storage reconciliation complete", {
+            peerId: peer.id,
+            startOffset,
+            completed: offset + lastBatchLength,
+            total: totalCoValueCount,
+          });
+          onComplete?.();
+        }
+      };
+
+      const processStorageBatch = (offset: number) => {
+        storage.getCoValueIDs(batchSize, offset, (batch) => {
+          this.buildStorageReconciliationEntries(batch, (entries) => {
+            if (entries.length === 0) {
+              triggerNextBatch(batch.length, offset);
+              return;
+            }
+
+            const batchId = base58.encode(this.local.crypto.randomBytes(12));
+            sendReconcileMessage(batchId, entries, offset);
+
+            this.reconciliationAckTracker.waitForAck(batchId, peer, () => {
+              triggerNextBatch(batch.length, offset);
+            });
+          });
+        });
+      };
+
+      logger.info("Starting storage reconciliation", {
+        peerId: peer.id,
+        startOffset,
+        total: totalCoValueCount,
+      });
+      processStorageBatch(startOffset);
+    });
+  }
+
+  private buildStorageReconciliationEntries(
+    batch: { id: RawCoID }[],
+    callback: (entries: [RawCoID, string][]) => void,
+  ): void {
+    const storage = this.local.storage;
+
+    if (!storage) {
+      callback([]);
+      return;
+    }
+
+    const pending = batch.filter(({ id }) => !this.local.isCoValueInMemory(id));
+
+    if (pending.length === 0) {
+      callback([]);
+      return;
+    }
+
+    let done = 0;
+    const entries: [RawCoID, string][] = [];
+
+    for (const coValue of pending) {
+      storage.loadKnownState(coValue.id, (storageKnownState) => {
+        if (storageKnownState) {
+          entries.push([
+            coValue.id,
+            this.hashKnownStateSessions(storageKnownState.sessions),
+          ]);
+        }
+
+        done += 1;
+        if (done === pending.length) {
+          callback(entries);
+        }
+      });
+    }
+  }
+
+  private maybeStartStorageReconciliationForPeer(peer: PeerState): void {
+    if (!this.fullStorageReconciliationEnabled) return;
+    if (!this.local.storage) return;
+
+    const sessionId = this.local.currentSessionID;
+    this.local.storage.tryAcquireStorageReconciliationLock(
+      sessionId,
+      peer.id,
+      (result) => {
+        if (!result.acquired) return;
+
+        const lastProcessedOffset = result.lastProcessedOffset;
+        this.startStorageReconciliation(peer, lastProcessedOffset, () => {
+          this.local.storage?.releaseStorageReconciliationLock(
+            sessionId,
+            peer.id,
+          );
+        });
+      },
+    );
+  }
+
   async resumeUnsyncedCoValues(): Promise<void> {
     if (!this.local.storage) {
       // No storage available, skip resumption
@@ -364,12 +586,19 @@ export class SyncManager {
     });
   }
 
+  /**
+   * Reconciles all in-memory CoValues with the given peer.
+   * Creates a subscription for each CoValue that is not already subscribed to.
+   */
   startPeerReconciliation(peer: PeerState) {
     if (isPersistentServerPeer(peer)) {
       // Resume syncing unsynced CoValues asynchronously
       this.resumeUnsyncedCoValues().catch((error) => {
         logger.warn("Failed to resume unsynced CoValues:", error);
       });
+
+      // Try to run full storage reconciliation for this peer (scheduled per peer, every 30 days)
+      this.maybeStartStorageReconciliationForPeer(peer);
     }
 
     const coValuesOrderedByDependency: CoValueCore[] = [];
@@ -549,6 +778,7 @@ export class SyncManager {
 
     peerState.addCloseListener(() => {
       unsubscribeFromKnownStatesUpdates();
+      this.ongoingStorageReconciliationTracker.clearPeer(peer.id);
       this.peersCounter.add(-1, { role: peer.role });
 
       if (!peer.persistent && this.peers[peer.id] === peerState) {
@@ -592,7 +822,7 @@ export class SyncManager {
 
     // Fast path: CoValue is already in memory
     if (coValue.isAvailable()) {
-      this.sendNewContent(msg.id, peer);
+      this.sendNewContent(msg.id, peer, true);
       return;
     }
 
@@ -608,7 +838,7 @@ export class SyncManager {
     coValue.getKnownStateFromStorage((storageKnownState) => {
       // Race condition: CoValue might have been loaded while we were waiting for storage
       if (coValue.isAvailable()) {
-        this.sendNewContent(msg.id, peer);
+        this.sendNewContent(msg.id, peer, true);
         return;
       }
 
@@ -631,7 +861,7 @@ export class SyncManager {
         // Even though we responded with KNOWN (client has everything), we need
         // to establish a subscription so that updates from core flow to us.
         const serverPeers = this.getServerPeers(msg.id, peer.id);
-        coValue.loadFromPeers(serverPeers);
+        coValue.loadFromPeers(serverPeers, "low-priority");
 
         return;
       }
@@ -652,7 +882,7 @@ export class SyncManager {
   ) {
     coValue.loadFromStorage((found) => {
       if (found && coValue.isAvailable()) {
-        this.sendNewContent(id, peer);
+        this.sendNewContent(id, peer, true);
       } else {
         this.loadFromPeersAndRespond(id, peer, coValue);
       }
@@ -668,7 +898,7 @@ export class SyncManager {
     coValue: CoValueCore,
   ) {
     const peers = this.getServerPeers(id, peer.id);
-    coValue.loadFromPeers(peers);
+    coValue.loadFromPeers(peers, "immediate");
 
     const handleLoadResult = () => {
       if (coValue.isAvailable()) {
@@ -734,9 +964,107 @@ export class SyncManager {
 
     if (coValue.isAvailable()) {
       this.sendNewContent(msg.id, peer);
+    } else if (coValue.isKnownStateAvailable()) {
+      // Validate if content is missing before loading it from storage
+      if (!this.syncState.isSynced(peer, msg.id)) {
+        this.local.loadCoValueCore(msg.id).then(() => {
+          this.sendNewContent(msg.id, peer);
+        });
+      }
     }
 
-    peer.trackLoadRequestComplete(coValue);
+    peer.trackLoadRequestComplete(coValue, "known");
+    this.maybeMarkCoValueAsReconciled(peer, msg.id);
+  }
+
+  handleReconcile(msg: ReconcileMessage, peer: PeerState): void {
+    let remaining = msg.values.length;
+    if (remaining === 0) {
+      this.trySendToPeer(peer, { action: "reconcile-ack", id: msg.id });
+      return;
+    }
+
+    const pending = new Set<RawCoID>();
+    const processEntryDone = () => {
+      remaining -= 1;
+
+      if (remaining !== 0) {
+        return;
+      }
+
+      if (pending.size === 0) {
+        this.trySendToPeer(peer, { action: "reconcile-ack", id: msg.id });
+        return;
+      }
+
+      this.ongoingStorageReconciliationTracker.trackBatch(
+        peer.id,
+        msg.id,
+        pending,
+      );
+    };
+
+    for (const [coValueId, clientSessionsHash] of msg.values) {
+      // Avoid creating a new coValue object if it's not already in memory
+      const inMemoryCoValue = this.local.isCoValueInMemory(coValueId)
+        ? this.local.getCoValue(coValueId)
+        : undefined;
+      if (inMemoryCoValue?.isErroredInPeer(peer.id)) {
+        processEntryDone();
+        continue;
+      }
+
+      const maybeSendLoadRequest = (
+        knownState: CoValueKnownState | undefined,
+      ) => {
+        if (!knownState) {
+          pending.add(coValueId);
+          peer.trackToldKnownState(coValueId);
+          this.trySendToPeer(peer, {
+            action: "load",
+            id: coValueId,
+            header: false,
+            sessions: {},
+          });
+        } else {
+          const serverSessionsHash = this.hashKnownStateSessions(
+            knownState.sessions,
+          );
+          if (serverSessionsHash !== clientSessionsHash) {
+            pending.add(coValueId);
+            peer.trackToldKnownState(coValueId);
+            this.trySendToPeer(peer, { action: "load", ...knownState });
+          }
+        }
+        processEntryDone();
+      };
+
+      if (
+        inMemoryCoValue?.isAvailable() ||
+        inMemoryCoValue?.loadingState === "onlyKnownState"
+      ) {
+        maybeSendLoadRequest(inMemoryCoValue.knownState());
+      } else {
+        this.local.storage
+          ? this.local.storage.loadKnownState(coValueId, maybeSendLoadRequest)
+          : maybeSendLoadRequest(undefined);
+      }
+    }
+  }
+
+  handleReconcileAck(msg: ReconcileAckMessage, peer: PeerState): void {
+    const nextOffset = this.reconciliationAckTracker.handleAck(msg.id, peer.id);
+    if (nextOffset !== undefined) {
+      this.local.storage?.renewStorageReconciliationLock(
+        this.local.currentSessionID,
+        peer.id,
+        nextOffset,
+      );
+    }
+  }
+
+  private hashKnownStateSessions(sessions: KnownStateSessions): string {
+    return this.local.crypto.shortHash(sessions);
   }
 
   recordTransactionsSize(newTransactions: Transaction[], source: string) {
@@ -1040,7 +1368,10 @@ export class SyncManager {
       }
     }
 
-    peer?.trackLoadRequestComplete(coValue);
+    peer?.trackLoadRequestComplete(coValue, "content");
+    if (peer && !coValue.isStreaming()) {
+      this.maybeMarkCoValueAsReconciled(peer, msg.id);
+    }
 
     for (const peer of this.getPeers(coValue.id)) {
       /**
@@ -1056,7 +1387,7 @@ export class SyncManager {
       if (peer.isCoValueSubscribedToPeer(coValue.id)) {
         this.sendNewContent(coValue.id, peer);
       } else if (peer.role === "server") {
-        peer.sendLoadRequest(coValue);
+        peer.sendLoadRequest(coValue, "low-priority");
       }
     }
   }
@@ -1067,6 +1398,20 @@ export class SyncManager {
     return this.sendNewContent(msg.id, peer);
   }
 
+  private maybeMarkCoValueAsReconciled(peer: PeerState, coValueId: RawCoID) {
+    const completedBatchIds =
+      this.ongoingStorageReconciliationTracker.markItemComplete(
+        peer.id,
+        coValueId,
+      );
+    for (const batchId of completedBatchIds) {
+      this.trySendToPeer(peer, {
+        action: "reconcile-ack",
+        id: batchId,
+      });
+    }
+  }
+
   private syncQueue = new LocalTransactionsSyncQueue((content) =>
     this.syncContent(content),
   );

package/src/tests/OngoingStorageReconciliationTracker.test.ts

@@ -0,0 +1,85 @@
+import { describe, expect, test } from "vitest";
+import { OngoingStorageReconciliationTracker } from "../OngoingStorageReconciliationTracker.js";
+import { RawCoID } from "../ids.js";
+
+function coID(value: string): RawCoID {
+  return value as RawCoID;
+}
+
+describe("OngoingStorageReconciliationTracker", () => {
+  test("does not track empty batches", () => {
+    const tracker = new OngoingStorageReconciliationTracker();
+
+    tracker.trackBatch("peer-1", "batch-1", new Set());
+
+    // @ts-expect-error - reconcileBatches is private
+    expect(tracker.reconcileBatches.size).toEqual(0);
+  });
+
+  test("completes a batch only when all its covalues are completed", () => {
+    const tracker = new OngoingStorageReconciliationTracker();
+
+    tracker.trackBatch(
+      "peer-1",
+      "batch-1",
+      new Set([coID("co_A"), coID("co_B")]),
+    );
+
+    expect(tracker.markItemComplete("peer-1", coID("co_A"))).toEqual([]);
+    expect(tracker.markItemComplete("peer-1", coID("co_B"))).toEqual([
+      "batch-1",
+    ]);
+    // @ts-expect-error - reconcileBatches is private
+    expect(tracker.reconcileBatches.size).toEqual(0);
+  });
+
+  test("returns all completed batch ids for a covalue that is in multiple batches", () => {
+    const tracker = new OngoingStorageReconciliationTracker();
+
+    tracker.trackBatch("peer-1", "batch-1", new Set([coID("co_A")]));
+    tracker.trackBatch(
+      "peer-1",
+      "batch-2",
+      new Set([coID("co_A"), coID("co_B")]),
+    );
+    tracker.trackBatch("peer-1", "batch-3", new Set([coID("co_B")]));
+
+    expect(tracker.markItemComplete("peer-1", coID("co_A"))).toEqual([
+      "batch-1",
+    ]);
+    expect(tracker.markItemComplete("peer-1", coID("co_B"))).toEqual([
+      "batch-2",
+      "batch-3",
+    ]);
+    // @ts-expect-error - reconcileBatches is private
+    expect(tracker.reconcileBatches.size).toEqual(0);
+  });
+
+  test("isolates state by peer", () => {
+    const tracker = new OngoingStorageReconciliationTracker();
+
+    tracker.trackBatch("peer-1", "batch-A", new Set([coID("co_A")]));
+    tracker.trackBatch("peer-2", "batch-B", new Set([coID("co_A")]));
+
+    expect(tracker.markItemComplete("peer-1", coID("co_A"))).toEqual([
+      "batch-A",
+    ]);
+    expect(tracker.markItemComplete("peer-2", coID("co_A"))).toEqual([
+      "batch-B",
+    ]);
+  });
+
+  test("clearPeer removes only that peer state", () => {
+    const tracker = new OngoingStorageReconciliationTracker();
+
+    tracker.trackBatch("peer-1", "batch-A", new Set([coID("co_A")]));
+    tracker.trackBatch("peer-2", "batch-B", new Set([coID("co_B")]));
+
+    tracker.clearPeer("peer-1");
+
+    expect(tracker.markItemComplete("peer-1", coID("co_A"))).toEqual([]);
+    expect(tracker.markItemComplete("peer-2", coID("co_B"))).toEqual([
+      "batch-B",
+    ]);
+  });
+});