@powersync/service-core 0.1.3 → 0.2.1

package/src/storage/ChecksumCache.ts

@@ -0,0 +1,294 @@
+import { BucketChecksum, OpId } from '@/util/protocol-types.js';
+import { ChecksumMap, addBucketChecksums } from '@/util/utils.js';
+import { LRUCache } from 'lru-cache/min';
+import { OrderedSet } from '@js-sdsl/ordered-set';
+
+interface ChecksumFetchContext {
+  fetch(bucket: string): Promise<BucketChecksum>;
+  checkpoint: bigint;
+}
+
+export interface FetchPartialBucketChecksum {
+  bucket: string;
+  start?: OpId;
+  end: OpId;
+}
+
+export type FetchChecksums = (batch: FetchPartialBucketChecksum[]) => Promise<ChecksumMap>;
+
+export interface ChecksumCacheOptions {
+  /**
+   * Upstream checksum implementation.
+   *
+   * This fetches a batch of either entire bucket checksums, or a partial range.
+   */
+  fetchChecksums: FetchChecksums;
+
+  /**
+   * Maximum number of cached checksums.
+   */
+  maxSize?: number;
+}
+
+// Approximately 5MB of memory, if we assume 50 bytes per entry
+const DEFAULT_MAX_SIZE = 100_000;
+
+/**
+ * Implement a LRU cache for checksum requests. Each (bucket, checkpoint) request is cached separately,
+ * while the lookups occur in batches.
+ *
+ * For each bucket, we keep a separate OrderedSet of cached checkpoints.
+ * This allows us to do incrementally update checksums by using the last cached checksum for the same bucket.
+ *
+ * We use the LRUCache fetchMethod to deduplicate in-progress requests.
+ */
+export class ChecksumCache {
+  /**
+   * The primary checksum cache, with key of `${checkpoint}/${bucket}`.
+   */
+  private cache: LRUCache<string, BucketChecksum, ChecksumFetchContext>;
+  /**
+   * For each bucket, an ordered set of cached checkpoints.
+   */
+  private bucketCheckpoints = new Map<string, OrderedSet<bigint>>();
+
+  private fetchChecksums: FetchChecksums;
+
+  constructor(options: ChecksumCacheOptions) {
+    this.fetchChecksums = options.fetchChecksums;
+
+    this.cache = new LRUCache<string, BucketChecksum, ChecksumFetchContext>({
+      max: options.maxSize ?? DEFAULT_MAX_SIZE,
+      fetchMethod: async (cacheKey, _staleValue, options) => {
+        // Called when this checksum hasn't been cached yet.
+        // Pass the call back to the request, which implements batch fetching.
+        const { bucket } = parseCacheKey(cacheKey);
+        const result = await options.context.fetch(bucket);
+
+        // Add to the set of cached checkpoints for the bucket.
+        let checkpointSet = this.bucketCheckpoints.get(bucket);
+        if (checkpointSet == null) {
+          checkpointSet = new OrderedSet();
+          this.bucketCheckpoints.set(bucket, checkpointSet);
+        }
+        checkpointSet.insert(options.context.checkpoint);
+        return result;
+      },
+
+      dispose: (value, key) => {
+        // Remove from the set of cached checkpoints for the bucket
+        const { checkpointString } = parseCacheKey(key);
+        const checkpoint = BigInt(checkpointString);
+        const checkpointSet = this.bucketCheckpoints.get(value.bucket);
+        if (checkpointSet == null) {
+          return;
+        }
+        checkpointSet.eraseElementByKey(checkpoint);
+        if (checkpointSet.length == 0) {
+          this.bucketCheckpoints.delete(value.bucket);
+        }
+      },
+
+      noDisposeOnSet: true,
+
+      // When we have more fetches than the cache size, complete the fetches instead
+      // of failing with Error('evicted').
+      ignoreFetchAbort: true
+    });
+  }
+
+  async getChecksums(checkpoint: OpId, buckets: string[]): Promise<BucketChecksum[]> {
+    const checksums = await this.getChecksumMap(checkpoint, buckets);
+    // Return results in the same order as the request
+    return buckets.map((bucket) => checksums.get(bucket)!);
+  }
+
+  /**
+   * Get bucket checksums for a checkpoint.
+   *
+   * Any checksums not found upstream are returned as zero checksums.
+   *
+   * @returns a Map with exactly one entry for each bucket requested
+   */
+  async getChecksumMap(checkpoint: OpId, buckets: string[]): Promise<ChecksumMap> {
+    // Buckets that don't have a cached checksum for this checkpoint yet
+    let toFetch = new Set<string>();
+
+    // Newly fetched results
+    let fetchResults = new Map<string, BucketChecksum>();
+
+    // Promise for the bactch new fetch requests
+    let resolveFetch!: () => void;
+    let rejectFetch!: (err: any) => void;
+    let fetchPromise = new Promise<void>((resolve, reject) => {
+      resolveFetch = resolve;
+      rejectFetch = reject;
+    });
+
+    // Accumulated results - both from cached checksums, and fetched checksums
+    let finalResults = new Map<string, BucketChecksum>();
+
+    const context: ChecksumFetchContext = {
+      async fetch(bucket) {
+        await fetchPromise;
+        if (!toFetch.has(bucket)) {
+          // Should never happen
+          throw new Error(`Expected to fetch ${bucket}`);
+        }
+        const checksum = fetchResults.get(bucket);
+        if (checksum == null) {
+          // Should never happen
+          throw new Error(`Failed to fetch checksum for bucket ${bucket}`);
+        }
+        return checksum;
+      },
+      checkpoint: BigInt(checkpoint)
+    };
+
+    // One promise to await to ensure all fetch requests completed.
+    let settledPromise: Promise<PromiseSettledResult<void>[]> | null = null;
+
+    try {
+      // Individual cache fetch promises
+      let cacheFetchPromises: Promise<void>[] = [];
+
+      for (let bucket of buckets) {
+        const cacheKey = makeCacheKey(checkpoint, bucket);
+        let status: LRUCache.Status<BucketChecksum> = {};
+        const p = this.cache.fetch(cacheKey, { context: context, status: status }).then((checksums) => {
+          if (checksums == null) {
+            // Should never happen
+            throw new Error(`Failed to get checksums for ${cacheKey}`);
+          }
+          finalResults.set(bucket, checksums);
+        });
+        cacheFetchPromises.push(p);
+        if (status.fetch == 'hit' || status.fetch == 'inflight') {
+          // The checksums is either cached already (hit), or another request is busy
+          // fetching (inflight).
+          // In either case, we don't need to fetch a new checksum.
+        } else {
+          // We need a new request for this checksum.
+          toFetch.add(bucket);
+        }
+      }
+      // We do this directly after creating the promises, otherwise
+      // we could end up with weird uncaught rejection errors.
+      settledPromise = Promise.allSettled(cacheFetchPromises);
+
+      if (toFetch.size == 0) {
+        // Nothing to fetch, but resolve in case
+        resolveFetch();
+      } else {
+        let bucketRequests: FetchPartialBucketChecksum[] = [];
+        // Partial checksum (previously cached) to add to the partial fetch
+        let add = new Map<string, BucketChecksum>();
+
+        for (let bucket of toFetch) {
+          let bucketRequest: FetchPartialBucketChecksum | null = null;
+          const checkpointSet = this.bucketCheckpoints.get(bucket);
+          if (checkpointSet != null) {
+            // Find smaller checkpoints, sorted in descending order
+            let iter = checkpointSet.reverseUpperBound(context.checkpoint);
+            const begin = checkpointSet.begin();
+            while (iter.isAccessible()) {
+              const cp = iter.pointer;
+              const cacheKey = makeCacheKey(cp, bucket);
+              // peek to avoid refreshing the key
+              const cached = this.cache.peek(cacheKey);
+              // As long as dispose() works correctly, the checkpointset should
+              // match up with the cache, and `cached` should also have a value here.
+              // However, we handle caces where it's not present either way.
+              // Test by disabling the `dispose()` callback.
+              if (cached != null) {
+                // Partial checksum found - make a partial checksum request
+                bucketRequest = {
+                  bucket,
+                  start: cp.toString(),
+                  end: checkpoint
+                };
+                add.set(bucket, cached);
+                break;
+              }
+
+              if (iter.equals(begin)) {
+                // Cannot iterate further
+                break;
+              }
+              // Iterate backwards
+              iter = iter.pre();
+            }
+          }
+
+          if (bucketRequest == null) {
+            // No partial checksum found - make a new full checksum request
+            bucketRequest = {
+              bucket,
+              end: checkpoint
+            };
+            add.set(bucket, {
+              bucket,
+              checksum: 0,
+              count: 0
+            });
+          }
+          bucketRequests.push(bucketRequest);
+        }
+
+        // Fetch partial checksums from upstream
+        const results = await this.fetchChecksums(bucketRequests);
+
+        for (let bucket of toFetch) {
+          const result = results.get(bucket);
+          const toAdd = add.get(bucket);
+          if (toAdd == null) {
+            // Should never happen
+            throw new Error(`toAdd null for ${bucket}`);
+          }
+          // Compute the full checksum from the two partials.
+          // No results returned are treated the same as a zero result.
+          const added = addBucketChecksums(toAdd, result ?? null);
+          fetchResults.set(bucket, added);
+        }
+
+        // fetchResults is fully populated, so we resolve the Promise
+        resolveFetch();
+      }
+    } catch (e) {
+      // Failure when fetching checksums - reject the Promise.
+      // This will reject all individual cache fetch requests, and each will be retried
+      // on the next request.
+      rejectFetch(e);
+
+      // Wait for the above rejection to propagate, otherwise we end up with "uncaught" errors.
+      // This promise never throws.
+      await settledPromise;
+
+      throw e;
+    }
+
+    // Wait for all cache fetch reqeusts to complete
+    const settledResults = (await settledPromise) ?? [];
+    // Check if any of them failed
+    for (let result of settledResults) {
+      if (result.status == 'rejected') {
+        throw result.reason;
+      }
+    }
+
+    if (finalResults.size != buckets.length) {
+      // Should not happen
+      throw new Error(`Bucket results mismatch: ${finalResults.size} != ${buckets.length}`);
+    }
+    return finalResults;
+  }
+}
+
+function makeCacheKey(checkpoint: bigint | string, bucket: string) {
+  return `${checkpoint}/${bucket}`;
+}
+
+function parseCacheKey(key: string) {
+  const index = key.indexOf('/');
+  return { checkpointString: key.substring(0, index), bucket: key.substring(index + 1) };
+}

