@powersync/service-core 0.0.0-dev-20240718134716 → 0.0.0-dev-20240918082156

package/src/storage/mongo/MongoCompactor.ts

@@ -0,0 +1,372 @@
+import { logger } from '@powersync/lib-services-framework';
+import { AnyBulkWriteOperation, MaxKey, MinKey } from 'mongodb';
+import { addChecksums } from '../../util/utils.js';
+import { PowerSyncMongo } from './db.js';
+import { BucketDataDocument, BucketDataKey } from './models.js';
+import { CompactOptions } from '../BucketStorage.js';
+import { cacheKey } from './OperationBatch.js';
+
+interface CurrentBucketState {
+  /** Bucket name */
+  bucket: string;
+  /**
+   * Rows seen in the bucket, with the last op_id of each.
+   */
+  seen: Map<string, bigint>;
+  /**
+   * Estimated memory usage of the seen Map.
+   */
+  trackingSize: number;
+
+  /**
+   * Last (lowest) seen op_id that is not a PUT.
+   */
+  lastNotPut: bigint | null;
+
+  /**
+   * Number of REMOVE/MOVE operations seen since lastNotPut.
+   */
+  opsSincePut: number;
+}
+
+/**
+ * Additional options, primarily for testing.
+ */
+export interface MongoCompactOptions extends CompactOptions {
+  /** Minimum of 2 */
+  clearBatchLimit?: number;
+  /** Minimum of 1 */
+  moveBatchLimit?: number;
+  /** Minimum of 1 */
+  moveBatchQueryLimit?: number;
+}
+
+const DEFAULT_CLEAR_BATCH_LIMIT = 5000;
+const DEFAULT_MOVE_BATCH_LIMIT = 2000;
+const DEFAULT_MOVE_BATCH_QUERY_LIMIT = 10_000;
+
+/** This default is primarily for tests. */
+const DEFAULT_MEMORY_LIMIT_MB = 64;
+
+export class MongoCompactor {
+  private updates: AnyBulkWriteOperation<BucketDataDocument>[] = [];
+
+  private idLimitBytes: number;
+  private moveBatchLimit: number;
+  private moveBatchQueryLimit: number;
+  private clearBatchLimit: number;
+  private maxOpId: bigint | undefined;
+  private buckets: string[] | undefined;
+
+  constructor(private db: PowerSyncMongo, private group_id: number, options?: MongoCompactOptions) {
+    this.idLimitBytes = (options?.memoryLimitMB ?? DEFAULT_MEMORY_LIMIT_MB) * 1024 * 1024;
+    this.moveBatchLimit = options?.moveBatchLimit ?? DEFAULT_MOVE_BATCH_LIMIT;
+    this.moveBatchQueryLimit = options?.moveBatchQueryLimit ?? DEFAULT_MOVE_BATCH_QUERY_LIMIT;
+    this.clearBatchLimit = options?.clearBatchLimit ?? DEFAULT_CLEAR_BATCH_LIMIT;
+    this.maxOpId = options?.maxOpId;
+    this.buckets = options?.compactBuckets;
+  }
+
+  /**
+   * Compact buckets by converting operations into MOVE and/or CLEAR operations.
+   *
+   * See /docs/compacting-operations.md for details.
+   */
+  async compact() {
+    if (this.buckets) {
+      for (let bucket of this.buckets) {
+        // We can make this more efficient later on by iterating
+        // through the buckets in a single query.
+        // That makes batching more tricky, so we leave for later.
+        await this.compactInternal(bucket);
+      }
+    } else {
+      await this.compactInternal(undefined);
+    }
+  }
+
+  async compactInternal(bucket: string | undefined) {
+    const idLimitBytes = this.idLimitBytes;
+
+    let currentState: CurrentBucketState | null = null;
+
+    // Constant lower bound
+    const lowerBound: BucketDataKey = {
+      g: this.group_id,
+      b: bucket ?? (new MinKey() as any),
+      o: new MinKey() as any
+    };
+
+    // Upper bound is adjusted for each batch
+    let upperBound: BucketDataKey = {
+      g: this.group_id,
+      b: bucket ?? (new MaxKey() as any),
+      o: new MaxKey() as any
+    };
+
+    while (true) {
+      // Query one batch at a time, to avoid cursor timeouts
+      const batch = await this.db.bucket_data
+        .find(
+          {
+            _id: {
+              $gte: lowerBound,
+              $lt: upperBound
+            }
+          },
+          {
+            projection: {
+              _id: 1,
+              op: 1,
+              table: 1,
+              row_id: 1,
+              source_table: 1,
+              source_key: 1
+            },
+            limit: this.moveBatchQueryLimit,
+            sort: { _id: -1 },
+            singleBatch: true
+          }
+        )
+        .toArray();
+
+      if (batch.length == 0) {
+        // We've reached the end
+        break;
+      }
+
+      // Set upperBound for the next batch
+      upperBound = batch[batch.length - 1]._id;
+
+      for (let doc of batch) {
+        if (currentState == null || doc._id.b != currentState.bucket) {
+          if (currentState != null && currentState.lastNotPut != null && currentState.opsSincePut >= 1) {
+            // Important to flush before clearBucket()
+            await this.flush();
+            logger.info(
+              `Inserting CLEAR at ${this.group_id}:${currentState.bucket}:${currentState.lastNotPut} to remove ${currentState.opsSincePut} operations`
+            );
+
+            const bucket = currentState.bucket;
+            const clearOp = currentState.lastNotPut;
+            // Free memory before clearing bucket
+            currentState = null;
+            await this.clearBucket(bucket, clearOp);
+          }
+          currentState = {
+            bucket: doc._id.b,
+            seen: new Map(),
+            trackingSize: 0,
+            lastNotPut: null,
+            opsSincePut: 0
+          };
+        }
+
+        if (this.maxOpId != null && doc._id.o > this.maxOpId) {
+          continue;
+        }
+
+        let isPersistentPut = doc.op == 'PUT';
+
+        if (doc.op == 'REMOVE' || doc.op == 'PUT') {
+          const key = `${doc.table}/${doc.row_id}/${cacheKey(doc.source_table!, doc.source_key!)}`;
+          const targetOp = currentState.seen.get(key);
+          if (targetOp) {
+            // Will convert to MOVE, so don't count as PUT
+            isPersistentPut = false;
+
+            this.updates.push({
+              updateOne: {
+                filter: {
+                  _id: doc._id
+                },
+                update: {
+                  $set: {
+                    op: 'MOVE',
+                    target_op: targetOp
+                  },
+                  $unset: {
+                    source_table: 1,
+                    source_key: 1,
+                    table: 1,
+                    row_id: 1,
+                    data: 1
+                  }
+                }
+              }
+            });
+          } else {
+            if (currentState.trackingSize >= idLimitBytes) {
+              // Reached memory limit.
+              // Keep the highest seen values in this case.
+            } else {
+              // flatstr reduces the memory usage by flattening the string
+              currentState.seen.set(flatstr(key), doc._id.o);
+              // length + 16 for the string
+              // 24 for the bigint
+              // 50 for map overhead
+              // 50 for additional overhead
+              currentState.trackingSize += key.length + 140;
+            }
+          }
+        }
+
+        if (isPersistentPut) {
+          currentState.lastNotPut = null;
+          currentState.opsSincePut = 0;
+        } else if (doc.op != 'CLEAR') {
+          if (currentState.lastNotPut == null) {
+            currentState.lastNotPut = doc._id.o;
+          }
+          currentState.opsSincePut += 1;
+        }
+
+        if (this.updates.length >= this.moveBatchLimit) {
+          await this.flush();
+        }
+      }
+    }
+
+    await this.flush();
+    currentState?.seen.clear();
+    if (currentState?.lastNotPut != null && currentState?.opsSincePut > 1) {
+      logger.info(
+        `Inserting CLEAR at ${this.group_id}:${currentState.bucket}:${currentState.lastNotPut} to remove ${currentState.opsSincePut} operations`
+      );
+      const bucket = currentState.bucket;
+      const clearOp = currentState.lastNotPut;
+      // Free memory before clearing bucket
+      currentState = null;
+      await this.clearBucket(bucket, clearOp);
+    }
+  }
+
+  private async flush() {
+    if (this.updates.length > 0) {
+      logger.info(`Compacting ${this.updates.length} ops`);
+      await this.db.bucket_data.bulkWrite(this.updates, {
+        // Order is not important.
+        // Since checksums are not affected, these operations can happen in any order,
+        // and it's fine if the operations are partially applied.
+        // Each individual operation is atomic.
+        ordered: false
+      });
+      this.updates = [];
+    }
+  }
+
+  /**
+   * Perform a CLEAR compact for a bucket.
+   *
+   * @param bucket bucket name
+   * @param op op_id of the last non-PUT operation, which will be converted to CLEAR.
+   */
+  private async clearBucket(bucket: string, op: bigint) {
+    const opFilter = {
+      _id: {
+        $gte: {
+          g: this.group_id,
+          b: bucket,
+          o: new MinKey() as any
+        },
+        $lte: {
+          g: this.group_id,
+          b: bucket,
+          o: op
+        }
+      }
+    };
+
+    const session = this.db.client.startSession();
+    try {
+      let done = false;
+      while (!done) {
+        // Do the CLEAR operation in batches, with each batch a separate transaction.
+        // The state after each batch is fully consistent.
+        // We need a transaction per batch to make sure checksums stay consistent.
+        await session.withTransaction(
+          async () => {
+            const query = this.db.bucket_data.find(opFilter, {
+              session,
+              sort: { _id: 1 },
+              projection: {
+                _id: 1,
+                op: 1,
+                checksum: 1,
+                target_op: 1
+              },
+              limit: this.clearBatchLimit
+            });
+            let checksum = 0;
+            let lastOpId: BucketDataKey | null = null;
+            let targetOp: bigint | null = null;
+            let gotAnOp = false;
+            for await (let op of query.stream()) {
+              if (op.op == 'MOVE' || op.op == 'REMOVE' || op.op == 'CLEAR') {
+                checksum = addChecksums(checksum, op.checksum);
+                lastOpId = op._id;
+                if (op.op != 'CLEAR') {
+                  gotAnOp = true;
+                }
+                if (op.target_op != null) {
+                  if (targetOp == null || op.target_op > targetOp) {
+                    targetOp = op.target_op;
+                  }
+                }
+              } else {
+                throw new Error(`Unexpected ${op.op} operation at ${op._id.g}:${op._id.b}:${op._id.o}`);
+              }
+            }
+            if (!gotAnOp) {
+              done = true;
+              return;
+            }
+
+            logger.info(`Flushing CLEAR at ${lastOpId?.o}`);
+            await this.db.bucket_data.deleteMany(
+              {
+                _id: {
+                  $gte: {
+                    g: this.group_id,
+                    b: bucket,
+                    o: new MinKey() as any
+                  },
+                  $lte: lastOpId!
+                }
+              },
+              { session }
+            );
+
+            await this.db.bucket_data.insertOne(
+              {
+                _id: lastOpId!,
+                op: 'CLEAR',
+                checksum: checksum,
+                data: null,
+                target_op: targetOp
+              },
+              { session }
+            );
+          },
+          {
+            writeConcern: { w: 'majority' },
+            readConcern: { level: 'snapshot' }
+          }
+        );
+      }
+    } finally {
+      await session.endSession();
+    }
+  }
+}
+
+/**
+ * Flattens string to reduce memory usage (around 320 bytes -> 120 bytes),
+ * at the cost of some upfront CPU usage.
+ *
+ * From: https://github.com/davidmarkclements/flatstr/issues/8
+ */
+function flatstr(s: string) {
+  s.match(/\n/g);
+  return s;
+}

