@powersync/service-module-mongodb-storage 0.0.0-dev-20260225160713 → 0.0.0-dev-20260511080634

package/dist/storage/implementation/MongoBucketBatch.js

@@ -1,17 +1,13 @@
 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 { BaseObserver, container, ErrorCode, errors, ReplicationAssertionError, ServiceError } from '@powersync/lib-services-framework';
+import { deserializeBson, isCompleteRow, PerformanceTracer, SaveOperationTag, storage, SyncRuleState, utils } from '@powersync/service-core';
 import * as timers from 'node:timers/promises';
-import { idPrefixFilter } from '../../utils/util.js';
+import { mongoTableId } from '../../utils/util.js';
+import { MAX_ROW_SIZE } from './MongoBucketBatchShared.js';
 import { MongoIdSequence } from './MongoIdSequence.js';
 import { batchCreateCustomWriteCheckpoints } from './MongoWriteCheckpointAPI.js';
-import { cacheKey, OperationBatch, RecordOperation } from './OperationBatch.js';
-import { PersistedBatch } from './PersistedBatch.js';
-/**
- * 15MB
- */
-export const MAX_ROW_SIZE = 15 * 1024 * 1024;
+import { OperationBatch, RecordOperation } from './OperationBatch.js';
 // Currently, we can only have a single flush() at a time, since it locks the op_id sequence.
 // While the MongoDB transaction retry mechanism handles this okay, using an in-process Mutex
 // makes it more fair and has less overhead.
