Improve parameter query filtering.

Author: rkistner

modules/module-mongodb-storage/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -18,7 +18,9 @@ import {
   ReplicationCheckpoint,
   storage,
   utils,
-  WatchWriteCheckpointOptions
+  WatchWriteCheckpointOptions,
+  CHECKPOINT_INVALIDATE_ALL,
+  deserializeParameterLookup
 } from '@powersync/service-core';
 import { SqliteJsonRow, SqliteJsonValue, SqlSyncRules } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
@@ -39,6 +41,7 @@ import { MongoBucketBatch } from './MongoBucketBatch.js';
 import { MongoCompactor } from './MongoCompactor.js';
 import { MongoWriteCheckpointAPI } from './MongoWriteCheckpointAPI.js';
 import { idPrefixFilter, mapOpEntry, readSingleBatch } from './util.js';
+import { JSONBig } from '@powersync/service-jsonbig';
 
 export class MongoSyncBucketStorage
   extends BaseObserver<storage.SyncRulesBucketStorageListener>
@@ -795,12 +798,7 @@ export class MongoSyncBucketStorage
 
       const updates: CheckpointChanges =
         lastCheckpoint == null
-          ? {
-              invalidateDataBuckets: true,
-              invalidateParameterBuckets: true,
-              updatedDataBuckets: [],
-              updatedParameterBucketDefinitions: []
-            }
+          ? CHECKPOINT_INVALIDATE_ALL
           : await this.getCheckpointChanges({
               lastCheckpoint: lastCheckpoint,
               nextCheckpoint: checkpoint
@@ -877,7 +875,7 @@ export class MongoSyncBucketStorage
     // Or we could store updated buckets in a separate collection, and query those.
     // For now, we ignore this optimization
 
-    // const dataBuckets = await this.db.bucket_data
+    // const dataBucketDocuments = await this.db.bucket_data
     //   .find(
     //     {
     //       '_id.g': this.group_id,
@@ -894,6 +892,14 @@ export class MongoSyncBucketStorage
     //   )
     //   .toArray();
 
+    // const buckets = dataBucketDocuments.map((doc) => doc._id.b);
+    // const invalidateDataBuckets = buckets.length > 1000;
+
+    // return {
+    //   invalidateDataBuckets: invalidateDataBuckets,
+    //   updatedDataBuckets: invalidateDataBuckets ? [] : buckets
+    // };
+
     return {
       invalidateDataBuckets: true,
       updatedDataBuckets: []
@@ -902,7 +908,7 @@ export class MongoSyncBucketStorage
 
   private async getParameterBucketChanges(
     options: GetCheckpointChangesOptions
-  ): Promise<Pick<CheckpointChanges, 'updatedParameterBucketDefinitions' | 'invalidateParameterBuckets'>> {
+  ): Promise<Pick<CheckpointChanges, 'updatedParameterLookups' | 'invalidateParameterBuckets'>> {
     // TODO: limit max query running time
     const parameterUpdates = await this.db.bucket_parameters
       .find(
@@ -924,9 +930,9 @@ export class MongoSyncBucketStorage
 
     return {
       invalidateParameterBuckets: invalidateParameterUpdates,
-      updatedParameterBucketDefinitions: invalidateParameterUpdates
-        ? []
-        : [...new Set<string>(parameterUpdates.map((p) => getLookupBucketDefinitionName(p.lookup)))]
+      updatedParameterLookups: invalidateParameterUpdates
+        ? new Set<string>()
+        : new Set<string>(parameterUpdates.map((p) => JSONBig.stringify(deserializeParameterLookup(p.lookup))))
     };
   }
 
@@ -937,7 +943,7 @@ export class MongoSyncBucketStorage
     max: 50,
     maxSize: 10 * 1024 * 1024,
     sizeCalculation: (value: CheckpointChanges) => {
-      return 100 + value.updatedParameterBucketDefinitions.reduce<number>((a, b) => a + b.length, 0);
+      return 100 + [...value.updatedParameterLookups].reduce<number>((a, b) => a + b.length, 0);
     },
     fetchMethod: async (_key, _staleValue, options) => {
       return this.getCheckpointChangesInternal(options.context.options);
@@ -966,7 +972,7 @@ export class MongoSyncBucketStorage
         invalidateDataBuckets: true,
         updatedDataBuckets: [],
         invalidateParameterBuckets: false,
-        updatedParameterBucketDefinitions: []
+        updatedParameterLookups: new Set<string>()
       };
     }
     const key = `${options.lastCheckpoint}_${options.nextCheckpoint}`;

packages/service-core/src/storage/SyncRulesBucketStorage.ts

@@ -253,13 +253,14 @@ export interface GetCheckpointChangesOptions {
 export interface CheckpointChanges {
   updatedDataBuckets: string[];
   invalidateDataBuckets: boolean;
-  updatedParameterBucketDefinitions: string[];
+  /** Serialized using JSONBig */
+  updatedParameterLookups: Set<string>;
   invalidateParameterBuckets: boolean;
 }
 
 export const CHECKPOINT_INVALIDATE_ALL: CheckpointChanges = {
   updatedDataBuckets: [],
   invalidateDataBuckets: true,
-  updatedParameterBucketDefinitions: [],
+  updatedParameterLookups: new Set<string>(),
   invalidateParameterBuckets: true
 };

packages/service-core/src/storage/bson.ts

@@ -25,22 +25,20 @@ export const BSON_DESERIALIZE_DATA_OPTIONS: bson.DeserializeOptions = {
  * @param lookup
  */
 export const serializeLookupBuffer = (lookup: SqliteJsonValue[]): NodeBuffer => {
-  const normalized = lookup.map((value) => {
-    if (typeof value == 'number' && Number.isInteger(value)) {
-      return BigInt(value);
-    } else {
-      return value;
-    }
-  });
-  return bson.serialize({ l: normalized }) as NodeBuffer;
+  return bson.serialize({ l: lookup }) as NodeBuffer;
 };
 
 export const serializeLookup = (lookup: SqliteJsonValue[]) => {
   return new bson.Binary(serializeLookupBuffer(lookup));
 };
 
-export const getLookupBucketDefinitionName = (lookup: bson.Binary) => {
+export const deserializeParameterLookup = (lookup: bson.Binary) => {
   const parsed = bson.deserialize(lookup.buffer, BSON_DESERIALIZE_INTERNAL_OPTIONS).l as SqliteJsonValue[];
+  return parsed;
+};
+
+export const getLookupBucketDefinitionName = (lookup: bson.Binary) => {
+  const parsed = deserializeParameterLookup(lookup);
   return parsed[0] as string;
 };
 

packages/service-core/src/sync/BucketChecksumState.ts

@@ -7,6 +7,8 @@ import { ErrorCode, logger, ServiceAssertionError, ServiceError } from '@powersy
 import { BucketParameterQuerier } from '@powersync/service-sync-rules/src/BucketParameterQuerier.js';
 import { BucketSyncState } from './sync.js';
 import { SyncContext } from './SyncContext.js';
+import { JSONBig } from '@powersync/service-jsonbig';
+import { hasIntersection } from './util.js';
 
 export interface BucketChecksumStateOptions {
   syncContext: SyncContext;
@@ -269,6 +271,9 @@ export class BucketParameterState {
   private readonly querier: BucketParameterQuerier;
   private readonly staticBuckets: Map<string, BucketDescription>;
   private cachedDynamicBuckets: BucketDescription[] | null = null;
+  private cachedDynamicBucketSet: Set<string> | null = null;
+
+  private readonly lookups: Set<string>;
 
   constructor(
     context: SyncContext,
@@ -283,6 +288,7 @@ export class BucketParameterState {
 
     this.querier = syncRules.getBucketParameterQuerier(this.syncParams);
     this.staticBuckets = new Map<string, BucketDescription>(this.querier.staticBuckets.map((b) => [b.bucket, b]));
+    this.lookups = new Set<string>(this.querier.parameterQueryLookups.map((l) => JSONBig.stringify(l)));
   }
 
   async getCheckpointUpdate(checkpoint: storage.StorageCheckpointUpdate): Promise<CheckpointUpdate | null> {
@@ -362,46 +368,61 @@ export class BucketParameterState {
     const staticBuckets = querier.staticBuckets;
     const update = checkpoint.update;
 
-    let hasDataChange = false;
     let hasParameterChange = false;
-    if (update.invalidateDataBuckets || update.updatedDataBuckets?.length > 0) {
-      hasDataChange = true;
+    let invalidateDataBuckets = false;
+    // If hasParameterChange == true, then invalidateDataBuckets = true
+    // If invalidateDataBuckets == true, we ignore updatedBuckets
+    let updatedBuckets = new Set<string>();
+
+    if (update.invalidateDataBuckets) {
+      invalidateDataBuckets = true;
     }
 
     if (update.invalidateParameterBuckets) {
       hasParameterChange = true;
     } else {
-      for (let bucket of update.updatedParameterBucketDefinitions) {
-        // This is a very coarse check, but helps
-        if (querier.dynamicBucketDefinitions.has(bucket)) {
-          hasParameterChange = true;
-          break;
-        }
+      if (hasIntersection(this.lookups, update.updatedParameterLookups)) {
+        // This is a very coarse re-check of all queries
+        hasParameterChange = true;
       }
     }
 
-    if (!hasDataChange && !hasParameterChange) {
-      return null;
-    }
-
     let dynamicBuckets: BucketDescription[];
-    if (hasParameterChange || this.cachedDynamicBuckets == null) {
+    if (hasParameterChange || this.cachedDynamicBuckets == null || this.cachedDynamicBucketSet == null) {
       dynamicBuckets = await querier.queryDynamicBucketDescriptions({
         getParameterSets(lookups) {
           return storage.getParameterSets(checkpoint.base.checkpoint, lookups);
         }
       });
       this.cachedDynamicBuckets = dynamicBuckets;
+      this.cachedDynamicBucketSet = new Set<string>(dynamicBuckets.map((b) => b.bucket));
+      invalidateDataBuckets = true;
     } else {
       dynamicBuckets = this.cachedDynamicBuckets;
+
+      if (!invalidateDataBuckets) {
+        // TODO: Do set intersection instead
+        for (let bucket of update.updatedDataBuckets ?? []) {
+          if (this.staticBuckets.has(bucket) || this.cachedDynamicBucketSet.has(bucket)) {
+            updatedBuckets.add(bucket);
+          }
+        }
+      }
     }
     const allBuckets = [...staticBuckets, ...dynamicBuckets];
 
-    return {
-      buckets: allBuckets,
-      // We cannot track individual bucket updates for dynamic lookups yet
-      updatedBuckets: null
-    };
+    if (invalidateDataBuckets) {
+      return {
+        buckets: allBuckets,
+        // We cannot track individual bucket updates for dynamic lookups yet
+        updatedBuckets: null
+      };
+    } else {
+      return {
+        buckets: allBuckets,
+        updatedBuckets: updatedBuckets
+      };
+    }
   }
 }
 

packages/service-core/src/sync/util.ts

@@ -153,3 +153,15 @@ export function settledPromise<T>(promise: Promise<T>): Promise<PromiseSettledRe
     }
   );
 }
+
+export function hasIntersection<T>(a: Set<T>, b: Set<T>) {
+  if (a.size > b.size) {
+    [a, b] = [b, a];
+  }
+  // Now, a is always smaller than b, so iterate over a
+  for (let value of a) {
+    if (b.has(value)) {
+      return true;
+    }
+  }
+}

packages/sync-rules/src/BucketParameterQuerier.ts

@@ -19,11 +19,11 @@ export interface BucketParameterQuerier {
    * True if there are dynamic buckets, meaning queryDynamicBucketDescriptions() should be used.
    *
    * If this is false, queryDynamicBucketDescriptions() will always return an empty array,
-   * and dynamicBucketDefinitions.size == 0.
+   * and parameterQueryLookups.length == 0.
    */
   readonly hasDynamicBuckets: boolean;
 
-  readonly dynamicBucketDefinitions: Set<string>;
+  readonly parameterQueryLookups: SqliteJsonValue[][];
 
   /**
    * These buckets depend on parameter storage, and needs to be retrieved dynamically for each checkpoint.
@@ -47,11 +47,11 @@ export interface QueryBucketDescriptorOptions extends ParameterLookupSource {
 }
 
 export function mergeBucketParameterQueriers(queriers: BucketParameterQuerier[]): BucketParameterQuerier {
-  const dynamicBucketDefinitions = new Set<string>(queriers.flatMap((q) => [...q.dynamicBucketDefinitions]));
+  const parameterQueryLookups = queriers.flatMap((q) => q.parameterQueryLookups);
   return {
     staticBuckets: queriers.flatMap((q) => q.staticBuckets),
-    hasDynamicBuckets: dynamicBucketDefinitions.size > 0,
-    dynamicBucketDefinitions,
+    hasDynamicBuckets: parameterQueryLookups.length > 0,
+    parameterQueryLookups: parameterQueryLookups,
     async queryDynamicBucketDescriptions(source: ParameterLookupSource) {
       let results: BucketDescription[] = [];
       for (let q of queriers) {

packages/sync-rules/src/SqlBucketDescriptor.ts

@@ -113,7 +113,7 @@ export class SqlBucketDescriptor {
     const staticQuerier = {
       staticBuckets,
       hasDynamicBuckets: false,
-      dynamicBucketDefinitions: new Set<string>(),
+      parameterQueryLookups: [],
       queryDynamicBucketDescriptions: async () => []
     } satisfies BucketParameterQuerier;
 

packages/sync-rules/src/SqlParameterQuery.ts

@@ -23,7 +23,7 @@ import {
   SqliteJsonValue,
   SqliteRow
 } from './types.js';
-import { filterJsonRow, getBucketId, isJsonValue, isSelectStatement } from './utils.js';
+import { filterJsonRow, getBucketId, isJsonValue, isSelectStatement, normalizeParameterValue } from './utils.js';
 
 /**
  * Represents a parameter query, such as:
@@ -308,7 +308,7 @@ export class SqlParameterQuery {
           const value = param.parametersToLookupValue(parameters);
 
           if (isJsonValue(value)) {
-            return value;
+            return normalizeParameterValue(value);
           } else {
             valid = false;
             return null;
@@ -339,17 +339,18 @@ export class SqlParameterQuery {
         .map((expandedValue) => {
           let lookup: SqliteJsonValue[] = [this.descriptor_name!, this.id!];
           let valid = true;
+          const normalizedExpandedValue = normalizeParameterValue(expandedValue);
           lookup.push(
             ...this.input_parameters!.map((param): SqliteJsonValue => {
               if (param == this.expanded_input_parameter) {
                 // Expand array value
-                return expandedValue;
+                return normalizedExpandedValue;
               } else {
                 // Scalar value
                 const value = param.parametersToLookupValue(parameters);
 
                 if (isJsonValue(value)) {
-                  return value;
+                  return normalizeParameterValue(value);
                 } else {
                   valid = false;
                   return null;
@@ -375,15 +376,15 @@ export class SqlParameterQuery {
       return {
         staticBuckets: [],
         hasDynamicBuckets: false,
-        dynamicBucketDefinitions: new Set<string>(),
+        parameterQueryLookups: [],
         queryDynamicBucketDescriptions: async () => []
       };
     }
 
     return {
       staticBuckets: [],
       hasDynamicBuckets: true,
-      dynamicBucketDefinitions: new Set<string>([this.descriptor_name!]),
+      parameterQueryLookups: lookups,
       queryDynamicBucketDescriptions: async (source: ParameterLookupSource) => {
         const bucketParameters = await source.getParameterSets(lookups);
         return this.resolveBucketDescriptions(bucketParameters, requestParameters);

packages/sync-rules/src/utils.ts

@@ -170,3 +170,13 @@ export const JSONBucketNameSerialize = {
     return stringifyRaw(value, replacer, space)!;
   }
 };
+
+/**
+ * Lookup serialization must be number-agnostic. I.e. normalize numbers, instead of preserving numbers.
+ */
+export function normalizeParameterValue(value: SqliteJsonValue): SqliteJsonValue {
+  if (typeof value == 'number' && Number.isInteger(value)) {
+    return BigInt(value);
+  }
+  return value;
+}