@powersync/service-module-mongodb-storage 0.14.0 → 0.15.1

package/dist/storage/implementation/MongoBucketBatch.js

@@ -1,9 +1,9 @@
 import { mongo } from '@powersync/lib-service-mongodb';
 import * as bson from 'bson';
 import { BaseObserver, container, logger as defaultLogger, ErrorCode, errors, ReplicationAssertionError, ServiceError } from '@powersync/lib-services-framework';
-import { deserializeBson, isCompleteRow, SaveOperationTag, storage, utils } from '@powersync/service-core';
+import { deserializeBson, isCompleteRow, SaveOperationTag, storage, SyncRuleState, utils } from '@powersync/service-core';
 import * as timers from 'node:timers/promises';
-import { idPrefixFilter } from '../../utils/util.js';
+import { idPrefixFilter, mongoTableId } from '../../utils/util.js';
 import { MongoIdSequence } from './MongoIdSequence.js';
 import { batchCreateCustomWriteCheckpoints } from './MongoWriteCheckpointAPI.js';
 import { cacheKey, OperationBatch, RecordOperation } from './OperationBatch.js';
@@ -18,6 +18,7 @@ export const MAX_ROW_SIZE = 15 * 1024 * 1024;
 //
 // In the future, we can investigate allowing multiple replication streams operating independently.
 const replicationMutex = new utils.Mutex();