package/src/storage/mongo/MongoPersistedSyncRulesContent.ts

@@ -1,7 +1,7 @@
 import { SqlSyncRules } from '@powersync/service-sync-rules';
 import * as mongo from 'mongodb';
 
-import { PersistedSyncRulesContent } from '../BucketStorage.js';
+import { ParseSyncRulesOptions, PersistedSyncRulesContent } from '../BucketStorage.js';
 import { MongoPersistedSyncRules } from './MongoPersistedSyncRules.js';
 import { MongoSyncRulesLock } from './MongoSyncRulesLock.js';
 import { PowerSyncMongo } from './db.js';
@@ -30,10 +30,10 @@ export class MongoPersistedSyncRulesContent implements PersistedSyncRulesContent
     this.last_keepalive_ts = doc.last_keepalive_ts;
   }
 
-  parsed() {
+  parsed(options: ParseSyncRulesOptions) {
     return new MongoPersistedSyncRules(
       this.id,
-      SqlSyncRules.fromYaml(this.sync_rules_content),
+      SqlSyncRules.fromYaml(this.sync_rules_content, options),
       this.last_checkpoint_lsn,
       this.slot_name
     );

package/src/storage/mongo/MongoStorageProvider.ts

@@ -0,0 +1,31 @@
+import * as db from '../../db/db-index.js';
+import { MongoBucketStorage } from '../MongoBucketStorage.js';
+import { BucketStorageProvider, ActiveStorage, GetStorageOptions } from '../StorageProvider.js';
+import { PowerSyncMongo } from './db.js';
+import { logger } from '@powersync/lib-services-framework';
+
+export class MongoStorageProvider implements BucketStorageProvider {
+  get type() {
+    return 'mongodb';
+  }
+
+  async getStorage(options: GetStorageOptions): Promise<ActiveStorage> {
+    const { resolvedConfig } = options;
+
+    const client = db.mongo.createMongoClient(resolvedConfig.storage);
+
+    const database = new PowerSyncMongo(client, { database: resolvedConfig.storage.database });
+
+    return {
+      storage: new MongoBucketStorage(database, {
+        // TODO currently need the entire resolved config due to this
+        slot_name_prefix: resolvedConfig.slot_name_prefix
+      }),
+      shutDown: () => client.close(),
+      tearDown: () => {
+        logger.info(`Tearing down storage: ${database.db.namespace}...`);
+        return database.db.dropDatabase();
+      }
+    } satisfies ActiveStorage;
+  }
+}

package/src/storage/mongo/MongoSyncBucketStorage.ts

@@ -3,26 +3,33 @@ import * as bson from 'bson';
 import * as mongo from 'mongodb';
 
 import * as db from '../../db/db-index.js';
-import * as replication from '../../replication/WalStream.js';
 import * as util from '../../util/util-index.js';
 import {
   BucketDataBatchOptions,
   BucketStorageBatch,
+  CompactOptions,
   DEFAULT_DOCUMENT_BATCH_LIMIT,
   DEFAULT_DOCUMENT_CHUNK_LIMIT_BYTES,
   FlushedResult,
+  ParseSyncRulesOptions,
+  PersistedSyncRules,
+  PersistedSyncRulesContent,
   ResolveTableOptions,
   ResolveTableResult,
+  StartBatchOptions,
+  SyncBucketDataBatch,
   SyncRulesBucketStorage,
-  SyncRuleStatus
+  SyncRuleStatus,
+  TerminateOptions
 } from '../BucketStorage.js';
+import { ChecksumCache, FetchPartialBucketChecksum } from '../ChecksumCache.js';
 import { MongoBucketStorage } from '../MongoBucketStorage.js';
 import { SourceTable } from '../SourceTable.js';
 import { PowerSyncMongo } from './db.js';
 import { BucketDataDocument, BucketDataKey, SourceKey, SyncRuleState } from './models.js';
 import { MongoBucketBatch } from './MongoBucketBatch.js';
-import { BSON_DESERIALIZE_OPTIONS, idPrefixFilter, readSingleBatch, serializeLookup } from './util.js';
-import { ChecksumCache, FetchPartialBucketChecksum } from '../ChecksumCache.js';
+import { MongoCompactor } from './MongoCompactor.js';
+import { BSON_DESERIALIZE_OPTIONS, idPrefixFilter, mapOpEntry, readSingleBatch, serializeLookup } from './util.js';
 
 export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
   private readonly db: PowerSyncMongo;
@@ -32,29 +39,38 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
     }
   });
 
