@powersync/service-module-mongodb-storage 0.12.7 → 0.12.9

package/src/storage/MongoReportStorage.ts

@@ -0,0 +1,174 @@
+import { storage } from '@powersync/service-core';
+import { event_types } from '@powersync/service-types';
+import { PowerSyncMongo } from './implementation/db.js';
+import { logger } from '@powersync/lib-services-framework';
+
+export class MongoReportStorage implements storage.ReportStorage {
+  public readonly db: PowerSyncMongo;
+
+  constructor(db: PowerSyncMongo) {
+    this.db = db;
+  }
+  async deleteOldConnectionData(data: event_types.DeleteOldConnectionData): Promise<void> {
+    const { date } = data;
+    const result = await this.db.connection_report_events.deleteMany({
+      connected_at: { $lt: date },
+      $or: [
+        { disconnected_at: { $exists: true } },
+        { jwt_exp: { $lt: new Date() }, disconnected_at: { $exists: false } }
+      ]
+    });
+    if (result.deletedCount > 0) {
+      logger.info(
+        `TTL from ${date.toISOString()}: ${result.deletedCount} MongoDB documents have been removed from connection_report_events.`
+      );
+    }
+  }
+
+  async getClientConnectionReports(
+    data: event_types.ClientConnectionReportRequest
+  ): Promise<event_types.ClientConnectionReportResponse> {
+    const { start, end } = data;
+    const result = await this.db.connection_report_events
+      .aggregate<event_types.ClientConnectionReportResponse>([
+        {
+          $match: {
+            connected_at: { $lte: end, $gte: start }
+          }
+        },
+        this.connectionsFacetPipeline(),
+        this.connectionsProjectPipeline()
+      ])
+      .toArray();
+    return result[0];
+  }
+
+  async reportClientConnection(data: event_types.ClientConnectionBucketData): Promise<void> {
+    const updateFilter = this.updateDocFilter(data.user_id, data.client_id!);
+    await this.db.connection_report_events.findOneAndUpdate(
+      updateFilter,
+      {
+        $set: data,
+        $unset: {
+          disconnected_at: ''
+        }
+      },
+      {
+        upsert: true
+      }
+    );
+  }
+  async reportClientDisconnection(data: event_types.ClientDisconnectionEventData): Promise<void> {
+    const { connected_at, user_id, client_id } = data;
+    await this.db.connection_report_events.findOneAndUpdate(
+      {
+        client_id,
+        user_id,
+        connected_at
+      },
+      {
+        $set: {
+          disconnected_at: data.disconnected_at
+        },
+        $unset: {
+          jwt_exp: ''
+        }
+      }
+    );
+  }
+  async getConnectedClients(): Promise<event_types.ClientConnectionReportResponse> {
+    const result = await this.db.connection_report_events
+      .aggregate<event_types.ClientConnectionReportResponse>([
+        {
+          $match: {
+            disconnected_at: { $exists: false },
+            jwt_exp: { $gt: new Date() }
+          }
+        },
+        this.connectionsFacetPipeline(),
+        this.connectionsProjectPipeline()
+      ])
+      .toArray();
+    return result[0];
+  }
+
+  async [Symbol.asyncDispose]() {
+    // No-op
+  }
+
+  private parseJsDate(date: Date) {
+    const year = date.getUTCFullYear();
+    const month = date.getUTCMonth();
+    const today = date.getUTCDate();
+    const day = date.getUTCDay();
+    return {
+      year,
+      month,
+      today,
+      day,
+      parsedDate: date
+    };
+  }
+
+  private connectionsFacetPipeline() {
+    return {
+      $facet: {
+        unique_users: [
+          {
+            $group: {
+              _id: '$user_id'
+            }
+          },
+          {
+            $count: 'count'
+          }
+        ],
+        sdk_versions_array: [
+          {
+            $group: {
+              _id: '$sdk',
+              total: { $sum: 1 },
+              client_ids: { $addToSet: '$client_id' },
+              user_ids: { $addToSet: '$user_id' }
+            }
+          },
+          {
+            $project: {
+              _id: 0,
+              sdk: '$_id',
+              users: { $size: '$user_ids' },
+              clients: { $size: '$client_ids' }
+            }
+          },
+          {
+            $sort: {
+              sdk: 1
+            }
+          }
+        ]
+      }
+    };
+  }
+
+  private connectionsProjectPipeline() {
+    return {
+      $project: {
+        users: { $ifNull: [{ $arrayElemAt: ['$unique_users.count', 0] }, 0] },
+        sdks: '$sdk_versions_array'
+      }
+    };
+  }
+
+  private updateDocFilter(userId: string, clientId: string) {
+    const { year, month, today } = this.parseJsDate(new Date());
+    const nextDay = today + 1;
+    return {
+      user_id: userId,
+      client_id: clientId,
+      connected_at: {
+        $gte: new Date(Date.UTC(year, month, today)),
+        $lt: new Date(Date.UTC(year, month, nextDay))
+      }
+    };
+  }
+}

