@powersync/service-module-mongodb-storage 0.24.0 → 0.24.2

package/src/storage/MongoBucketStorage.ts

@@ -52,6 +52,8 @@ export interface MongoBucketStorageOptions {
    */
   replicationStreamNamePrefix: string;
   readPreference?: mongo.ReadPreference;
+  /** Read preference for collection discovery and collection stats used by storage metrics. */
+  statsReadPreference?: mongo.ReadPreference;
   checksumCacheTtlMs?: number;
   clearBatchThrottleRate?: number;
   defaultStorageVersion?: number;
@@ -714,7 +716,12 @@ export class MongoBucketStorage extends storage.BucketStorageFactory {
 
     if (storageConfig.incrementalReprocessing) {
       const v3 = doc as ReplicationStreamDocumentV3;
-      const matching = v3.sync_configs.filter((c) => stateFilter.includes(c.state));
+      // Cleanup of a stopped stream requires all embedded configs, regardless of
+      // their individual states, so every config's storage can be terminated.
+      const matching =
+        v3.state == storage.SyncRuleState.STOP && stateFilter.includes(storage.SyncRuleState.STOP)
+          ? v3.sync_configs
+          : v3.sync_configs.filter((c) => stateFilter.includes(c.state));
       if (matching.length == 0) {
         return null;
       }
@@ -828,89 +835,47 @@ export class MongoBucketStorage extends storage.BucketStorageFactory {
   }
 
   async getStorageMetrics(): Promise<storage.StorageMetrics> {
-    const ignoreNotExisting = (e: unknown) => {
-      if (lib_mongo.isMongoNamespaceNotFoundError(e)) {
-        // Collection doesn't exist - return 0
-        return [{ storageStats: { size: 0 } }];
-      } else {
-        return Promise.reject(e);
-      }
-    };
-
-    const aggregateStaticCollection = async <T extends mongo.Document>(collection: mongo.Collection<T>) => {
-      // We check whether the collection exists before getting the statistics. This avoids repeated
-      // errors in the MongoDB logs if the collection hasn't been created yet.
-      const exists =
-        (await this.db.db.listCollections({ name: collection.collectionName }, { nameOnly: true }).toArray()).length >
-        0;
-      if (!exists) {
-        return [{ storageStats: { size: 0 } }];
-      }
-
-      return collection
-        .aggregate([
-          {
-            $collStats: {
-              storageStats: {}
-            }
+    // We check which collections exist before getting statistics to avoid repeated errors in
+    // MongoDB logs for collections that have not been created yet.
+    const collections = await this.db.db
+      .listCollections(
+        {
+          name: {
+            $regex: '^(bucket_data|bucket_parameters|current_data)$|^(bucket_data_|parameter_index_|source_records_)'
           }
-        ])
-        .toArray()
-        .catch(ignoreNotExisting);
-    };
-
-    const operations_aggregate = await aggregateStaticCollection(this.db.bucket_data);
-    const v3OperationCollections = await this.db.listBucketDataCollectionsV3();
-    const v3_operation_aggregates = await Promise.all(
-      v3OperationCollections.map((collection) =>
-        collection
-          .aggregate([
-            {
-              $collStats: {
-                storageStats: {}
-              }
-            }
-          ])
-          .toArray()
-          .catch(ignoreNotExisting)
-      )
-    );
-
-    const parameters_aggregate = await aggregateStaticCollection(this.db.bucket_parameters);
-
-    const v3ParameterCollections = await this.db.listAllParameterIndexCollectionsV3();
-    const v3_parameter_aggregates = await Promise.all(
-      v3ParameterCollections.map((collection) =>
-        collection
-          .aggregate([
-            {
-              $collStats: {
-                storageStats: {}
-              }
-            }
-          ])
-          .toArray()
-          .catch(ignoreNotExisting)
-      )
-    );
-
-    const v1_source_record_aggregate = await aggregateStaticCollection(this.db.current_data);
-
-    const v3SourceRecordCollections = await this.db.listAllSourceRecordCollectionsV3();
-    const source_record_aggregates = await Promise.all(
-      v3SourceRecordCollections.map((collection) =>
-        collection
-          .aggregate([
-            {
-              $collStats: {
-                storageStats: {}
-              }
-            }
-          ])
-          .toArray()
-          .catch(ignoreNotExisting)
+        },
+        { nameOnly: true, readPreference: this.options.statsReadPreference }
       )
-    );
+      .toArray();
+    const sizesByCollection = new Map<string, number>();
+    const statsPipeline: mongo.Document[] = [
+      { $collStats: { storageStats: {} } },
+      { $project: { ns: 1, size: '$storageStats.size' } }
+    ];
+    // MongoDB has a limit of 1000 stages per aggregation pipeline. Additionally, we want
+    // to avoid too much work per aggregation, so we limit the number of collections per batch.
+    const COLLECTIONS_PER_BATCH = 100;
+    // Query batches sequentially so metrics collection does not fill the connection pool on standby pods.
+    // If a collection disappears during aggregation, let the next metrics collection retry.
+    for (let offset = 0; offset < collections.length; offset += COLLECTIONS_PER_BATCH) {
+      const batch = collections.slice(offset, offset + COLLECTIONS_PER_BATCH);
+      const pipeline = [
+        ...statsPipeline,
+        ...batch.slice(1).map((collection) => ({
+          $unionWith: { coll: collection.name, pipeline: statsPipeline }
+        }))
+      ];
+      const stats = await this.db.db
+        .collection(batch[0].name)
+        .aggregate<{ ns: string; size: number | bigint }>(pipeline, {
+          readPreference: this.options.statsReadPreference
+        })
+        .toArray();
+      for (const stat of stats) {
+        const name = stat.ns.slice(this.db.db.databaseName.length + 1);
+        sizesByCollection.set(name, (sizesByCollection.get(name) ?? 0) + Number(stat.size));
+      }
+    }
 
     const v3StorageConfig = getMongoStorageConfig(storage.STORAGE_VERSION_3) as StorageConfig & {
       incrementalReprocessing: true;
@@ -922,29 +887,8 @@ export class MongoBucketStorage extends storage.BucketStorageFactory {
       .find({ storage_version: { $gte: storage.STORAGE_VERSION_3 } }, { projection: { _id: 1, sync_configs: 1 } })
       .toArray()) as unknown as Pick<ReplicationStreamDocumentV3, '_id' | 'sync_configs'>[];
 
-    const collectionSizes = new Map<string, number>();
-    const addCollectionSize = (
-      collection: { collectionName: string },
-      aggregate: { storageStats?: { size?: number | bigint } }[],
-      prefix: string,
-      sizes: Map<string, number>
-    ) => {
-      const match = collection.collectionName.match(new RegExp(`^${prefix}(\\d+)_(.+)$`));
-      if (match == null) {
-        return;
-      }
-      sizes.set(`${prefix}${match[1]}:${match[2]}`, Number(aggregate[0]?.storageStats?.size ?? 0));
-    };
-
-    v3OperationCollections.forEach((collection, index) =>
-      addCollectionSize(collection, v3_operation_aggregates[index], 'bucket_data_', collectionSizes)
-    );
-    v3ParameterCollections.forEach((collection, index) =>
-      addCollectionSize(collection, v3_parameter_aggregates[index], 'parameter_index_', collectionSizes)
-    );
-    v3SourceRecordCollections.forEach((collection, index) =>
-      addCollectionSize(collection, source_record_aggregates[index], 'source_records_', collectionSizes)
-    );
+    const sumCollectionSizesByPrefix = (prefix: string) =>
+      [...sizesByCollection].reduce((total, [name, size]) => total + (name.startsWith(prefix) ? size : 0), 0);
 
     const syncConfigIds = v3StreamDocs.flatMap((stream) => (stream.sync_configs ?? []).map((config) => config._id));
     const syncConfigDefinitions =
@@ -961,15 +905,13 @@ export class MongoBucketStorage extends storage.BucketStorageFactory {
       number,
       { _id: ObjectId; bucket_data_source_ids: string[]; parameter_lookup_source_ids: string[] }[]
     >();
-    await Promise.all(
-      v3StreamDocs.map(async (stream) => {
-        const sourceTables = await v3Db
-          .sourceTables(stream._id)
-          .find({}, { projection: { _id: 1, bucket_data_source_ids: 1, parameter_lookup_source_ids: 1 } })
-          .toArray();
-        sourceTablesByStream.set(stream._id, sourceTables);
-      })
-    );
+    for (const stream of v3StreamDocs) {
+      const sourceTables = await v3Db
+        .sourceTables(stream._id)
+        .find({}, { projection: { _id: 1, bucket_data_source_ids: 1, parameter_lookup_source_ids: 1 } })
+        .toArray();
+      sourceTablesByStream.set(stream._id, sourceTables);
+    }
 
     const objectStorageSizeByDefinition = new Map<string, number>();
     for (const usage of objectStorageDefinitionUsage) {
@@ -980,7 +922,7 @@ export class MongoBucketStorage extends storage.BucketStorageFactory {
     }
 
     const sumCollectionSizes = (prefix: string, streamId: number, ids: ReadonlySet<string>) =>
-      [...ids].reduce((total, id) => total + (collectionSizes.get(`${prefix}${streamId}:${id}`) ?? 0), 0);
+      [...ids].reduce((total, id) => total + (sizesByCollection.get(`${prefix}${streamId}_${id}`) ?? 0), 0);
 
     const syncConfigMetrics: storage.StorageSyncConfigMetrics[] = [];
     for (const stream of v3StreamDocs) {
@@ -1003,7 +945,7 @@ export class MongoBucketStorage extends storage.BucketStorageFactory {
           ) {
             return total;
           }
-          return total + (collectionSizes.get(`source_records_${stream._id}:${sourceTable._id.toHexString()}`) ?? 0);
+          return total + (sizesByCollection.get(`source_records_${stream._id}_${sourceTable._id.toHexString()}`) ?? 0);
         }, 0);
 
         syncConfigMetrics.push({
@@ -1026,15 +968,11 @@ export class MongoBucketStorage extends storage.BucketStorageFactory {
       0n
     );
     return {
-      operations_size_bytes:
-        Number(operations_aggregate[0].storageStats.size) +
-        v3_operation_aggregates.reduce((total, aggregate) => total + Number(aggregate[0].storageStats.size), 0),
+      operations_size_bytes: (sizesByCollection.get('bucket_data') ?? 0) + sumCollectionSizesByPrefix('bucket_data_'),
       parameters_size_bytes:
-        Number(parameters_aggregate[0].storageStats.size) +
-        v3_parameter_aggregates.reduce((total, aggregate) => total + Number(aggregate[0].storageStats.size), 0),
+        (sizesByCollection.get('bucket_parameters') ?? 0) + sumCollectionSizesByPrefix('parameter_index_'),
       replication_size_bytes:
-        Number(v1_source_record_aggregate[0]?.storageStats?.size ?? 0) +
-        source_record_aggregates.reduce((total, aggregate) => total + Number(aggregate[0]?.storageStats?.size ?? 0), 0),
+        (sizesByCollection.get('current_data') ?? 0) + sumCollectionSizesByPrefix('source_records_'),
       object_storage_size_bytes: Number(totalObjectStorageSize),
       sync_config_metrics: syncConfigMetrics
     };

package/src/storage/implementation/MongoStorageProvider.ts

@@ -78,6 +78,10 @@ export class MongoStorageProvider implements storage.StorageProvider {
       chunkCompactionConcurrency,
       replicationStreamNamePrefix: resolvedConfig.slot_name_prefix,
       readPreference,
+      statsReadPreference:
+        decodedConfig.stats_read_preference == null
+          ? undefined
+          : new lib_mongo.mongo.ReadPreference(decodedConfig.stats_read_preference),
       clearBatchThrottleRate: normalizeClearBatchThrottleRate(decodedConfig.clear_batch_throttle_rate),
       checksumCacheTtlMs: resolvedConfig.api_parameters.bucket_count_cache_ttl_minutes * 60_000,
       defaultStorageVersion: decodedConfig.default_storage_version,

package/src/storage/implementation/v1/MongoBucketBatchV1.ts

@@ -10,6 +10,7 @@ import { MongoIdSequence, OpIdRangeExhausted } from '../MongoIdSequence.js';
 import { MongoSyncRulesLock } from '../MongoSyncRulesLock.js';
 import { MongoWriteBatch } from '../MongoWriteBatch.js';
 import { OperationBatch } from '../OperationBatch.js';
+import { stopReplicationStreamPipeline } from '../SyncRuleStateUpdate.js';
 import { SourceRecordStore } from '../common/SourceRecordStore.js';
 import { SyncRuleDocumentBase } from '../models.js';
 import { PersistedBatchV1 } from './PersistedBatchV1.js';
@@ -853,11 +854,7 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
             _id: { $ne: this.replicationStreamId },
             state: { $in: [storage.SyncRuleState.ACTIVE, storage.SyncRuleState.ERRORED] }
           },
-          {
-            $set: {
-              state: storage.SyncRuleState.STOP
-            }
-          },
+          stopReplicationStreamPipeline(),
           { session }
         );
         activated = true;

package/src/types/types.ts

@@ -54,6 +54,12 @@ export const MongoStorageConfig = service_types.configFile.BaseStorageConfig.and
      * This is an experimental option, and may be removed in a future release.
      */
     bulk_read_preference: MongoStorageReadPreference.optional(),
+    /**
+     * Read preference for collection discovery and collection stats used by storage metrics.
+     * For example, use secondaryPreferred to offload collection stats from the primary.
+     * If unset, MongoDB driver defaults are used.
+     */
+    stats_read_preference: MongoStorageReadPreference.optional(),
     /**
      * Throttle clear operations after sync config deploys, by pausing between batches.
      * The rate produces a pause proportional to the previous batch duration.

package/test/src/storage_sync.test.ts

@@ -7,7 +7,13 @@ import {
   storage,
   updateSyncRulesFromYaml
 } from '@powersync/service-core';
-import { bucketRequest, getTestStorage, register, test_utils } from '@powersync/service-core-tests';
+import {
+  bucketRequest,
+  getTestStorage,
+  register,
+  releaseTestStorageLease,
+  test_utils
+} from '@powersync/service-core-tests';
 import { parseSyncConfig, RequestParameters, ScopedParameterLookup } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
 import { describe, expect, test } from 'vitest';
@@ -1951,6 +1957,24 @@ streams:
     }
   );
 
+  test.runIf(storageVersion >= 3)('storage metrics collect sizes across union batches', async () => {
+    await using factory = await storageConfig.factory();
+    const mongoFactory = factory as MongoBucketStorage;
+    const before = await factory.getStorageMetrics();
+    let expectedSize = 0;
+    for (let index = 0; index < 101; index++) {
+      const collection = mongoFactory.db.db.collection(`bucket_data_999_metrics_${index}`);
+      await collection.insertOne({ value: 'x'.repeat(index + 1) });
+      const [stats] = await collection.aggregate([{ $collStats: { storageStats: {} } }]).toArray();
+      expectedSize += Number(stats.storageStats.size);
+    }
+
+    const after = await factory.getStorageMetrics();
+    expect(after.operations_size_bytes - before.operations_size_bytes).toBe(expectedSize);
+    expect(after.parameters_size_bytes).toBe(before.parameters_size_bytes);
+    expect(after.replication_size_bytes).toBe(before.replication_size_bytes);
+  });
+
   test.runIf(storageVersion < 3)('storage metrics include v1 current_data', async () => {
     await using factory = await storageConfig.factory();
     const syncRules = await factory.updateSyncRules(
@@ -2160,6 +2184,91 @@ streams:
 }
 
 describe('sync - mongodb', () => {
+  test.each([
+    [1, false],
+    [2, false],
+    [1, true],
+    [2, true]
+  ] as const)('v%i activation cleans up v4 streams (previously stranded: %s)', async (storageVersion, stranded) => {
+    await using factory = await INITIALIZED_MONGO_STORAGE_FACTORY.factory();
+    const yaml = /* yaml */ `
+      # sync config fixture.
+      config:
+        edition: 3
+      streams:
+        global:
+          query: SELECT id FROM test
+    `;
+
+    async function activate(version: number, lsn: string) {
+      const stream = await factory.updateSyncRules(updateSyncRulesFromYaml(yaml, { storageVersion: version }));
+      const bucketStorage = await getTestStorage(factory, stream);
+      await using writer = await bucketStorage.createWriter(test_utils.BATCH_OPTIONS);
+      const sourceTable = await test_utils.resolveTestTable(writer, 'test', ['id'], INITIALIZED_MONGO_STORAGE_FACTORY);
+      await writer.save({
+        sourceTable,
+        tag: storage.SaveOperationTag.INSERT,
+        after: { id: 'item-1' },
+        afterReplicaId: test_utils.rid('item-1')
+      });
+      await writer.markAllSnapshotDone(lsn);
+      await writer.commit(lsn);
+      return stream;
+    }
+
+    const previous = await activate(storage.STORAGE_VERSION_4, '1/1');
+    await releaseTestStorageLease(factory, previous.replicationStreamId);
+    const oldCollections = await factory.db.db
+      .listCollections(
+        {
+          name: new RegExp(
+            `^(bucket_data|source_records|source_table|bucket_state|parameter_index)_${previous.replicationStreamId}(_|$)`
+          )
+        },
+        { nameOnly: true }
+      )
+      .toArray();
+    expect(oldCollections.some((collection) => collection.name.startsWith('bucket_data_'))).toBe(true);
+    expect(oldCollections.some((collection) => collection.name.startsWith('source_records_'))).toBe(true);
+
+    const replacement = await activate(storageVersion, '2/1');
+    const previousDoc = (await factory.db.sync_rules.findOne({
+      _id: previous.replicationStreamId
+    })) as ReplicationStreamDocumentV3;
+    expect(previousDoc.state).toBe(storage.SyncRuleState.STOP);
+    if (stranded) {
+      // Reproduce databases left behind by older services: only the stream was stopped.
+      await factory.db.sync_rules.updateOne(
+        { _id: previous.replicationStreamId },
+        { $set: { 'sync_configs.$[].state': storage.SyncRuleState.ACTIVE } }
+      );
+    } else {
+      expect(previousDoc.sync_configs.map((config) => config.state)).toEqual([storage.SyncRuleState.STOP]);
+    }
+
+    const stopped = await factory.getStoppedReplicationStreams();
+    expect(stopped.map((stream) => stream.replicationStreamId)).toEqual([previous.replicationStreamId]);
+    const lock = await stopped[0].lock();
+    try {
+      await factory.getInstance(stopped[0], { skipLifecycleHooks: true }).terminate({ clearStorage: true });
+    } finally {
+      await lock.release();
+    }
+
+    expect((await factory.db.sync_rules.findOne({ _id: previous.replicationStreamId }))?.state).toBe(
+      storage.SyncRuleState.TERMINATED
+    );
+    const remaining = new Set(
+      (await factory.db.db.listCollections({}, { nameOnly: true }).toArray()).map((c) => c.name)
+    );
+    for (const collection of oldCollections) expect(remaining.has(collection.name)).toBe(false);
+    expect(await factory.getStoppedReplicationStreams()).toEqual([]);
+    expect((await factory.getActiveSyncConfig())?.replicationStream.replicationStreamId).toBe(
+      replacement.replicationStreamId
+    );
+    expect(await factory.db.current_data.countDocuments({ '_id.g': replacement.replicationStreamId })).toBe(1);
+  });
+
   test('v3 activation stops legacy active sync rules', async () => {
     await using factory = await INITIALIZED_MONGO_STORAGE_FACTORY.factory();
     const mongoFactory = factory as MongoBucketStorage;