@powersync/service-module-mongodb-storage 0.0.0-dev-20260803144812 → 0.0.0-dev-20260929095716

package/src/storage/implementation/v3/object-storage/ObjectStorageLifecycle.ts

@@ -1,11 +1,11 @@
-import { mongo } from '@powersync/lib-service-mongodb';
 import { Logger } from '@powersync/lib-services-framework';
 import * as bson from 'bson';
 import { createHash, randomUUID } from 'node:crypto';
+import { MongoWriteBatch } from '../../MongoWriteBatch.js';
 import { ObjectStorageDeletionMarker } from '../models.js';
 import { VersionedPowerSyncMongoV3 } from '../VersionedPowerSyncMongoV3.js';
 import { BucketDataObjectStorage } from './BucketDataObjectStorage.js';
-import { ObjectStorage } from './ObjectStorage.js';
+import { ObjectStorage, ObjectStorageOperationOptions } from './ObjectStorage.js';
 
 export const OBJECT_STORAGE_UPLOAD_LEASE_MS = 60 * 60 * 1000;
 export const OBJECT_STORAGE_PUBLICATION_SAFETY_MARGIN_MS = 2 * 60 * 1000;
@@ -93,7 +93,7 @@ export class ObjectStorageLifecycle {
    *
    * Call this in the same transaction as the one creating the references.
    */
-  async publishUploads(uploads: PreparedObjectStorageUpload[], session: mongo.ClientSession): Promise<void> {
+  publishUploads(uploads: PreparedObjectStorageUpload[], writes: MongoWriteBatch): void {
     for (const upload of uploads) {
       if (!this.canPublish(upload)) {
         throw new Error(`Object storage publication lease expired for ${upload.path}`);
@@ -103,12 +103,13 @@ export class ObjectStorageLifecycle {
       return;
     }
 
-    const result = await this.db
-      .pendingObjectStorageDeletes(this.replicationStreamId)
-      .deleteMany({ _id: { $in: uploads.map((upload) => upload.markerId) } }, { session });
-    if (result.deletedCount !== uploads.length) {
-      throw new Error(`Missing object storage publication markers`);
-    }
+    writes.deleteMany(
+      this.db.pendingObjectStorageDeletes(this.replicationStreamId),
+      { _id: { $in: uploads.map((upload) => upload.markerId) } },
+      (deletedCount) => {
+        if (deletedCount !== uploads.length) throw new Error('Missing object storage publication markers');
+      }
+    );
   }
 
   /**
@@ -122,14 +123,14 @@ export class ObjectStorageLifecycle {
    *
    * Call this in the same transaction as the one removing the references to these files.
    */
-  async retire(paths: Iterable<string>, session: mongo.ClientSession, now = new Date()): Promise<void> {
+  retire(paths: Iterable<string>, writes: MongoWriteBatch, now = new Date()): void {
     const markers: ObjectStorageDeletionMarker[] = Array.from(paths, (path) => ({
       _id: new bson.ObjectId(),
       path,
       delete_after: new Date(now.getTime() + OBJECT_STORAGE_REFERENCE_GRACE_MS)
     }));
     if (markers.length) {
-      await this.db.pendingObjectStorageDeletes(this.replicationStreamId).insertMany(markers, { session });
+      writes.insertMany(this.db.pendingObjectStorageDeletes(this.replicationStreamId), markers);
     }
   }
 
@@ -137,11 +138,12 @@ export class ObjectStorageLifecycle {
    * Discover and delete objects by their storage prefix
    * without touching the database.
    */
-  async deletePrefix(prefix: string, options?: { signal?: AbortSignal }): Promise<{ objectCount: number }> {
+  async deletePrefix(prefix: string, options?: ObjectStorageOperationOptions): Promise<{ objectCount: number }> {
     return this.objectStorage.deletePrefix(prefix, { signal: options?.signal });
   }
 
-  async cleanup(logger: Logger): Promise<void> {
+  async cleanup(logger: Logger, options?: ObjectStorageOperationOptions): Promise<void> {
+    const signal = options?.signal;
     const markers = this.db.pendingObjectStorageDeletes(this.replicationStreamId);
     const cursor = markers.find({ delete_after: { $lte: new Date() } });
     let batch: ObjectStorageDeletionMarker[] = [];
@@ -153,14 +155,24 @@ export class ObjectStorageLifecycle {
       const deleting = batch;
       batch = [];
       try {
-        await this.bucketData.delete(deleting.map((marker) => marker.path));
+        await this.bucketData.delete(
+          deleting.map((marker) => marker.path),
+          { signal }
+        );
         await markers.deleteMany({ _id: { $in: deleting.map((marker) => marker._id) } });
       } catch (error) {
+        if (signal?.aborted) {
+          // Shutting down: stop instead of working through the remaining markers. They stay
+          // pending, and are cleaned up during the next compaction. Checked on the signal rather
+          // than the error, since a caller may abort with any reason.
+          throw error;
+        }
         logger.warn('Failed to clean up object storage deletion markers; will retry during the next compaction', error);
       }
     };
 
     for await (const marker of cursor) {
+      signal?.throwIfAborted();
       batch.push(marker);
       if (batch.length === 500) {
         await deleteBatch();

package/src/storage/implementation/v3/object-storage/ObjectStorageUsage.ts

@@ -0,0 +1,229 @@
+import * as lib_mongo from '@powersync/lib-service-mongodb';
+import { mongo } from '@powersync/lib-service-mongodb';
+import { ReplicationAssertionError } from '@powersync/lib-services-framework';
+import { BucketDefinitionId } from '@powersync/service-sync-rules';
+import { randomUUID } from 'node:crypto';
+import { MongoWriteBatch } from '../../MongoWriteBatch.js';
+import { BucketDataDocumentV3, ObjectStorageUsageDocument } from '../models.js';
+import { VersionedPowerSyncMongoV3 } from '../VersionedPowerSyncMongoV3.js';
+
+export const OBJECT_STORAGE_USAGE_BASE_WRITER_ID = '__base__';
+export const DEFAULT_OBJECT_STORAGE_USAGE_STALE_WRITER_MS = 30 * 60 * 1000;
+export const DEFAULT_OBJECT_STORAGE_USAGE_FOLD_LIMIT = 100;
+
+export interface ReplicationStreamObjectStorageDefinitionUsageResult {
+  replication_stream_id: number;
+  definition_id: BucketDefinitionId;
+  active_bytes: bigint;
+}
+
+export function createObjectStorageUsageWriterId(): string {
+  return randomUUID();
+}
+
+/**
+ * Tracks active S3 references. Values are signed writer-local deltas, so reference changes and
+ * accounting updates can be committed together without making replication and compaction share a
+ * hot counter document.
+ */
+export class ObjectStorageUsage {
+  constructor(
+    private readonly db: VersionedPowerSyncMongoV3,
+    private readonly replicationStreamId: number,
+    readonly writerId?: string
+  ) {
+    if (writerId === OBJECT_STORAGE_USAGE_BASE_WRITER_ID) {
+      throw new ReplicationAssertionError('The reserved object-storage usage writer id cannot be used by a writer');
+    }
+  }
+
+  static bytes(document: Pick<BucketDataDocumentV3, 'storage_ref'>): bigint {
+    return BigInt(document.storage_ref?.file_size ?? 0);
+  }
+
+  static async readAllDefinitionUsage(
+    db: VersionedPowerSyncMongoV3
+  ): Promise<ReplicationStreamObjectStorageDefinitionUsageResult[]> {
+    return db.client.withSession({ snapshot: true }, async (session) => {
+      // Deliberately read the whole usage collection. The number of replication streams is
+      // expected to stay low, so a collection scan is cheaper than issuing one _id range query
+      // per stream (and the usage collection is bounded by streams and writers, not buckets).
+      const entries = await db.objectStorageUsage
+        .aggregate<{
+          _id: { replication_stream_id: number; definition_id: BucketDefinitionId };
+          active_bytes: bigint;
+        }>(
+          [
+            { $project: { replication_stream_id: '$_id.g', definitions: { $objectToArray: '$definitions' } } },
+            { $unwind: '$definitions' },
+            {
+              $group: {
+                _id: {
+                  replication_stream_id: '$replication_stream_id',
+                  definition_id: '$definitions.k'
+                },
+                active_bytes: { $sum: '$definitions.v' }
+              }
+            }
+          ],
+          { session, readConcern: 'snapshot' }
+        )
+        .toArray()
+        .catch((error) => {
+          if (lib_mongo.isMongoNamespaceNotFoundError(error)) {
+            return [];
+          }
+          throw error;
+        });
+
+      return entries.map((entry) => {
+        const activeBytes = BigInt(entry.active_bytes ?? 0);
+        if (activeBytes < 0n) {
+          return {
+            replication_stream_id: entry._id.replication_stream_id,
+            definition_id: entry._id.definition_id,
+            active_bytes: 0n
+          };
+        }
+        return {
+          replication_stream_id: entry._id.replication_stream_id,
+          definition_id: entry._id.definition_id,
+          active_bytes: activeBytes
+        };
+      });
+    });
+  }
+
+  applyDelta(definitionId: BucketDefinitionId, delta: bigint, writes: MongoWriteBatch): void {
+    this.applyDeltas(new Map([[definitionId, delta]]), writes);
+  }
+
+  applyDeltas(deltas: ReadonlyMap<BucketDefinitionId, bigint>, writes: MongoWriteBatch): void {
+    if (this.writerId == null) {
+      throw new ReplicationAssertionError('A writer id is required to apply object-storage usage');
+    }
+
+    const increments: Record<string, bigint> = {};
+    for (const [definitionId, delta] of deltas) {
+      if (delta === 0n) {
+        continue;
+      }
+      increments[this.definitionPath(definitionId)] = delta;
+    }
+    if (Object.keys(increments).length === 0) {
+      return;
+    }
+
+    writes.updateOne(
+      this.db.objectStorageUsage,
+      { _id: this.documentId() },
+      { $inc: increments, $currentDate: { updated_at: true } },
+      { upsert: true }
+    );
+  }
+
+  async removeDefinition(definitionId: BucketDefinitionId, session: mongo.ClientSession): Promise<void> {
+    await this.db.objectStorageUsage.updateMany(
+      { '_id.g': this.replicationStreamId },
+      { $unset: { [this.definitionPath(definitionId)]: 1 } },
+      { session }
+    );
+  }
+
+  async removeStream(session: mongo.ClientSession): Promise<void> {
+    await this.db.objectStorageUsage.deleteMany({ '_id.g': this.replicationStreamId }, { session });
+  }
+
+  async foldStaleWriterDeltas(
+    options: {
+      staleWriterMs?: number;
+      limit?: number;
+    } = {}
+  ): Promise<void> {
+    const staleWriterMs = options.staleWriterMs ?? DEFAULT_OBJECT_STORAGE_USAGE_STALE_WRITER_MS;
+    const limit = options.limit ?? DEFAULT_OBJECT_STORAGE_USAGE_FOLD_LIMIT;
+    if (limit <= 0) {
+      return;
+    }
+
+    await this.db.client.withSession((session) =>
+      session.withTransaction(
+        async () => {
+          const staleDocuments = await this.db.objectStorageUsage
+            .find(
+              {
+                '_id.g': this.replicationStreamId,
+                '_id.w': { $ne: OBJECT_STORAGE_USAGE_BASE_WRITER_ID },
+                $expr: {
+                  $lt: [
+                    '$updated_at',
+                    { $dateSubtract: { startDate: '$$NOW', unit: 'millisecond', amount: staleWriterMs } }
+                  ]
+                }
+              },
+              { session, sort: { updated_at: 1 }, limit }
+            )
+            .toArray();
+          if (staleDocuments.length === 0) {
+            return;
+          }
+
+          const deltas = new Map<BucketDefinitionId, bigint>();
+          for (const document of staleDocuments) {
+            for (const [definitionId, delta] of Object.entries(document.definitions ?? {})) {
+              this.validateDefinitionId(definitionId);
+              deltas.set(definitionId, (deltas.get(definitionId) ?? 0n) + BigInt(delta));
+            }
+          }
+
+          const increments: Record<string, bigint> = {};
+          for (const [definitionId, delta] of deltas) {
+            if (delta !== 0n) {
+              increments[this.definitionPath(definitionId)] = delta;
+            }
+          }
+          if (Object.keys(increments).length > 0) {
+            await this.db.objectStorageUsage.updateOne(
+              { _id: this.documentId(OBJECT_STORAGE_USAGE_BASE_WRITER_ID) },
+              {
+                $inc: increments,
+                $currentDate: { updated_at: true }
+              },
+              { upsert: true, session }
+            );
+          }
+
+          await this.db.objectStorageUsage.deleteMany(
+            {
+              '_id.g': this.replicationStreamId,
+              _id: { $in: staleDocuments.map((document) => document._id) }
+            },
+            { session }
+          );
+        },
+        { readConcern: { level: 'snapshot' }, writeConcern: { w: 'majority' } }
+      )
+    );
+  }
+
+  private definitionPath(definitionId: BucketDefinitionId): string {
+    this.validateDefinitionId(definitionId);
+    return `definitions.${definitionId}`;
+  }
+
+  private documentId(writerId = this.writerId): ObjectStorageUsageDocument['_id'] {
+    if (writerId == null) {
+      throw new ReplicationAssertionError('A writer id is required to identify object-storage usage');
+    }
+    return {
+      g: this.replicationStreamId,
+      w: writerId
+    };
+  }
+
+  private validateDefinitionId(definitionId: BucketDefinitionId): void {
+    if (definitionId.length === 0 || definitionId.includes('.') || definitionId.includes('$')) {
+      throw new ReplicationAssertionError(`Invalid bucket definition id for object-storage usage: ${definitionId}`);
+    }
+  }
+}