package/src/storage/implementation/MongoBucketBatch.ts

@@ -28,7 +28,7 @@ import { MongoIdSequence } from './MongoIdSequence.js';
 import { batchCreateCustomWriteCheckpoints } from './MongoWriteCheckpointAPI.js';
 import { cacheKey, OperationBatch, RecordOperation } from './OperationBatch.js';
 import { PersistedBatch } from './PersistedBatch.js';
-import { idPrefixFilter } from './util.js';
+import { idPrefixFilter } from '../../utils/util.js';
 
 /**
  * 15MB

package/src/storage/implementation/MongoCompactor.ts

@@ -1,6 +1,13 @@
 import { mongo, MONGO_OPERATION_TIMEOUT_MS } from '@powersync/lib-service-mongodb';
 import { logger, ReplicationAssertionError, ServiceAssertionError } from '@powersync/lib-services-framework';
-import { addChecksums, InternalOpId, isPartialChecksum, storage, utils } from '@powersync/service-core';
+import {
+  addChecksums,
+  InternalOpId,
+  isPartialChecksum,
+  PopulateChecksumCacheResults,
+  storage,
+  utils
+} from '@powersync/service-core';
 
 import { PowerSyncMongo } from './db.js';
 import { BucketDataDocument, BucketDataKey, BucketStateDocument } from './models.js';
@@ -10,6 +17,7 @@ import { cacheKey } from './OperationBatch.js';
 interface CurrentBucketState {
   /** Bucket name */
   bucket: string;
+
   /**
    * Rows seen in the bucket, with the last op_id of each.
    */
@@ -96,67 +104,56 @@ export class MongoCompactor {
         // We can make this more efficient later on by iterating
         // through the buckets in a single query.
         // That makes batching more tricky, so we leave for later.
-        await this.compactInternal(bucket);
+        await this.compactSingleBucket(bucket);
       }
     } else {
-      await this.compactInternal(undefined);
+      await this.compactDirtyBuckets();
     }
   }
 