+  private parsedSyncRulesCache: SqlSyncRules | undefined;
+
   constructor(
     public readonly factory: MongoBucketStorage,
     public readonly group_id: number,
-    public readonly sync_rules: SqlSyncRules,
+    private readonly sync_rules: PersistedSyncRulesContent,
     public readonly slot_name: string
   ) {
     this.db = factory.db;
   }
 
+  getParsedSyncRules(options: ParseSyncRulesOptions): SqlSyncRules {
+    this.parsedSyncRulesCache ??= this.sync_rules.parsed(options).sync_rules;
+    return this.parsedSyncRulesCache;
+  }
+
   async getCheckpoint() {
     const doc = await this.db.sync_rules.findOne(
       { _id: this.group_id },
       {
-        projection: { last_checkpoint: 1, last_checkpoint_lsn: 1 }
+        projection: { last_checkpoint: 1 }
       }
     );
     return {
-      checkpoint: util.timestampToOpId(doc?.last_checkpoint ?? 0n),
-      lsn: doc?.last_checkpoint_lsn ?? replication.ZERO_LSN
+      checkpoint: util.timestampToOpId(doc?.last_checkpoint ?? 0n)
     };
   }
 
-  async startBatch(options: {}, callback: (batch: BucketStorageBatch) => Promise<void>): Promise<FlushedResult | null> {
+  async startBatch(
+    options: StartBatchOptions,
+    callback: (batch: BucketStorageBatch) => Promise<void>
+  ): Promise<FlushedResult | null> {
     const doc = await this.db.sync_rules.findOne(
       {
         _id: this.group_id
@@ -65,11 +81,11 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
 
     const batch = new MongoBucketBatch(
       this.db,
-      this.sync_rules,
+      this.sync_rules.parsed(options).sync_rules,
       this.group_id,
       this.slot_name,
       checkpoint_lsn,
-      doc?.no_checkpoint_before ?? null
+      doc?.no_checkpoint_before ?? options.zeroLSN
     );
     try {
       await callback(batch);
@@ -87,11 +103,15 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
   }
 
   async resolveTable(options: ResolveTableOptions): Promise<ResolveTableResult> {
-    const { group_id, connection_id, connection_tag, relation } = options;
+    const { group_id, connection_id, connection_tag, entity_descriptor } = options;
 
-    const { schema, name: table, relationId, replicationColumns } = relation;
+    const { schema, name: table, objectId, replicationColumns } = entity_descriptor;
 
-    const columns = replicationColumns.map((column) => ({ name: column.name, type_oid: column.typeOid }));
+    const columns = replicationColumns.map((column) => ({
+      name: column.name,
+      type: column.type,
+      type_oid: column.typeId
+    }));
     let result: ResolveTableResult | null = null;
     await this.db.client.withSession(async (session) => {
       const col = this.db.source_tables;
@@ -99,7 +119,7 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
         {
           group_id: group_id,
           connection_id: connection_id,
-          relation_id: relationId,
+          relation_id: objectId,
           schema_name: schema,
           table_name: table,
           replica_id_columns2: columns
@@ -111,7 +131,7 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
           _id: new bson.ObjectId(),
           group_id: group_id,
           connection_id: connection_id,
-          relation_id: relationId,
+          relation_id: objectId,
           schema_name: schema,
           table_name: table,
           replica_id_columns: null,
@@ -124,7 +144,7 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
       const sourceTable = new SourceTable(
         doc._id,
         connection_tag,
-        relationId,
+        objectId,
         schema,
         table,
         replicationColumns,
@@ -139,7 +159,7 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
             group_id: group_id,
             connection_id: connection_id,
             _id: { $ne: doc._id },
-            $or: [{ relation_id: relationId }, { schema_name: schema, table_name: table }]
+            $or: [{ relation_id: objectId }, { schema_name: schema, table_name: table }]
           },
           { session }
         )
@@ -154,7 +174,7 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
               doc.relation_id ?? 0,
               doc.schema_name,
               doc.table_name,
-              doc.replica_id_columns2?.map((c) => ({ name: c.name, typeOid: c.type_oid })) ?? [],
+              doc.replica_id_columns2?.map((c) => ({ name: c.name, typeOid: c.type_oid, type: c.type })) ?? [],
               doc.snapshot_done ?? true
             )
         )
@@ -201,7 +221,7 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
     checkpoint: util.OpId,
     dataBuckets: Map<string, string>,
     options?: BucketDataBatchOptions
-  ): AsyncIterable<util.SyncBucketData> {
+  ): AsyncIterable<SyncBucketDataBatch> {
     if (dataBuckets.size == 0) {
       return;
     }
@@ -267,6 +287,7 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
 
     let batchSize = 0;
     let currentBatch: util.SyncBucketData | null = null;
+    let targetOp: bigint | null = null;
 
     // Ordered by _id, meaning buckets are grouped together
     for (let rawData of data) {
@@ -284,7 +305,8 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
           start = currentBatch.after;
           currentBatch = null;
           batchSize = 0;
-          yield yieldBatch;
+          yield { batch: yieldBatch, targetOp: targetOp };
+          targetOp = null;
         }
 
         start ??= dataBuckets.get(bucket);
@@ -298,17 +320,18 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
           data: [],
           next_after: start
         };
+        targetOp = null;
+      }
+
+      const entry = mapOpEntry(row);
+
+      if (row.target_op != null) {
+        // MOVE, CLEAR
+        if (targetOp == null || row.target_op > targetOp) {
+          targetOp = row.target_op;
+        }
       }
 
-      const entry: util.OplogEntry = {
-        op_id: util.timestampToOpId(row._id.o),
-        op: row.op,
-        object_type: row.table,
-        object_id: row.row_id,
-        checksum: Number(row.checksum),
-        subkey: `${row.source_table}/${row.source_key.toHexString()}`,
-        data: row.data
-      };
       currentBatch.data.push(entry);
       currentBatch.next_after = entry.op_id;
 
@@ -318,7 +341,8 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
     if (currentBatch != null) {
       const yieldBatch = currentBatch;
       currentBatch = null;
-      yield yieldBatch;
+      yield { batch: yieldBatch, targetOp: targetOp };
+      targetOp = null;
     }
   }
 
@@ -379,9 +403,11 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
     );
   }
 
-  async terminate() {
-    await this.clear();
-
+  async terminate(options?: TerminateOptions) {
+    // Default is to clear the storage except when explicitly requested not to.
+    if (!options || options?.clearStorage) {
+      await this.clear();
+    }
     await this.db.sync_rules.updateOne(
       {
         _id: this.group_id
@@ -530,4 +556,8 @@ export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
       }
     );
   }
+
+  async compact(options?: CompactOptions) {
+    return new MongoCompactor(this.db, this.group_id, options).compact();
+  }
 }

package/src/storage/mongo/MongoSyncRulesLock.ts

@@ -30,7 +30,7 @@ export class MongoSyncRulesLock implements ReplicationLock {
     );
 
     if (doc == null) {
-      throw new Error(`Replication slot ${sync_rules.slot_name} is locked by another process`);
+      throw new Error(`Sync rules: ${sync_rules.id} have been locked by another process for replication.`);
     }
     return new MongoSyncRulesLock(db, sync_rules.id, lockId);
   }