package/src/storage/mongo/MongoSyncBucketStorage.ts

@@ -22,9 +22,15 @@ import { PowerSyncMongo } from './db.js';
 import { BucketDataDocument, BucketDataKey, SourceKey, SyncRuleState } from './models.js';
 import { MongoBucketBatch } from './MongoBucketBatch.js';
 import { BSON_DESERIALIZE_OPTIONS, idPrefixFilter, readSingleBatch, serializeLookup } from './util.js';
+import { ChecksumCache, FetchPartialBucketChecksum } from '../ChecksumCache.js';
 
 export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
   private readonly db: PowerSyncMongo;
+  private checksumCache = new ChecksumCache({
+    fetchChecksums: (batch) => {
+      return this.getChecksumsInternal(batch);
+    }
+  });
 
   constructor(
     public readonly factory: MongoBucketStorage,
@@ -316,23 +322,28 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
     }
   }
 
-  async getChecksums(checkpoint: util.OpId, buckets: string[]): Promise<util.BucketChecksum[]> {
-    if (buckets.length == 0) {
-      return [];
+  async getChecksums(checkpoint: util.OpId, buckets: string[]): Promise<util.ChecksumMap> {
+    return this.checksumCache.getChecksumMap(checkpoint, buckets);
+  }
+
+  private async getChecksumsInternal(batch: FetchPartialBucketChecksum[]): Promise<util.ChecksumMap> {
+    if (batch.length == 0) {
+      return new Map();
     }
+
     const filters: any[] = [];
-    for (let name of buckets) {
+    for (let request of batch) {
       filters.push({
         _id: {
           $gt: {
             g: this.group_id,
-            b: name,
-            o: new bson.MinKey()
+            b: request.bucket,
+            o: request.start ? BigInt(request.start) : new bson.MinKey()
           },
           $lte: {
             g: this.group_id,
-            b: name,
-            o: BigInt(checkpoint)
+            b: request.bucket,
+            o: BigInt(request.end)
           }
         }
       });
@@ -354,13 +365,18 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
       )
       .toArray();
 
-    return aggregate.map((doc) => {
-      return {
-        bucket: doc._id,
-        count: doc.count,
-        checksum: Number(BigInt(doc.checksum_total) & 0xffffffffn) & 4294967295
-      };
-    });
+    return new Map<string, util.BucketChecksum>(
+      aggregate.map((doc) => {
+        return [
+          doc._id,
+          {
+            bucket: doc._id,
+            count: doc.count,
+            checksum: Number(BigInt(doc.checksum_total) & 0xffffffffn) & 0xffffffff
+          } satisfies util.BucketChecksum
+        ];
+      })
+    );
   }
 
   async terminate() {

package/src/sync/sync.ts

@@ -78,8 +78,8 @@ async function* streamResponseInner(
   // This starts with the state from the client. May contain buckets that the user do not have access to (anymore).
   let dataBuckets = new Map<string, string>();
 
-  let last_checksums: util.BucketChecksum[] | null = null;
-  let last_write_checkpoint: bigint | null = null;
+  let lastChecksums: util.ChecksumMap | null = null;
+  let lastWriteCheckpoint: bigint | null = null;
 
   const { raw_data, binary_data } = params;
 
@@ -113,39 +113,42 @@ async function* streamResponseInner(
       throw new Error(`Too many buckets: ${allBuckets.length}`);
     }
 
-    let checksums: util.BucketChecksum[] | undefined = undefined;
-
     let dataBucketsNew = new Map<string, string>();
     for (let bucket of allBuckets) {
       dataBucketsNew.set(bucket, dataBuckets.get(bucket) ?? '0');
     }
     dataBuckets = dataBucketsNew;
 
-    checksums = await storage.getChecksums(checkpoint, [...dataBuckets.keys()]);
+    const bucketList = [...dataBuckets.keys()];
+    const checksumMap = await storage.getChecksums(checkpoint, bucketList);
+    // Subset of buckets for which there may be new data in this batch.
+    let bucketsToFetch: string[];
 
-    if (last_checksums) {
-      const diff = util.checksumsDiff(last_checksums, checksums);
+    if (lastChecksums) {
+      const diff = util.checksumsDiff(lastChecksums, checksumMap);
 
       if (
-        last_write_checkpoint == writeCheckpoint &&
-        diff.removed_buckets.length == 0 &&
-        diff.updated_buckets.length == 0
+        lastWriteCheckpoint == writeCheckpoint &&
+        diff.removedBuckets.length == 0 &&
+        diff.updatedBuckets.length == 0
       ) {
         // No changes - don't send anything to the client
         continue;
       }
+      bucketsToFetch = diff.updatedBuckets.map((c) => c.bucket);
 
       let message = `Updated checkpoint: ${checkpoint} | write: ${writeCheckpoint} | `;
       message += `buckets: ${allBuckets.length} | `;
-      message += `updated: ${limitedBuckets(diff.updated_buckets, 20)} | `;
-      message += `removed: ${limitedBuckets(diff.removed_buckets, 20)} | `;
+      message += `updated: ${limitedBuckets(diff.updatedBuckets, 20)} | `;
+      message += `removed: ${limitedBuckets(diff.removedBuckets, 20)} | `;
       micro.logger.info(message);
 
       const checksum_line: util.StreamingSyncCheckpointDiff = {
         checkpoint_diff: {
           last_op_id: checkpoint,
           write_checkpoint: writeCheckpoint ? String(writeCheckpoint) : undefined,
-          ...diff
+          removed_buckets: diff.removedBuckets,
+          updated_buckets: diff.updatedBuckets
         }
       };
 
@@ -154,35 +157,41 @@ async function* streamResponseInner(
       let message = `New checkpoint: ${checkpoint} | write: ${writeCheckpoint} | `;
       message += `buckets: ${allBuckets.length} ${limitedBuckets(allBuckets, 20)}`;
       micro.logger.info(message);
+      bucketsToFetch = allBuckets;
       const checksum_line: util.StreamingSyncCheckpoint = {
         checkpoint: {
           last_op_id: checkpoint,
           write_checkpoint: writeCheckpoint ? String(writeCheckpoint) : undefined,
-          buckets: checksums
+          buckets: [...checksumMap.values()]
         }
       };
       yield checksum_line;
     }
+    lastChecksums = checksumMap;
+    lastWriteCheckpoint = writeCheckpoint;
 
-    last_checksums = checksums;
-    last_write_checkpoint = writeCheckpoint;
-
-    yield* bucketDataInBatches(storage, checkpoint, dataBuckets, raw_data, binary_data, signal);
+    // This incrementally updates dataBuckets with each individual bucket position.
+    // At the end of this, we can be sure that all buckets have data up to the checkpoint.
+    yield* bucketDataInBatches({ storage, checkpoint, bucketsToFetch, dataBuckets, raw_data, binary_data, signal });
 
     await new Promise((resolve) => setTimeout(resolve, 10));
   }
 }
 
-async function* bucketDataInBatches(
-  storage: storage.SyncRulesBucketStorage,
-  checkpoint: string,
-  dataBuckets: Map<string, string>,
-  raw_data: boolean | undefined,
-  binary_data: boolean | undefined,
-  signal: AbortSignal
-) {
+interface BucketDataRequest {
+  storage: storage.SyncRulesBucketStorage;
+  checkpoint: string;
+  bucketsToFetch: string[];
+  /** Bucket data position, modified by the request.  */
+  dataBuckets: Map<string, string>;
+  raw_data: boolean | undefined;
+  binary_data: boolean | undefined;
+  signal: AbortSignal;
+}
+
+async function* bucketDataInBatches(request: BucketDataRequest) {
   let isDone = false;
-  while (!signal.aborted && !isDone) {
+  while (!request.signal.aborted && !isDone) {
     // The code below is functionally the same as this for-await loop below.
     // However, the for-await loop appears to have a memory leak, so we avoid it.
     // for await (const { done, data } of bucketDataBatch(storage, checkpoint, dataBuckets, raw_data, signal)) {
@@ -192,7 +201,7 @@ async function* bucketDataInBatches(
     //   }
     //   break;
     // }
-    const iter = bucketDataBatch(storage, checkpoint, dataBuckets, raw_data, binary_data, signal);
+    const iter = bucketDataBatch(request);
     try {
       while (true) {
         const { value, done: iterDone } = await iter.next();
@@ -215,17 +224,15 @@ async function* bucketDataInBatches(
 /**
  * Extracted as a separate internal function just to avoid memory leaks.
  */
-async function* bucketDataBatch(
-  storage: storage.SyncRulesBucketStorage,
-  checkpoint: string,
-  dataBuckets: Map<string, string>,
-  raw_data: boolean | undefined,
-  binary_data: boolean | undefined,
-  signal: AbortSignal
-) {
+async function* bucketDataBatch(request: BucketDataRequest) {
+  const { storage, checkpoint, bucketsToFetch, dataBuckets, raw_data, binary_data, signal } = request;
+
   const [_, release] = await syncSemaphore.acquire();
   try {
-    const data = storage.getBucketDataBatch(checkpoint, dataBuckets);
+    // Optimization: Only fetch buckets for which the checksums have changed since the last checkpoint
+    // For the first batch, this will be all buckets.
+    const filteredBuckets = new Map(bucketsToFetch.map((bucket) => [bucket, dataBuckets.get(bucket)!]));
+    const data = storage.getBucketDataBatch(checkpoint, filteredBuckets);
 
     let has_more = false;
 

package/src/util/protocol-types.ts

@@ -89,7 +89,12 @@ export const StreamingSyncRequest = t.object({
   /**
    * Data is received in a serialized BSON Buffer
    */
-  binary_data: t.boolean.optional()
+  binary_data: t.boolean.optional(),
+
+  /**
+   * Client parameters to be passed to the sync rules.
+   */
+  parameters: t.record(t.any).optional()
 });
 
 export type StreamingSyncRequest = t.Decoded<typeof StreamingSyncRequest>;

package/src/util/utils.ts

@@ -7,6 +7,8 @@ import * as storage from '@/storage/storage-index.js';
 import { BucketChecksum, OpId } from './protocol-types.js';
 import { retriedQuery } from './pgwire_utils.js';
 
+export type ChecksumMap = Map<string, BucketChecksum>;
+
 export function hashData(type: string, id: string, data: string): number {
   const hash = crypto.createHash('sha256');
   hash.update(`put.${type}.${id}.${data}`);
@@ -30,32 +32,50 @@ export function timestampToOpId(ts: bigint): OpId {
   return ts.toString(10);
 }
 
-export function checksumsDiff(previous: BucketChecksum[], current: BucketChecksum[]) {
-  const updated_buckets: BucketChecksum[] = [];
+export function checksumsDiff(previous: ChecksumMap, current: ChecksumMap) {
+  // All changed ones
+  const updatedBuckets = new Map<string, BucketChecksum>();
 
-  const previousBuckets = new Map<string, BucketChecksum>();
-  for (let checksum of previous) {
-    previousBuckets.set(checksum.bucket, checksum);
-  }
-  for (let checksum of current) {
-    if (!previousBuckets.has(checksum.bucket)) {
-      updated_buckets.push(checksum);
+  const toRemove = new Set<string>(previous.keys());
+
+  for (let checksum of current.values()) {
+    const p = previous.get(checksum.bucket);
+    if (p == null) {
+      // Added
+      updatedBuckets.set(checksum.bucket, checksum);
     } else {
-      const p = previousBuckets.get(checksum.bucket);
-      if (p?.checksum != checksum.checksum || p?.count != checksum.count) {
-        updated_buckets.push(checksum);
+      toRemove.delete(checksum.bucket);
+      if (checksum.checksum != p.checksum || checksum.count != p.count) {
+        // Updated
+        updatedBuckets.set(checksum.bucket, checksum);
+      } else {
+        // No change
       }
-      previousBuckets.delete(checksum.bucket);
     }
   }
 
-  const removed_buckets: string[] = [...previousBuckets.keys()];
   return {
-    updated_buckets,
-    removed_buckets
+    updatedBuckets: [...updatedBuckets.values()],
+    removedBuckets: [...toRemove]
   };
 }
 
+export function addChecksums(a: number, b: number) {
+  return (a + b) & 0xffffffff;
+}
+
+export function addBucketChecksums(a: BucketChecksum, b: BucketChecksum | null): BucketChecksum {
+  if (b == null) {
+    return a;
+  } else {
+    return {
+      bucket: a.bucket,
+      count: a.count + b.count,
+      checksum: addChecksums(a.checksum, b.checksum)
+    };
+  }
+}
+
 export async function getClientCheckpoint(
   db: pgwire.PgClient,
   bucketStorage: storage.BucketStorageFactory,

package/test/src/__snapshots__/sync.test.ts.snap

@@ -12,7 +12,13 @@ exports[`sync - mongodb > expiring token 1`] = `
 [
   {
     "checkpoint": {
-      "buckets": [],
+      "buckets": [
+        {
+          "bucket": "mybucket[]",
+          "checksum": 0,
+          "count": 0,
+        },
+      ],
       "last_op_id": "0",
       "write_checkpoint": undefined,
     },
@@ -135,7 +141,13 @@ exports[`sync - mongodb > sync updates to global data 1`] = `
 [
   {
     "checkpoint": {
-      "buckets": [],
+      "buckets": [
+        {
+          "bucket": "mybucket[]",
+          "checksum": 0,
+          "count": 0,
+        },
+      ],
       "last_op_id": "0",
       "write_checkpoint": undefined,
     },