-  async compactInternal(bucket: string | undefined) {
-    const idLimitBytes = this.idLimitBytes;
+  private async compactDirtyBuckets() {
+    while (!this.signal?.aborted) {
+      // Process all buckets with 1 or more changes since last time
+      const buckets = await this.dirtyBucketBatch({ minBucketChanges: 1 });
+      if (buckets.length == 0) {
+        // All done
+        break;
+      }
+      for (let bucket of buckets) {
+        await this.compactSingleBucket(bucket);
+      }
+    }
+  }
 
-    let currentState: CurrentBucketState | null = null;
+  private async compactSingleBucket(bucket: string) {
+    const idLimitBytes = this.idLimitBytes;
 
-    let bucketLower: string | mongo.MinKey;
-    let bucketUpper: string | mongo.MaxKey;
+    let currentState: CurrentBucketState = {
+      bucket,
+      seen: new Map(),
+      trackingSize: 0,
+      lastNotPut: null,
+      opsSincePut: 0,
 
-    if (bucket == null) {
-      bucketLower = new mongo.MinKey();
-      bucketUpper = new mongo.MaxKey();
-    } else if (bucket.includes('[')) {
-      // Exact bucket name
-      bucketLower = bucket;
-      bucketUpper = bucket;
-    } else {
-      // Bucket definition name
-      bucketLower = `${bucket}[`;
-      bucketUpper = `${bucket}[\uFFFF`;
-    }
+      checksum: 0,
+      opCount: 0,
+      opBytes: 0
+    };
 
     // Constant lower bound
     const lowerBound: BucketDataKey = {
       g: this.group_id,
-      b: bucketLower as string,
+      b: bucket,
       o: new mongo.MinKey() as any
     };
 
     // Upper bound is adjusted for each batch
     let upperBound: BucketDataKey = {
       g: this.group_id,
-      b: bucketUpper as string,
+      b: bucket,
       o: new mongo.MaxKey() as any
     };
 
-    const doneWithBucket = async () => {
-      if (currentState == null) {
-        return;
-      }
-      // Free memory before clearing bucket
-      currentState.seen.clear();
-      if (currentState.lastNotPut != null && currentState.opsSincePut >= 1) {
-        logger.info(
-          `Inserting CLEAR at ${this.group_id}:${currentState.bucket}:${currentState.lastNotPut} to remove ${currentState.opsSincePut} operations`
-        );
-        // Need flush() before clear()
-        await this.flush();
-        await this.clearBucket(currentState);
-      }
-
-      // Do this _after_ clearBucket so that we have accurate counts.
-      this.updateBucketChecksums(currentState);
-    };
-
     while (!this.signal?.aborted) {
       // Query one batch at a time, to avoid cursor timeouts
       const cursor = this.db.bucket_data.aggregate<BucketDataDocument & { size: number | bigint }>(
@@ -184,7 +181,11 @@ export class MongoCompactor {
             }
           }
         ],
-        { batchSize: this.moveBatchQueryLimit }
+        {
+          // batchSize is 1 more than limit to auto-close the cursor.
+          // See https://github.com/mongodb/node-mongodb-native/pull/4580
+          batchSize: this.moveBatchQueryLimit + 1
+        }
       );
       // We don't limit to a single batch here, since that often causes MongoDB to scan through more than it returns.
       // Instead, we load up to the limit.
@@ -199,22 +200,6 @@ export class MongoCompactor {
       upperBound = batch[batch.length - 1]._id;
 
       for (let doc of batch) {
-        if (currentState == null || doc._id.b != currentState.bucket) {
-          await doneWithBucket();
-
-          currentState = {
-            bucket: doc._id.b,
-            seen: new Map(),
-            trackingSize: 0,
-            lastNotPut: null,
-            opsSincePut: 0,
-
-            checksum: 0,
-            opCount: 0,
-            opBytes: 0
-          };
-        }
-
         if (doc._id.o > this.maxOpId) {
           continue;
         }
@@ -285,12 +270,22 @@ export class MongoCompactor {
         }
       }
 
-      if (currentState != null) {
-        logger.info(`Processed batch of length ${batch.length} current bucket: ${currentState.bucket}`);
-      }
+      logger.info(`Processed batch of length ${batch.length} current bucket: ${bucket}`);
+    }
+
+    // Free memory before clearing bucket
+    currentState.seen.clear();
+    if (currentState.lastNotPut != null && currentState.opsSincePut >= 1) {
+      logger.info(
+        `Inserting CLEAR at ${this.group_id}:${bucket}:${currentState.lastNotPut} to remove ${currentState.opsSincePut} operations`
+      );
+      // Need flush() before clear()
+      await this.flush();
+      await this.clearBucket(currentState);
     }
 
-    await doneWithBucket();
+    // Do this _after_ clearBucket so that we have accurate counts.
+    this.updateBucketChecksums(currentState);
 
     // Need another flush after updateBucketChecksums()
     await this.flush();
@@ -478,50 +473,55 @@ export class MongoCompactor {
   /**
    * Subset of compact, only populating checksums where relevant.
    */
-  async populateChecksums() {
-    // This is updated after each batch
-    let lowerBound: BucketStateDocument['_id'] = {
-      g: this.group_id,
-      b: new mongo.MinKey() as any
-    };
-    // This is static
-    const upperBound: BucketStateDocument['_id'] = {
-      g: this.group_id,
-      b: new mongo.MaxKey() as any
-    };
+  async populateChecksums(options: { minBucketChanges: number }): Promise<PopulateChecksumCacheResults> {
+    let count = 0;
     while (!this.signal?.aborted) {
-      // By filtering buckets, we effectively make this "resumeable".
-      const filter: mongo.Filter<BucketStateDocument> = {
-        _id: {
-          $gt: lowerBound,
-          $lt: upperBound
-        },
-        compacted_state: { $exists: false }
-      };
+      const buckets = await this.dirtyBucketBatch(options);
+      if (buckets.length == 0) {
+        // All done
+        break;
+      }
+      const start = Date.now();
+      logger.info(`Calculating checksums for batch of ${buckets.length} buckets, starting at ${buckets[0]}`);
+
+      await this.updateChecksumsBatch(buckets);
+      logger.info(`Updated checksums for batch of ${buckets.length} buckets in ${Date.now() - start}ms`);
+      count += buckets.length;
+    }
+    return { buckets: count };
+  }
 
-      const bucketsWithoutChecksums = await this.db.bucket_state
-        .find(filter, {
+  /**
+   * Returns a batch of dirty buckets - buckets with most changes first.
+   *
+   * This cannot be used to iterate on its own - the client is expected to process these buckets and
+   * set estimate_since_compact.count: 0 when done, before fetching the next batch.
+   */
+  private async dirtyBucketBatch(options: { minBucketChanges: number }): Promise<string[]> {
+    if (options.minBucketChanges <= 0) {
+      throw new ReplicationAssertionError('minBucketChanges must be >= 1');
+    }
+    // We make use of an index on {_id.g: 1, 'estimate_since_compact.count': -1}
+    const dirtyBuckets = await this.db.bucket_state
+      .find(
+        {
+          '_id.g': this.group_id,
+          'estimate_since_compact.count': { $gte: options.minBucketChanges }
+        },
+        {
           projection: {
             _id: 1
           },
           sort: {
-            _id: 1
+            'estimate_since_compact.count': -1
           },
           limit: 5_000,
           maxTimeMS: MONGO_OPERATION_TIMEOUT_MS
-        })
-        .toArray();
-      if (bucketsWithoutChecksums.length == 0) {
-        // All done
-        break;
-      }
-
-      logger.info(`Calculating checksums for batch of ${bucketsWithoutChecksums.length} buckets`);
-
-      await this.updateChecksumsBatch(bucketsWithoutChecksums.map((b) => b._id.b));
+        }
+      )
+      .toArray();
 
-      lowerBound = bucketsWithoutChecksums[bucketsWithoutChecksums.length - 1]._id;
-    }
+    return dirtyBuckets.map((bucket) => bucket._id.b);
   }
 
   private async updateChecksumsBatch(buckets: string[]) {
@@ -555,6 +555,10 @@ export class MongoCompactor {
                 count: bucketChecksum.count,
                 checksum: BigInt(bucketChecksum.checksum),
                 bytes: null
+              },
+              estimate_since_compact: {
+                count: 0,
+                bytes: 0
               }
             }
           },

package/src/storage/implementation/MongoStorageProvider.ts

@@ -4,8 +4,9 @@ import { POWERSYNC_VERSION, storage } from '@powersync/service-core';
 import { MongoStorageConfig } from '../../types/types.js';
 import { MongoBucketStorage } from '../MongoBucketStorage.js';
 import { PowerSyncMongo } from './db.js';
+import { MongoReportStorage } from '../MongoReportStorage.js';
 
-export class MongoStorageProvider implements storage.BucketStorageProvider {
+export class MongoStorageProvider implements storage.StorageProvider {
   get type() {
     return lib_mongo.MONGO_CONNECTION_TYPE;
   }
@@ -37,15 +38,19 @@ export class MongoStorageProvider implements storage.BucketStorageProvider {
     await client.connect();
 
     const database = new PowerSyncMongo(client, { database: resolvedConfig.storage.database });
-    const factory = new MongoBucketStorage(database, {
+    const syncStorageFactory = new MongoBucketStorage(database, {
       // TODO currently need the entire resolved config due to this
       slot_name_prefix: resolvedConfig.slot_name_prefix
     });
+
+    // Storage factory for reports
+    const reportStorageFactory = new MongoReportStorage(database);
     return {
-      storage: factory,
+      storage: syncStorageFactory,
+      reportStorage: reportStorageFactory,
       shutDown: async () => {
         shuttingDown = true;
-        await factory[Symbol.asyncDispose]();
+        await syncStorageFactory[Symbol.asyncDispose]();
         await client.close();
       },
       tearDown: () => {

package/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -16,6 +16,8 @@ import {
   InternalOpId,
   internalToExternalOpId,
   maxLsn,
+  PopulateChecksumCacheOptions,
+  PopulateChecksumCacheResults,
   ProtocolOpId,
   ReplicationCheckpoint,
   storage,
@@ -35,7 +37,8 @@ import { MongoChecksumOptions, MongoChecksums } from './MongoChecksums.js';
 import { MongoCompactor } from './MongoCompactor.js';
 import { MongoParameterCompactor } from './MongoParameterCompactor.js';
 import { MongoWriteCheckpointAPI } from './MongoWriteCheckpointAPI.js';
-import { idPrefixFilter, mapOpEntry, readSingleBatch, setSessionSnapshotTime } from './util.js';
+import { idPrefixFilter, mapOpEntry, readSingleBatch, setSessionSnapshotTime } from '../../utils/util.js';
+
 
 export interface MongoSyncBucketStorageOptions {
   checksumOptions?: MongoChecksumOptions;
@@ -403,7 +406,9 @@ export class MongoSyncBucketStorage
         limit: batchLimit,
         // Increase batch size above the default 101, so that we can fill an entire batch in
         // one go.
-        batchSize: batchLimit,
+        // batchSize is 1 more than limit to auto-close the cursor.
+        // See https://github.com/mongodb/node-mongodb-native/pull/4580
+        batchSize: batchLimit + 1,
         // Raw mode is returns an array of Buffer instead of parsed documents.
         // We use it so that:
         // 1. We can calculate the document size accurately without serializing again.
@@ -663,7 +668,7 @@ export class MongoSyncBucketStorage
     }
   }
 
-  async populatePersistentChecksumCache(options: Required<Pick<CompactOptions, 'signal' | 'maxOpId'>>): Promise<void> {
+  async populatePersistentChecksumCache(options: PopulateChecksumCacheOptions): Promise<PopulateChecksumCacheResults> {
     logger.info(`Populating persistent checksum cache...`);
     const start = Date.now();
     // We do a minimal compact here.
@@ -674,9 +679,14 @@ export class MongoSyncBucketStorage
       memoryLimitMB: 0
     });
 
-    await compactor.populateChecksums();
+    const result = await compactor.populateChecksums({
+      // There are cases with millions of small buckets, in which case it can take very long to
+      // populate the checksums, with minimal benefit. We skip the small buckets here.
+      minBucketChanges: options.minBucketChanges ?? 10
+    });
     const duration = Date.now() - start;
     logger.info(`Populated persistent checksum cache in ${(duration / 1000).toFixed(1)}s`);
+    return result;
   }
 
   /**
@@ -905,7 +915,9 @@ export class MongoSyncBucketStorage
             '_id.b': 1
           },
           limit: limit + 1,
-          batchSize: limit + 1,
+          // batchSize is 1 more than limit to auto-close the cursor.
+          // See https://github.com/mongodb/node-mongodb-native/pull/4580
+          batchSize: limit + 2,
           singleBatch: true
         }
       )
@@ -935,7 +947,9 @@ export class MongoSyncBucketStorage
             lookup: 1
           },
           limit: limit + 1,
-          batchSize: limit + 1,
+          // batchSize is 1 more than limit to auto-close the cursor.
+          // See https://github.com/mongodb/node-mongodb-native/pull/4580
+          batchSize: limit + 2,
           singleBatch: true
         }
       )

package/src/storage/implementation/MongoWriteCheckpointAPI.ts

@@ -111,7 +111,9 @@ export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
         },
         {
           limit: limit + 1,
-          batchSize: limit + 1,
+          // batchSize is 1 more than limit to auto-close the cursor.
+          // See https://github.com/mongodb/node-mongodb-native/pull/4580
+          batchSize: limit + 2,
           singleBatch: true
         }
       )
@@ -140,7 +142,9 @@ export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
         },
         {
           limit: limit + 1,
-          batchSize: limit + 1,
+          // batchSize is 1 more than limit to auto-close the cursor.
+          // See https://github.com/mongodb/node-mongodb-native/pull/4580
+          batchSize: limit + 2,
           singleBatch: true
         }
       )

package/src/storage/implementation/PersistedBatch.ts

@@ -16,7 +16,7 @@ import {
   CurrentDataDocument,
   SourceKey
 } from './models.js';
-import { replicaIdToSubkey } from './util.js';
+import { replicaIdToSubkey } from '../../utils/util.js';
 
 /**
  * Maximum size of operations we write in a single transaction.

package/src/storage/implementation/db.ts

@@ -8,6 +8,7 @@ import {
   BucketParameterDocument,
   BucketStateDocument,
   CheckpointEventDocument,
+  ClientConnectionDocument,
   CurrentDataDocument,
   CustomWriteCheckpointDocument,
   IdSequenceDocument,
@@ -37,6 +38,7 @@ export class PowerSyncMongo {
   readonly locks: mongo.Collection<lib_mongo.locks.Lock>;
   readonly bucket_state: mongo.Collection<BucketStateDocument>;
   readonly checkpoint_events: mongo.Collection<CheckpointEventDocument>;
+  readonly connection_report_events: mongo.Collection<ClientConnectionDocument>;
 
   readonly client: mongo.MongoClient;
   readonly db: mongo.Db;
@@ -61,6 +63,7 @@ export class PowerSyncMongo {
     this.locks = this.db.collection('locks');
     this.bucket_state = this.db.collection('bucket_state');
     this.checkpoint_events = this.db.collection('checkpoint_events');
+    this.connection_report_events = this.db.collection('connection_report_events');
   }
 
   /**
@@ -128,6 +131,20 @@ export class PowerSyncMongo {
     });
   }
 
+  /**
+   * Only use in migrations and tests.
+   */
+  async createConnectionReportingCollection() {
+    const existingCollections = await this.db
+      .listCollections({ name: 'connection_report_events' }, { nameOnly: false })
+      .toArray();
+    const collection = existingCollections[0];
+    if (collection != null) {
+      return;
+    }
+    await this.db.createCollection('connection_report_events');
+  }
+
   /**
    * Only use in migrations and tests.
    */
@@ -141,6 +158,19 @@ export class PowerSyncMongo {
       { name: 'bucket_updates', unique: true }
     );
   }
+  /**
+   * Only use in migrations and tests.
+   */
+  async createBucketStateIndex2() {
+    // TODO: Implement a better mechanism to use migrations in tests
+    await this.bucket_state.createIndex(
+      {
+        '_id.g': 1,
+        'estimate_since_compact.count': -1
+      },
+      { name: 'dirty_count' }
+    );
+  }
 }
 
 export function createPowerSyncMongo(config: MongoStorageConfig, options?: lib_mongo.MongoConnectionOptions) {

package/src/storage/implementation/models.ts

@@ -1,6 +1,7 @@
 import { InternalOpId, storage } from '@powersync/service-core';
 import { SqliteJsonValue } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
+import { event_types } from '@powersync/service-types';
 
 /**
  * Replica id uniquely identifying a row on the source database.
@@ -238,3 +239,5 @@ export interface InstanceDocument {
   // The instance UUID
   _id: string;
 }
+
+export interface ClientConnectionDocument extends event_types.ClientConnection {}