+export const EMPTY_DATA = new bson.Binary(bson.serialize({}));
 export class MongoBucketBatch extends BaseObserver {
     logger;
     client;
@@ -40,7 +41,6 @@ export class MongoBucketBatch extends BaseObserver {
      * 2. A keepalive message LSN.
      */
     last_checkpoint_lsn = null;
-    no_checkpoint_before_lsn;
     persisted_op = null;
     /**
      * Last written op, if any. This may not reflect a consistent checkpoint.
@@ -65,7 +65,6 @@ export class MongoBucketBatch extends BaseObserver {
         this.db = options.db;
         this.group_id = options.groupId;
         this.last_checkpoint_lsn = options.lastCheckpointLsn;
-        this.no_checkpoint_before_lsn = options.noCheckpointBeforeLsn;
         this.resumeFromLsn = options.resumeFromLsn;
         this.session = this.client.startSession();
         this.slot_name = options.slotName;
@@ -85,9 +84,6 @@ export class MongoBucketBatch extends BaseObserver {
     get lastCheckpointLsn() {
         return this.last_checkpoint_lsn;
     }
-    get noCheckpointBeforeLsn() {
-        return this.no_checkpoint_before_lsn;
-    }
     async flush(options) {
         let result = null;
         // One flush may be split over multiple transactions.
@@ -139,10 +135,10 @@ export class MongoBucketBatch extends BaseObserver {
             // the order of processing, which then becomes really tricky to manage.
             // This now takes 2+ queries, but doesn't have any issues with order of operations.
             const sizeLookups = batch.batch.map((r) => {
-                return { g: this.group_id, t: r.record.sourceTable.id, k: r.beforeId };
+                return { g: this.group_id, t: mongoTableId(r.record.sourceTable.id), k: r.beforeId };
             });
             sizes = new Map();
-            const sizeCursor = this.db.current_data.aggregate([
+            const sizeCursor = this.db.common_current_data.aggregate([
                 {
                     $match: {
                         _id: { $in: sizeLookups }
@@ -174,18 +170,18 @@ export class MongoBucketBatch extends BaseObserver {
                 continue;
             }
             const lookups = b.map((r) => {
-                return { g: this.group_id, t: r.record.sourceTable.id, k: r.beforeId };
+                return { g: this.group_id, t: mongoTableId(r.record.sourceTable.id), k: r.beforeId };
             });
             let current_data_lookup = new Map();
             // With skipExistingRows, we only need to know whether or not the row exists.
             const projection = this.skipExistingRows ? { _id: 1 } : undefined;
-            const cursor = this.db.current_data.find({
+            const cursor = this.db.common_current_data.find({
                 _id: { $in: lookups }
             }, { session, projection });
             for await (let doc of cursor.stream()) {
                 current_data_lookup.set(cacheKey(doc._id.t, doc._id.k), doc);
             }
-            let persistedBatch = new PersistedBatch(this.group_id, transactionSize, {
+            let persistedBatch = new PersistedBatch(this.db, this.group_id, transactionSize, {
                 logger: this.logger
             });
             for (let op of b) {
@@ -207,7 +203,7 @@ export class MongoBucketBatch extends BaseObserver {
                 if (persistedBatch.shouldFlushTransaction()) {
                     // Transaction is getting big.
                     // Flush, and resume in a new transaction.
-                    const { flushedAny } = await persistedBatch.flush(this.db, this.session, options);
+                    const { flushedAny } = await persistedBatch.flush(this.session, options);
                     didFlush ||= flushedAny;
                     persistedBatch = null;
                     // Computing our current progress is a little tricky here, since
@@ -218,7 +214,7 @@ export class MongoBucketBatch extends BaseObserver {
             }
             if (persistedBatch) {
                 transactionSize = persistedBatch.currentSize;
-                const { flushedAny } = await persistedBatch.flush(this.db, this.session, options);
+                const { flushedAny } = await persistedBatch.flush(this.session, options);
                 didFlush ||= flushedAny;
             }
         }
@@ -237,7 +233,7 @@ export class MongoBucketBatch extends BaseObserver {
         let new_buckets = [];
         let existing_lookups = [];
         let new_lookups = [];
-        const before_key = { g: this.group_id, t: record.sourceTable.id, k: beforeId };
+        const before_key = { g: this.group_id, t: mongoTableId(record.sourceTable.id), k: beforeId };
         if (this.skipExistingRows) {
             if (record.tag == SaveOperationTag.INSERT) {
                 if (current_data != null) {
@@ -298,7 +294,7 @@ export class MongoBucketBatch extends BaseObserver {
         }
         let afterData;
         if (afterId != null && !this.storeCurrentData) {
-            afterData = new bson.Binary(bson.serialize({}));
+            afterData = EMPTY_DATA;
         }
         else if (afterId != null) {
             try {
@@ -419,7 +415,7 @@ export class MongoBucketBatch extends BaseObserver {
         // 5. TOAST: Update current data and bucket list.
         if (afterId) {
             // Insert or update
-            const after_key = { g: this.group_id, t: sourceTable.id, k: afterId };
+            const after_key = { g: this.group_id, t: mongoTableId(sourceTable.id), k: afterId };
             batch.upsertCurrentData(after_key, {
                 data: afterData,
                 buckets: new_buckets,
@@ -434,7 +430,10 @@ export class MongoBucketBatch extends BaseObserver {
         }
         if (afterId == null || !storage.replicaIdEquals(beforeId, afterId)) {
             // Either a delete (afterId == null), or replaced the old replication id
-            batch.deleteCurrentData(before_key);
+            // Note that this is a soft delete.
+            // We don't specifically need a new or unique op_id here, but it must be greater than the
+            // last checkpoint, so we use next().
+            batch.softDeleteCurrentData(before_key, opSeq.next());
         }
         return result;
     }
@@ -505,59 +504,23 @@ export class MongoBucketBatch extends BaseObserver {
         });
     }
     async [Symbol.asyncDispose]() {
+        if (this.batch != null || this.write_checkpoint_batch.length > 0) {
+            // We don't error here, since:
+            // 1. In error states, this is expected (we can't distinguish between disposing after success or error).
+            // 2. SuppressedError is messy to deal with.
+            this.logger.warn('Disposing writer with unflushed changes');
+        }
         await this.session.endSession();
         super.clearListeners();
     }
+    async dispose() {
+        await this[Symbol.asyncDispose]();
+    }
     lastWaitingLogThottled = 0;
     async commit(lsn, options) {
         const { createEmptyCheckpoints } = { ...storage.DEFAULT_BUCKET_BATCH_COMMIT_OPTIONS, ...options };
         await this.flush(options);
-        if (this.last_checkpoint_lsn != null && lsn < this.last_checkpoint_lsn) {
-            // When re-applying transactions, don't create a new checkpoint until
-            // we are past the last transaction.
-            this.logger.info(`Re-applied transaction ${lsn} - skipping checkpoint`);
-            // Cannot create a checkpoint yet - return false
-            return false;
-        }
-        if (lsn < this.no_checkpoint_before_lsn) {
-            if (Date.now() - this.lastWaitingLogThottled > 5_000) {
-                this.logger.info(`Waiting until ${this.no_checkpoint_before_lsn} before creating checkpoint, currently at ${lsn}. Persisted op: ${this.persisted_op}`);
-                this.lastWaitingLogThottled = Date.now();
-            }
-            // Edge case: During initial replication, we have a no_checkpoint_before_lsn set,
-            // and don't actually commit the snapshot.
-            // The first commit can happen from an implicit keepalive message.
-            // That needs the persisted_op to get an accurate checkpoint, so
-            // we persist that in keepalive_op.
-            await this.db.sync_rules.updateOne({
-                _id: this.group_id
-            }, {
-                $set: {
-                    keepalive_op: this.persisted_op == null ? null : String(this.persisted_op)
-                }
-            }, { session: this.session });
-            await this.db.notifyCheckpoint();
-            // Cannot create a checkpoint yet - return false
-            return false;
-        }
-        if (!createEmptyCheckpoints && this.persisted_op == null) {
-            // Nothing to commit - also return true
-            await this.autoActivate(lsn);
-            return true;
-        }
         const now = new Date();
-        const update = {
-            last_checkpoint_lsn: lsn,
-            last_checkpoint_ts: now,
-            last_keepalive_ts: now,
-            snapshot_done: true,
-            last_fatal_error: null,
-            last_fatal_error_ts: null,
-            keepalive_op: null
-        };
-        if (this.persisted_op != null) {
-            update.last_checkpoint = this.persisted_op;
-        }
         // Mark relevant write checkpoints as "processed".
         // This makes it easier to identify write checkpoints that are "valid" in order.
         await this.db.write_checkpoints.updateMany({
@@ -570,17 +533,143 @@ export class MongoBucketBatch extends BaseObserver {
         }, {
             session: this.session
         });
-        await this.db.sync_rules.updateOne({
-            _id: this.group_id
-        }, {
-            $set: update,
-            $unset: { snapshot_lsn: 1 }
-        }, { session: this.session });
-        await this.autoActivate(lsn);
-        await this.db.notifyCheckpoint();
-        this.persisted_op = null;
-        this.last_checkpoint_lsn = lsn;
-        return true;
+        const can_checkpoint = {
+            $and: [
+                { $eq: ['$snapshot_done', true] },
+                {
+                    $or: [{ $eq: ['$last_checkpoint_lsn', null] }, { $lte: ['$last_checkpoint_lsn', { $literal: lsn }] }]
+                },
+                {
+                    $or: [{ $eq: ['$no_checkpoint_before', null] }, { $lte: ['$no_checkpoint_before', { $literal: lsn }] }]
+                }
+            ]
+        };
+        const new_keepalive_op = {
+            $cond: [
+                can_checkpoint,
+                { $literal: null },
+                {
+                    $toString: {
+                        $max: [{ $toLong: '$keepalive_op' }, { $literal: this.persisted_op }, 0n]
+                    }
+                }
+            ]
+        };
+        const new_last_checkpoint = {
+            $cond: [
+                can_checkpoint,
+                {
+                    $max: ['$last_checkpoint', { $literal: this.persisted_op }, { $toLong: '$keepalive_op' }, 0n]
+                },
+                '$last_checkpoint'
+            ]
+        };
+        // For this query, we need to handle multiple cases, depending on the state:
+        // 1. Normal commit - advance last_checkpoint to this.persisted_op.
+        // 2. Commit delayed by no_checkpoint_before due to snapshot. In this case we only advance keepalive_op.
+        // 3. Commit with no new data - here may may set last_checkpoint = keepalive_op, if a delayed commit is relevant.
+        // We want to do as much as possible in a single atomic database operation, which makes this somewhat complex.
+        let preUpdateDocument = await this.db.sync_rules.findOneAndUpdate({ _id: this.group_id }, [
+            {
+                $set: {
+                    _can_checkpoint: can_checkpoint,
+                    _not_empty: createEmptyCheckpoints
+                        ? true
+                        : {
+                            $or: [
+                                { $literal: createEmptyCheckpoints },
+                                { $ne: ['$keepalive_op', new_keepalive_op] },
+                                { $ne: ['$last_checkpoint', new_last_checkpoint] }
+                            ]
+                        }
+                }
+            },
+            {
+                $set: {
+                    last_checkpoint_lsn: {
+                        $cond: [{ $and: ['$_can_checkpoint', '$_not_empty'] }, { $literal: lsn }, '$last_checkpoint_lsn']
+                    },
+                    last_checkpoint_ts: {
+                        $cond: [{ $and: ['$_can_checkpoint', '$_not_empty'] }, { $literal: now }, '$last_checkpoint_ts']
+                    },
+                    last_keepalive_ts: { $literal: now },
+                    last_fatal_error: { $literal: null },
+                    last_fatal_error_ts: { $literal: null },
+                    keepalive_op: new_keepalive_op,
+                    last_checkpoint: new_last_checkpoint,
+                    // Unset snapshot_lsn on checkpoint
+                    snapshot_lsn: {
+                        $cond: [{ $and: ['$_can_checkpoint', '$_not_empty'] }, { $literal: null }, '$snapshot_lsn']
+                    }
+                }
+            },
+            {
+                $unset: ['_can_checkpoint', '_not_empty']
+            }
+        ], {
+            session: this.session,
+            // We return the before document, so that we can check the previous state to determine if a checkpoint was actually created or if we were blocked by snapshot/no_checkpoint_before.
+            returnDocument: 'before',
+            projection: {
+                snapshot_done: 1,
+                last_checkpoint_lsn: 1,
+                no_checkpoint_before: 1,
+                keepalive_op: 1,
+                last_checkpoint: 1
+            }
+        });
+        if (preUpdateDocument == null) {
+            throw new ReplicationAssertionError('Failed to update checkpoint - no matching sync_rules document for _id: ' + this.group_id);
+        }
+        // This re-implements the same logic as in the pipeline, to determine what was actually updated.
+        // Unfortunately we cannot return these from the pipeline directly, so we need to re-implement the logic.
+        const canCheckpoint = preUpdateDocument.snapshot_done === true &&
+            (preUpdateDocument.last_checkpoint_lsn == null || preUpdateDocument.last_checkpoint_lsn <= lsn) &&
+            (preUpdateDocument.no_checkpoint_before == null || preUpdateDocument.no_checkpoint_before <= lsn);
+        const keepaliveOp = preUpdateDocument.keepalive_op == null ? null : BigInt(preUpdateDocument.keepalive_op);
+        const maxKeepalive = [keepaliveOp ?? 0n, this.persisted_op ?? 0n, 0n].reduce((a, b) => (a > b ? a : b));
+        const newKeepaliveOp = canCheckpoint ? null : maxKeepalive.toString();
+        const newLastCheckpoint = canCheckpoint
+            ? [preUpdateDocument.last_checkpoint ?? 0n, this.persisted_op ?? 0n, keepaliveOp ?? 0n, 0n].reduce((a, b) => a > b ? a : b)
+            : preUpdateDocument.last_checkpoint;
+        const notEmpty = createEmptyCheckpoints ||
+            preUpdateDocument.keepalive_op !== newKeepaliveOp ||
+            preUpdateDocument.last_checkpoint !== newLastCheckpoint;
+        const checkpointCreated = canCheckpoint && notEmpty;
+        const checkpointBlocked = !canCheckpoint;
+        if (checkpointBlocked) {
+            // Failed on snapshot_done or no_checkpoint_before.
+            if (Date.now() - this.lastWaitingLogThottled > 5_000) {
+                this.logger.info(`Waiting before creating checkpoint, currently at ${lsn} / ${preUpdateDocument.keepalive_op}. Current state: ${JSON.stringify({
+                    snapshot_done: preUpdateDocument.snapshot_done,
+                    last_checkpoint_lsn: preUpdateDocument.last_checkpoint_lsn,
+                    no_checkpoint_before: preUpdateDocument.no_checkpoint_before
+                })}`);
+                this.lastWaitingLogThottled = Date.now();
+            }
+        }
+        else {
+            if (checkpointCreated) {
+                this.logger.debug(`Created checkpoint at ${lsn} / ${newLastCheckpoint}`);
+            }
+            await this.autoActivate(lsn);
+            await this.db.notifyCheckpoint();
+            this.persisted_op = null;
+            this.last_checkpoint_lsn = lsn;
+            if (this.db.storageConfig.softDeleteCurrentData && newLastCheckpoint != null) {
+                await this.cleanupCurrentData(newLastCheckpoint);
+            }
+        }
+        return { checkpointBlocked, checkpointCreated };
+    }
+    async cleanupCurrentData(lastCheckpoint) {
+        const result = await this.db.v3_current_data.deleteMany({
+            '_id.g': this.group_id,
+            pending_delete: { $exists: true, $lte: lastCheckpoint }
+        });
+        if (result.deletedCount > 0) {
+            this.logger.info(`Cleaned up ${result.deletedCount} pending delete current_data records for checkpoint ${lastCheckpoint}`);
+        }
     }
     /**
      * Switch from processing -> active if relevant.
@@ -597,7 +686,7 @@ export class MongoBucketBatch extends BaseObserver {
         let activated = false;
         await session.withTransaction(async () => {
             const doc = await this.db.sync_rules.findOne({ _id: this.group_id }, { session });
-            if (doc && doc.state == 'PROCESSING') {
+            if (doc && doc.state == SyncRuleState.PROCESSING && doc.snapshot_done && doc.last_checkpoint != null) {
                 await this.db.sync_rules.updateOne({
                     _id: this.group_id
                 }, {
@@ -615,53 +704,18 @@ export class MongoBucketBatch extends BaseObserver {
                 }, { session });
                 activated = true;
             }
+            else if (doc?.state != SyncRuleState.PROCESSING) {
+                this.needsActivation = false;
+            }
         });
         if (activated) {
             this.logger.info(`Activated new sync rules at ${lsn}`);
             await this.db.notifyCheckpoint();
+            this.needsActivation = false;
         }
-        this.needsActivation = false;
     }
     async keepalive(lsn) {
-        if (this.last_checkpoint_lsn != null && lsn < this.last_checkpoint_lsn) {
-            // No-op
-            return false;
-        }
-        if (lsn < this.no_checkpoint_before_lsn) {
-            return false;
-        }
-        if (this.persisted_op != null) {
-            // The commit may have been skipped due to "no_checkpoint_before_lsn".
-            // Apply it now if relevant
-            this.logger.info(`Commit due to keepalive at ${lsn} / ${this.persisted_op}`);
-            return await this.commit(lsn);
-        }
-        await this.db.write_checkpoints.updateMany({
-            processed_at_lsn: null,
-            'lsns.1': { $lte: lsn }
-        }, {
-            $set: {
-                processed_at_lsn: lsn
-            }
-        }, {
-            session: this.session
-        });
-        await this.db.sync_rules.updateOne({
-            _id: this.group_id
-        }, {
-            $set: {
-                last_checkpoint_lsn: lsn,
-                snapshot_done: true,
-                last_fatal_error: null,
-                last_fatal_error_ts: null,
-                last_keepalive_ts: new Date()
-            },
-            $unset: { snapshot_lsn: 1 }
-        }, { session: this.session });
-        await this.autoActivate(lsn);
-        await this.db.notifyCheckpoint();
-        this.last_checkpoint_lsn = lsn;
-        return true;
+        return await this.commit(lsn, { createEmptyCheckpoints: true });
     }
     async setResumeLsn(lsn) {
         const update = {
@@ -712,7 +766,7 @@ export class MongoBucketBatch extends BaseObserver {
         const result = await this.flush();
         await this.withTransaction(async () => {
             for (let table of sourceTables) {
-                await this.db.source_tables.deleteOne({ _id: table.id });
+                await this.db.source_tables.deleteOne({ _id: mongoTableId(table.id) });
             }
         });
         return result;
@@ -742,9 +796,12 @@ export class MongoBucketBatch extends BaseObserver {
         while (lastBatchCount == BATCH_LIMIT) {
             await this.withReplicationTransaction(`Truncate ${sourceTable.qualifiedName}`, async (session, opSeq) => {
                 const current_data_filter = {
-                    _id: idPrefixFilter({ g: this.group_id, t: sourceTable.id }, ['k'])
+                    _id: idPrefixFilter({ g: this.group_id, t: mongoTableId(sourceTable.id) }, ['k']),
+                    // Skip soft-deleted data
+                    // Works for both v1 and v3 current_data schemas
+                    pending_delete: { $exists: false }
                 };
-                const cursor = this.db.current_data.find(current_data_filter, {
+                const cursor = this.db.common_current_data.find(current_data_filter, {
                     projection: {
                         _id: 1,
                         buckets: 1,
@@ -754,7 +811,7 @@ export class MongoBucketBatch extends BaseObserver {
                     session: session
                 });
                 const batch = await cursor.toArray();
-                const persistedBatch = new PersistedBatch(this.group_id, 0, { logger: this.logger });
+                const persistedBatch = new PersistedBatch(this.db, this.group_id, 0, { logger: this.logger });
                 for (let value of batch) {
                     persistedBatch.saveBucketData({
                         op_seq: opSeq,
@@ -770,9 +827,10 @@ export class MongoBucketBatch extends BaseObserver {
                         sourceTable: sourceTable,
                         sourceKey: value._id.k
                     });
-                    persistedBatch.deleteCurrentData(value._id);
+                    // Since this is not from streaming replication, we can do a hard delete
+                    persistedBatch.hardDeleteCurrentData(value._id);
                 }
-                await persistedBatch.flush(this.db, session);
+                await persistedBatch.flush(session);
                 lastBatchCount = batch.length;
                 last_op = opSeq.last();
             });
@@ -788,7 +846,7 @@ export class MongoBucketBatch extends BaseObserver {
         };
         copy.snapshotStatus = snapshotStatus;
         await this.withTransaction(async () => {
-            await this.db.source_tables.updateOne({ _id: table.id }, {
+            await this.db.source_tables.updateOne({ _id: mongoTableId(table.id) }, {
                 $set: {
                     snapshot_status: {
                         last_key: snapshotStatus.lastKey == null ? null : new bson.Binary(snapshotStatus.lastKey),
@@ -800,9 +858,31 @@ export class MongoBucketBatch extends BaseObserver {
         });
         return copy;
     }
-    async markSnapshotDone(tables, no_checkpoint_before_lsn) {
+    async markAllSnapshotDone(no_checkpoint_before_lsn) {
+        await this.db.sync_rules.updateOne({
+            _id: this.group_id
+        }, {
+            $set: {
+                snapshot_done: true,
+                last_keepalive_ts: new Date()
+            },
+            $max: {
+                no_checkpoint_before: no_checkpoint_before_lsn
+            }
+        }, { session: this.session });
+    }
+    async markTableSnapshotRequired(table) {
+        await this.db.sync_rules.updateOne({
+            _id: this.group_id
+        }, {
+            $set: {
+                snapshot_done: false
+            }
+        }, { session: this.session });
+    }
+    async markTableSnapshotDone(tables, no_checkpoint_before_lsn) {
         const session = this.session;
-        const ids = tables.map((table) => table.id);
+        const ids = tables.map((table) => mongoTableId(table.id));
         await this.withTransaction(async () => {
             await this.db.source_tables.updateMany({ _id: { $in: ids } }, {
                 $set: {
@@ -812,14 +892,15 @@ export class MongoBucketBatch extends BaseObserver {
                     snapshot_status: 1
                 }
             }, { session });
-            if (no_checkpoint_before_lsn > this.no_checkpoint_before_lsn) {
-                this.no_checkpoint_before_lsn = no_checkpoint_before_lsn;
+            if (no_checkpoint_before_lsn != null) {
                 await this.db.sync_rules.updateOne({
                     _id: this.group_id
                 }, {
                     $set: {
-                        no_checkpoint_before: no_checkpoint_before_lsn,
                         last_keepalive_ts: new Date()
+                    },
+                    $max: {
+                        no_checkpoint_before: no_checkpoint_before_lsn
                     }
                 }, { session: this.session });
             }