@@ -28,10 +24,12 @@ export class MongoBucketBatch extends BaseObserver {
     slot_name;
     storeCurrentData;
     skipExistingRows;
+    mapping;
     batch = null;
     write_checkpoint_batch = [];
     markRecordUnavailable;
     clearedError = false;
+    tracer;
     /**
      * Last LSN received associated with a checkpoint.
      *
@@ -40,7 +38,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.
@@ -60,21 +57,22 @@ export class MongoBucketBatch extends BaseObserver {
     needsActivation = true;
     constructor(options) {
         super();
-        this.logger = options.logger ?? defaultLogger;
+        this.logger = options.logger;
         this.client = options.db.client;
         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;
         this.sync_rules = options.syncRules;
         this.storeCurrentData = options.storeCurrentData;
+        this.mapping = options.mapping;
         this.skipExistingRows = options.skipExistingRows;
         this.markRecordUnavailable = options.markRecordUnavailable;
         this.batch = new OperationBatch();
         this.persisted_op = options.keepaliveOp ?? null;
+        this.tracer = options.tracer ?? new PerformanceTracer('MongoDB storage');
     }
     addCustomWriteCheckpoint(checkpoint) {
         this.write_checkpoint_batch.push({
@@ -85,9 +83,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.
@@ -104,6 +99,7 @@ export class MongoBucketBatch extends BaseObserver {
         const batch = this.batch;
         let last_op = null;
         let resumeBatch = null;
+        using _ = this.tracer.span('storage', 'flush');
         await this.withReplicationTransaction(`Flushing ${batch?.length ?? 0} ops`, async (session, opSeq) => {
             if (batch != null) {
                 resumeBatch = await this.replicateBatch(session, batch, opSeq, options);
@@ -126,6 +122,7 @@ export class MongoBucketBatch extends BaseObserver {
     }
     async replicateBatch(session, batch, op_seq, options) {
         let sizes = undefined;
+        using _ = this.tracer.span('storage', 'replicate_batch');
         if (this.storeCurrentData && !this.skipExistingRows) {
             // We skip this step if we don't store current_data, since the sizes will
             // always be small in that case.
@@ -138,27 +135,11 @@ export class MongoBucketBatch extends BaseObserver {
             // (automatically limited to 48MB(?) per batch by MongoDB). The issue is that it changes
             // 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 };
-            });
-            sizes = new Map();
-            const sizeCursor = this.db.current_data.aggregate([
-                {
-                    $match: {
-                        _id: { $in: sizeLookups }
-                    }
-                },
-                {
-                    $project: {
-                        _id: 1,
-                        size: { $bsonSize: '$$ROOT' }
-                    }
-                }
-            ], { session });
-            for await (let doc of sizeCursor.stream()) {
-                const key = cacheKey(doc._id.t, doc._id.k);
-                sizes.set(key, doc.size);
-            }
+            const sizeLookups = batch.batch.map((r) => ({
+                sourceTableId: mongoTableId(r.record.sourceTable.id),
+                replicaId: r.beforeId
+            }));
+            sizes = await this.sourceRecordStore.loadSizes(session, sizeLookups);
         }
         // If set, we need to start a new transaction with this batch.
         let resumeBatch = null;
@@ -173,61 +154,64 @@ export class MongoBucketBatch extends BaseObserver {
                 }
                 continue;
             }
-            const lookups = b.map((r) => {
-                return { g: this.group_id, t: 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({
-                _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, {
-                logger: this.logger
-            });
+            using lookupSpan = this.tracer.span('storage', 'lookup');
+            const lookups = b.map((r) => ({
+                sourceTableId: mongoTableId(r.record.sourceTable.id),
+                replicaId: r.beforeId
+            }));
+            let sourceRecordLookup = await this.sourceRecordStore.loadDocuments(session, lookups, this.skipExistingRows);
+            lookupSpan.end();
+            let persistedBatch = this.createPersistedBatch(transactionSize);
+            // The current code structure makes it tricky to cleanly split this span from the one
+            // where fluhsing. So we manually end and re-create this span whenever we flush.
+            let evalSpan = this.tracer.span('evaluate');
             for (let op of b) {
                 if (resumeBatch) {
                     resumeBatch.push(op);
                     continue;
                 }
-                const currentData = current_data_lookup.get(op.internalBeforeKey) ?? null;
-                if (currentData != null) {
+                const sourceRecord = sourceRecordLookup.get(op.internalBeforeKey) ?? null;
+                if (sourceRecord != null) {
                     // If it will be used again later, it will be set again using nextData below
-                    current_data_lookup.delete(op.internalBeforeKey);
+                    sourceRecordLookup.delete(op.internalBeforeKey);
                 }
-                const nextData = this.saveOperation(persistedBatch, op, currentData, op_seq);
+                const nextData = this.saveOperation(persistedBatch, op, sourceRecord, op_seq);
                 if (nextData != null) {
                     // Update our current_data and size cache
-                    current_data_lookup.set(op.internalAfterKey, nextData);
-                    sizes?.set(op.internalAfterKey, nextData.data.length());
+                    sourceRecordLookup.set(op.internalAfterKey, nextData);
+                    sizes?.set(op.internalAfterKey, nextData.data?.length() ?? 0);
                 }
                 if (persistedBatch.shouldFlushTransaction()) {
+                    evalSpan.end();
                     // Transaction is getting big.
                     // Flush, and resume in a new transaction.
-                    const { flushedAny } = await persistedBatch.flush(this.db, this.session, options);
+                    using persistSpan = this.tracer.span('storage', 'persist_flush');
+                    const { flushedAny } = await persistedBatch.flush(this.session, options);
                     didFlush ||= flushedAny;
                     persistedBatch = null;
                     // Computing our current progress is a little tricky here, since
                     // we're stopping in the middle of a batch.
                     // We create a new batch, and push any remaining operations to it.
                     resumeBatch = new OperationBatch();
+                    persistSpan.end();
+                    evalSpan = this.tracer.span('evaluate');
                 }
             }
+            evalSpan.end();
             if (persistedBatch) {
                 transactionSize = persistedBatch.currentSize;
-                const { flushedAny } = await persistedBatch.flush(this.db, this.session, options);
+                using _ = this.tracer.span('storage', 'persist_flush');
+                const { flushedAny } = await persistedBatch.flush(this.session, options);
                 didFlush ||= flushedAny;
             }
         }
         if (didFlush) {
+            using _ = this.tracer.span('storage', 'clear_error');
             await this.clearError();
         }
         return resumeBatch?.hasData() ? resumeBatch : null;
     }
-    saveOperation(batch, operation, current_data, opSeq) {
+    saveOperation(batch, operation, sourceRecord, opSeq) {
         const record = operation.record;
         const beforeId = operation.beforeId;
         const afterId = operation.afterId;
@@ -237,10 +221,10 @@ 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 sourceTableId = mongoTableId(record.sourceTable.id);
         if (this.skipExistingRows) {
             if (record.tag == SaveOperationTag.INSERT) {
-                if (current_data != null) {
+                if (sourceRecord != null) {
                     // Initial replication, and we already have the record.
                     // This may be a different version of the record, but streaming replication
                     // will take care of that.
@@ -253,7 +237,7 @@ export class MongoBucketBatch extends BaseObserver {
             }
         }
         if (record.tag == SaveOperationTag.UPDATE) {
-            const result = current_data;
+            const result = sourceRecord;
             if (result == null) {
                 // Not an error if we re-apply a transaction
                 existing_buckets = [];
@@ -274,14 +258,14 @@ export class MongoBucketBatch extends BaseObserver {
             else {
                 existing_buckets = result.buckets;
                 existing_lookups = result.lookups;
-                if (this.storeCurrentData) {
+                if (this.storeCurrentData && result.data != null) {
                     const data = deserializeBson(result.data.buffer);
                     after = storage.mergeToast(after, data);
                 }
             }
         }
         else if (record.tag == SaveOperationTag.DELETE) {
-            const result = current_data;
+            const result = sourceRecord;
             if (result == null) {
                 // Not an error if we re-apply a transaction
                 existing_buckets = [];
@@ -296,9 +280,9 @@ export class MongoBucketBatch extends BaseObserver {
                 existing_lookups = result.lookups;
             }
         }
-        let afterData;
+        let afterData = null;
         if (afterId != null && !this.storeCurrentData) {
-            afterData = new bson.Binary(bson.serialize({}));
+            afterData = null;
         }
         else if (afterId != null) {
             try {
@@ -382,13 +366,7 @@ export class MongoBucketBatch extends BaseObserver {
                     table: sourceTable,
                     before_buckets: existing_buckets
                 });
-                new_buckets = evaluated.map((e) => {
-                    return {
-                        bucket: e.bucket,
-                        table: e.table,
-                        id: e.id
-                    };
-                });
+                new_buckets = this.sourceRecordStore.mapEvaluatedBuckets(evaluated);
             }
             if (sourceTable.syncParameters) {
                 // Parameters
@@ -410,36 +388,42 @@ export class MongoBucketBatch extends BaseObserver {
                     evaluated: paramEvaluated,
                     existing_lookups
                 });
-                new_lookups = paramEvaluated.map((p) => {
-                    return storage.serializeLookup(p.lookup);
-                });
+                new_lookups = this.sourceRecordStore.mapParameterLookups(paramEvaluated);
             }
         }
         let result = null;
         // 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 };
-            batch.upsertCurrentData(after_key, {
+            batch.upsertCurrentData({
+                sourceTableId,
+                replicaId: afterId,
                 data: afterData,
                 buckets: new_buckets,
                 lookups: new_lookups
             });
             result = {
-                _id: after_key,
+                sourceTableId,
+                replicaId: afterId,
                 data: afterData,
                 buckets: new_buckets,
-                lookups: new_lookups
+                lookups: new_lookups,
+                cacheKey: operation.internalAfterKey
             };
         }
         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(sourceTableId, beforeId, opSeq.next());
         }
         return result;
     }
     async withTransaction(cb) {
+        using lockSpan = this.tracer.span('storage', 'internal_lock');
         await replicationMutex.exclusiveLock(async () => {
+            lockSpan.end();
             await this.session.withTransaction(async () => {
                 try {
                     await cb();
@@ -451,7 +435,9 @@ export class MongoBucketBatch extends BaseObserver {
                     else {
                         this.logger.warn('Transaction error', e);
                     }
-                    await timers.setTimeout(Math.random() * 50);
+                    const delay = Math.random() * 50;
+                    using _ = this.tracer.span('storage', 'retry_delay');
+                    await timers.setTimeout(delay);
                     throw e;
                 }
             }, { maxCommitTimeMS: 10000 });
@@ -505,59 +491,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 +520,134 @@ 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 (newLastCheckpoint != null) {
+                await this.sourceRecordStore.postCommitCleanup(newLastCheckpoint, this.logger);
+            }
+        }
+        return { checkpointBlocked, checkpointCreated };
     }
     /**
      * Switch from processing -> active if relevant.
@@ -597,7 +664,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 +682,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}`);
+            this.logger.info(`Activated new replication stream 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,9 +744,10 @@ 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.commonSourceTables(this.group_id).deleteOne({ _id: mongoTableId(table.id) });
             }
         });
+        await this.cleanupDroppedSourceTables(sourceTables);
         return result;
     }
     async truncate(sourceTables) {
@@ -741,38 +774,31 @@ export class MongoBucketBatch extends BaseObserver {
         let lastBatchCount = BATCH_LIMIT;
         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'])
-                };
-                const cursor = this.db.current_data.find(current_data_filter, {
-                    projection: {
-                        _id: 1,
-                        buckets: 1,
-                        lookups: 1
-                    },
-                    limit: BATCH_LIMIT,
-                    session: session
-                });
-                const batch = await cursor.toArray();
-                const persistedBatch = new PersistedBatch(this.group_id, 0, { logger: this.logger });
+                using evalSpan = this.tracer.span('evaluate');
+                const sourceTableId = mongoTableId(sourceTable.id);
+                const batch = await this.sourceRecordStore.loadTruncateBatch(session, sourceTableId, BATCH_LIMIT);
+                const persistedBatch = this.createPersistedBatch(0);
                 for (let value of batch) {
                     persistedBatch.saveBucketData({
                         op_seq: opSeq,
                         before_buckets: value.buckets,
                         evaluated: [],
                         table: sourceTable,
-                        sourceKey: value._id.k
+                        sourceKey: value.replicaId
                     });
                     persistedBatch.saveParameterData({
                         op_seq: opSeq,
                         existing_lookups: value.lookups,
                         evaluated: [],
                         sourceTable: sourceTable,
-                        sourceKey: value._id.k
+                        sourceKey: value.replicaId
                     });
-                    persistedBatch.deleteCurrentData(value._id);
+                    // Since this is not from streaming replication, we can do a hard delete
+                    persistedBatch.hardDeleteCurrentData(sourceTableId, value.replicaId);
                 }
-                await persistedBatch.flush(this.db, session);
+                evalSpan.end();
+                using _ = this.tracer.span('storage', 'persist_flush');
+                await persistedBatch.flush(session);
                 lastBatchCount = batch.length;
                 last_op = opSeq.last();
             });
@@ -788,7 +814,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.commonSourceTables(this.group_id).updateOne({ _id: mongoTableId(table.id) }, {
                 $set: {
                     snapshot_status: {
                         last_key: snapshotStatus.lastKey == null ? null : new bson.Binary(snapshotStatus.lastKey),
@@ -800,11 +826,33 @@ 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 } }, {
+            await this.db.commonSourceTables(this.group_id).updateMany({ _id: { $in: ids } }, {
                 $set: {
                     snapshot_done: true
                 },
@@ -812,14 +860,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 });
             }
@@ -852,7 +901,4 @@ export class MongoBucketBatch extends BaseObserver {
         return this.sync_rules.eventDescriptors.filter((evt) => [...evt.getSourceTables()].some((sourceTable) => sourceTable.matches(table)));
     }
 }
-export function currentBucketKey(b) {
-    return `${b.bucket}/${b.table}/${b.id}`;
-}
 //# sourceMappingURL=MongoBucketBatch.js.map