@powersync/service-core 0.0.0-dev-20240620165206

package/src/storage/mongo/OperationBatch.ts

@@ -0,0 +1,115 @@
+import * as bson from 'bson';
+import { ToastableSqliteRow } from '@powersync/service-sync-rules';
+
+import * as util from '../../util/util-index.js';
+import { SaveOptions } from '../BucketStorage.js';
+
+/**
+ * Maximum number of operations in a batch.
+ */
+const MAX_BATCH_COUNT = 2000;
+
+/**
+ * Maximum size of operations in the batch (estimated).
+ */
+const MAX_RECORD_BATCH_SIZE = 5_000_000;
+
+/**
+ * Maximum size of size of current_data documents we lookup at a time.
+ */
+const MAX_CURRENT_DATA_BATCH_SIZE = 16_000_000;
+
+/**
+ * Batch of input operations.
+ *
+ * We accumulate operations up to MAX_RECORD_BATCH_SIZE,
+ * then further split into sub-batches if MAX_CURRENT_DATA_BATCH_SIZE is exceeded.
+ */
+export class OperationBatch {
+  batch: RecordOperation[] = [];
+  currentSize: number = 0;
+
+  get length() {
+    return this.batch.length;
+  }
+
+  push(op: RecordOperation) {
+    this.batch.push(op);
+    this.currentSize += op.estimatedSize;
+  }
+
+  shouldFlush() {
+    return this.batch.length >= MAX_BATCH_COUNT || this.currentSize > MAX_RECORD_BATCH_SIZE;
+  }
+
+  *batched(sizes: Map<string, number>): Generator<RecordOperation[]> {
+    let currentBatch: RecordOperation[] = [];
+    let currentBatchSize = 0;
+    for (let op of this.batch) {
+      const key = op.internalBeforeKey;
+      const size = sizes.get(key) ?? 0;
+      if (currentBatchSize + size > MAX_CURRENT_DATA_BATCH_SIZE && currentBatch.length > 0) {
+        yield currentBatch;
+        currentBatch = [];
+        currentBatchSize = 0;
+      }
+      currentBatchSize += size;
+      currentBatch.push(op);
+    }
+    if (currentBatch.length > 0) {
+      yield currentBatch;
+    }
+  }
+}
+
+export class RecordOperation {
+  public readonly afterId: bson.UUID | null;
+  public readonly beforeId: bson.UUID;
+  public readonly internalBeforeKey: string;
+  public readonly internalAfterKey: string | null;
+  public readonly estimatedSize: number;
+
+  constructor(public readonly record: SaveOptions) {
+    const after = record.after;
+    const afterId = after ? util.getUuidReplicaIdentityBson(after, record.sourceTable.replicaIdColumns!) : null;
+    const beforeId = record.before
+      ? util.getUuidReplicaIdentityBson(record.before, record.sourceTable.replicaIdColumns!)
+      : afterId!;
+    this.afterId = afterId;
+    this.beforeId = beforeId;
+    this.internalBeforeKey = cacheKey(record.sourceTable.id, beforeId);
+    this.internalAfterKey = afterId ? cacheKey(record.sourceTable.id, afterId) : null;
+
+    this.estimatedSize = estimateRowSize(record.before) + estimateRowSize(record.after);
+  }
+}
+
+export function cacheKey(table: bson.ObjectId, id: bson.UUID) {
+  return `${table.toHexString()}.${id.toHexString()}`;
+}
+
+/**
+ * Estimate in-memory size of row.
+ */
+function estimateRowSize(record: ToastableSqliteRow | undefined) {
+  if (record == null) {
+    return 12;
+  }
+  let size = 0;
+  for (let [key, value] of Object.entries(record)) {
+    size += 12 + key.length;
+    // number | string | null | bigint | Uint8Array
+    if (value == null) {
+      size += 4;
+    } else if (typeof value == 'number') {
+      size += 8;
+    } else if (typeof value == 'bigint') {
+      size += 8;
+    } else if (typeof value == 'string') {
+      size += value.length;
+    } else if (value instanceof Uint8Array) {
+      size += value.byteLength;
+    }
+  }
+  return size;
+}

package/src/storage/mongo/PersistedBatch.ts

@@ -0,0 +1,268 @@
+import { JSONBig } from '@powersync/service-jsonbig';
+import { EvaluatedParameters, EvaluatedRow } from '@powersync/service-sync-rules';
+import * as bson from 'bson';
+import * as mongo from 'mongodb';
+
+import * as util from '../../util/util-index.js';
+import { SourceTable } from '../SourceTable.js';
+import { currentBucketKey } from './MongoBucketBatch.js';
+import { MongoIdSequence } from './MongoIdSequence.js';
+import { PowerSyncMongo } from './db.js';
+import {
+  BucketDataDocument,
+  BucketParameterDocument,
+  CurrentBucket,
+  CurrentDataDocument,
+  SourceKey
+} from './models.js';
+import { serializeLookup } from './util.js';
+import { logger } from '@powersync/lib-services-framework';
+
+/**
+ * Maximum size of operations we write in a single transaction.
+ *
+ * It's tricky to find the exact limit, but from experience, over 100MB
+ * can cause an error:
+ * > transaction is too large and will not fit in the storage engine cache
+ *
+ * Additionally, unbounded size here can balloon our memory usage in some edge
+ * cases.
+ *
+ * When we reach this threshold, we commit the transaction and start a new one.
+ */
+const MAX_TRANSACTION_BATCH_SIZE = 30_000_000;
+
+/**
+ * Keeps track of bulkwrite operations within a transaction.
+ *
+ * There may be multiple of these batches per transaction, but it may not span
+ * multiple transactions.
+ */
+export class PersistedBatch {
+  bucketData: mongo.AnyBulkWriteOperation<BucketDataDocument>[] = [];
+  bucketParameters: mongo.AnyBulkWriteOperation<BucketParameterDocument>[] = [];
+  currentData: mongo.AnyBulkWriteOperation<CurrentDataDocument>[] = [];
+
+  /**
+   * For debug logging only.
+   */
+  debugLastOpId: bigint | null = null;
+
+  /**
+   * Very rough estimate of transaction size.
+   */
+  currentSize = 0;
+
+  constructor(private group_id: number, writtenSize: number) {
+    this.currentSize = writtenSize;
+  }
+
+  saveBucketData(options: {
+    op_seq: MongoIdSequence;
+    sourceKey: bson.UUID;
+    table: SourceTable;
+    evaluated: EvaluatedRow[];
+    before_buckets: CurrentBucket[];
+  }) {
+    const remaining_buckets = new Map<string, CurrentBucket>();
+    for (let b of options.before_buckets) {
+      const key = currentBucketKey(b);
+      remaining_buckets.set(key, b);
+    }
+
+    const dchecksum = util.hashDelete(`${options.table.id}/${options.sourceKey}`);
+
+    for (let k of options.evaluated) {
+      const key = currentBucketKey(k);
+      remaining_buckets.delete(key);
+
+      // INSERT
+      const recordData = JSONBig.stringify(k.data);
+      const checksum = util.hashData(k.table, k.id, recordData);
+      this.currentSize += recordData.length + 200;
+
+      const op_id = options.op_seq.next();
+      this.debugLastOpId = op_id;
+
+      this.bucketData.push({
+        insertOne: {
+          document: {
+            _id: {
+              g: this.group_id,
+              b: k.bucket,
+              o: op_id
+            },
+            op: 'PUT',
+            source_table: options.table.id,
+            source_key: options.sourceKey,
+            table: k.table,
+            row_id: k.id,
+            checksum: checksum,
+            data: recordData
+          }
+        }
+      });
+    }
+
+    for (let bd of remaining_buckets.values()) {
+      // REMOVE
+
+      const op_id = options.op_seq.next();
+      this.debugLastOpId = op_id;
+
+      this.bucketData.push({
+        insertOne: {
+          document: {
+            _id: {
+              g: this.group_id,
+              b: bd.bucket,
+              o: op_id
+            },
+            op: 'REMOVE',
+            source_table: options.table.id,
+            source_key: options.sourceKey,
+            table: bd.table,
+            row_id: bd.id,
+            checksum: dchecksum,
+            data: null
+          }
+        }
+      });
+      this.currentSize += 200;
+    }
+  }
+
+  saveParameterData(data: {
+    op_seq: MongoIdSequence;
+    sourceKey: bson.UUID;
+    sourceTable: SourceTable;
+    evaluated: EvaluatedParameters[];
+    existing_lookups: bson.Binary[];
+  }) {
+    // This is similar to saving bucket data.
+    // A key difference is that we don't need to keep the history intact.
+    // We do need to keep track of recent history though - enough that we can get consistent data for any specific checkpoint.
+    // Instead of storing per bucket id, we store per "lookup".
+    // A key difference is that we don't need to store or keep track of anything per-bucket - the entire record is
+    // either persisted or removed.
+    // We also don't need to keep history intact.
+    const { sourceTable, sourceKey, evaluated } = data;
+
+    const remaining_lookups = new Map<string, bson.Binary>();
+    for (let l of data.existing_lookups) {
+      remaining_lookups.set(l.toString('base64'), l);
+    }
+
+    // 1. Insert new entries
+    for (let result of evaluated) {
+      const binLookup = serializeLookup(result.lookup);
+      const hex = binLookup.toString('base64');
+      remaining_lookups.delete(hex);
+
+      const op_id = data.op_seq.next();
+      this.debugLastOpId = op_id;
+      this.bucketParameters.push({
+        insertOne: {
+          document: {
+            _id: op_id,
+            key: {
+              g: this.group_id,
+              t: sourceTable.id,
+              k: sourceKey
+            },
+            lookup: binLookup,
+            bucket_parameters: result.bucket_parameters
+          }
+        }
+      });
+
+      this.currentSize += 200;
+    }
+
+    // 2. "REMOVE" entries for any lookup not touched.
+    for (let lookup of remaining_lookups.values()) {
+      const op_id = data.op_seq.next();
+      this.debugLastOpId = op_id;
+      this.bucketParameters.push({
+        insertOne: {
+          document: {
+            _id: op_id,
+            key: {
+              g: this.group_id,
+              t: sourceTable.id,
+              k: sourceKey
+            },
+            lookup: lookup,
+            bucket_parameters: []
+          }
+        }
+      });
+
+      this.currentSize += 200;
+    }
+  }
+
+  deleteCurrentData(id: SourceKey) {
+    const op: mongo.AnyBulkWriteOperation<CurrentDataDocument> = {
+      deleteOne: {
+        filter: { _id: id }
+      }
+    };
+    this.currentData.push(op);
+    this.currentSize += 50;
+  }
+
+  upsertCurrentData(id: SourceKey, values: Partial<CurrentDataDocument>) {
+    const op: mongo.AnyBulkWriteOperation<CurrentDataDocument> = {
+      updateOne: {
+        filter: { _id: id },
+        update: {
+          $set: values
+        },
+        upsert: true
+      }
+    };
+    this.currentData.push(op);
+    this.currentSize += (values.data?.length() ?? 0) + 100;
+  }
+
+  shouldFlushTransaction() {
+    return this.currentSize >= MAX_TRANSACTION_BATCH_SIZE;
+  }
+
+  async flush(db: PowerSyncMongo, session: mongo.ClientSession) {
+    if (this.bucketData.length > 0) {
+      await db.bucket_data.bulkWrite(this.bucketData, {
+        session,
+        // inserts only - order doesn't matter
+        ordered: false
+      });
+    }
+    if (this.bucketParameters.length > 0) {
+      await db.bucket_parameters.bulkWrite(this.bucketParameters, {
+        session,
+        // inserts only - order doesn't matter
+        ordered: false
+      });
+    }
+    if (this.currentData.length > 0) {
+      await db.current_data.bulkWrite(this.currentData, {
+        session,
+        // may update and delete data within the same batch - order matters
+        ordered: true
+      });
+    }
+
+    logger.info(
+      `powersync_${this.group_id} Flushed ${this.bucketData.length} + ${this.bucketParameters.length} + ${
+        this.currentData.length
+      } updates, ${Math.round(this.currentSize / 1024)}kb. Last op_id: ${this.debugLastOpId}`
+    );
+
+    this.bucketData = [];
+    this.bucketParameters = [];
+    this.currentData = [];
+    this.currentSize = 0;
+    this.debugLastOpId = null;
+  }
+}

package/src/storage/mongo/db.ts

@@ -0,0 +1,73 @@
+import * as mongo from 'mongodb';
+
+import * as db from '../../db/db-index.js';
+import * as locks from '../../locks/locks-index.js';
+import {
+  BucketDataDocument,
+  BucketParameterDocument,
+  CurrentDataDocument,
+  IdSequenceDocument,
+  InstanceDocument,
+  SourceTableDocument,
+  SyncRuleDocument,
+  WriteCheckpointDocument
+} from './models.js';
+import { BSON_DESERIALIZE_OPTIONS } from './util.js';
+import { configFile } from '@powersync/service-types';
+
+export interface PowerSyncMongoOptions {
+  /**
+   * Optional - uses the database from the MongoClient connection URI if not specified.
+   */
+  database?: string;
+}
+
+export function createPowerSyncMongo(config: configFile.PowerSyncConfig['storage']) {
+  return new PowerSyncMongo(db.mongo.createMongoClient(config), { database: config.database });
+}
+
+export class PowerSyncMongo {
+  readonly current_data: mongo.Collection<CurrentDataDocument>;
+  readonly bucket_data: mongo.Collection<BucketDataDocument>;
+  readonly bucket_parameters: mongo.Collection<BucketParameterDocument>;
+  readonly op_id_sequence: mongo.Collection<IdSequenceDocument>;
+  readonly sync_rules: mongo.Collection<SyncRuleDocument>;
+  readonly source_tables: mongo.Collection<SourceTableDocument>;
+  readonly write_checkpoints: mongo.Collection<WriteCheckpointDocument>;
+  readonly instance: mongo.Collection<InstanceDocument>;
+  readonly locks: mongo.Collection<locks.Lock>;
+
+  readonly client: mongo.MongoClient;
+  readonly db: mongo.Db;
+
+  constructor(client: mongo.MongoClient, options?: PowerSyncMongoOptions) {
+    this.client = client;
+
+    const db = client.db(options?.database, {
+      ...BSON_DESERIALIZE_OPTIONS
+    });
+    this.db = db;
+
+    this.current_data = db.collection<CurrentDataDocument>('current_data');
+    this.bucket_data = db.collection('bucket_data');
+    this.bucket_parameters = db.collection('bucket_parameters');
+    this.op_id_sequence = db.collection('op_id_sequence');
+    this.sync_rules = db.collection('sync_rules');
+    this.source_tables = db.collection('source_tables');
+    this.write_checkpoints = db.collection('write_checkpoints');
+    this.instance = db.collection('instance');
+    this.locks = this.db.collection('locks');
+  }
+
+  async clear() {
+    await this.current_data.deleteMany({});
+    await this.bucket_data.deleteMany({});
+    await this.bucket_parameters.deleteMany({});
+    await this.op_id_sequence.deleteMany({});
+    await this.sync_rules.deleteMany({});
+    await this.source_tables.deleteMany({});
+    await this.write_checkpoints.deleteMany({});
+    await this.instance.deleteOne({});
+    await this.locks.deleteMany({});
+  }
+}

package/src/storage/mongo/models.ts

@@ -0,0 +1,162 @@
+import * as bson from 'bson';
+import { SqliteJsonValue } from '@powersync/service-sync-rules';
+
+export interface SourceKey {
+  /** group_id */
+  g: number;
+  /** source table id */
+  t: bson.ObjectId;
+  /** source key */
+  k: bson.UUID;
+}
+
+export interface BucketDataKey {
+  /** group_id */
+  g: number;
+  /** bucket name */
+  b: string;
+  /** op_id */
+  o: bigint;
+}
+
+export interface CurrentDataDocument {
+  _id: SourceKey;
+  data: bson.Binary;
+  buckets: CurrentBucket[];
+  lookups: bson.Binary[];
+}
+
+export interface CurrentBucket {
+  bucket: string;
+  table: string;
+  id: string;
+}
+
+export interface BucketParameterDocument {
+  _id: bigint;
+  key: SourceKey;
+  lookup: bson.Binary;
+  bucket_parameters: Record<string, SqliteJsonValue>[];
+}
+
+export interface BucketDataDocument {
+  _id: BucketDataKey;
+  op: OpType;
+  source_table: bson.ObjectId;
+  source_key: bson.UUID;
+  table: string;
+  row_id: string;
+  checksum: number;
+  data: string | null;
+}
+
+export type OpType = 'PUT' | 'REMOVE' | 'MOVE' | 'CLEAR';
+
+export interface SourceTableDocument {
+  _id: bson.ObjectId;
+  group_id: number;
+  connection_id: number;
+  relation_id: number | undefined;
+  schema_name: string;
+  table_name: string;
+  replica_id_columns: string[] | null;
+  replica_id_columns2: { name: string; type_oid: number }[] | undefined;
+  snapshot_done: boolean | undefined;
+}
+
+export interface IdSequenceDocument {
+  _id: string;
+  op_id: bigint;
+}
+
+export enum SyncRuleState {
+  /**
+   * New sync rules - needs to be processed (initial replication).
+   *
+   * While multiple sets of sync rules _can_ be in PROCESSING,
+   * it's generally pointless, so we only keep one in that state.
+   */
+  PROCESSING = 'PROCESSING',
+
+  /**
+   * Sync rule processing is done, and can be used for sync.
+   *
+   * Only one set of sync rules should be in ACTIVE state.
+   */
+  ACTIVE = 'ACTIVE',
+  /**
+   * This state is used when the sync rules has been replaced,
+   * and replication is or should be stopped.
+   */
+  STOP = 'STOP',
+  /**
+   * After sync rules have been stopped, the data needs to be
+   * deleted. Once deleted, the state is TERMINATED.
+   */
+  TERMINATED = 'TERMINATED'
+}
+
+export interface SyncRuleDocument {
+  _id: number;
+
+  state: SyncRuleState;
+
+  /**
+   * True if initial snapshot has been replicated.
+   *
+   * Can only be false if state == PROCESSING.
+   */
+  snapshot_done: boolean;
+
+  /**
+   * The last consistent checkpoint.
+   *
+   * There may be higher OpIds used in the database if we're in the middle of replicating a large transaction.
+   */
+  last_checkpoint: bigint | null;
+
+  /**
+   * The LSN associated with the last consistent checkpoint.
+   */
+  last_checkpoint_lsn: string | null;
+
+  /**
+   * If set, no new checkpoints may be created < this value.
+   */
+  no_checkpoint_before: string | null;
+
+  slot_name: string | null;
+
+  /**
+   * Last time we persisted a checkpoint.
+   *
+   * This may be old if no data is incoming.
+   */
+  last_checkpoint_ts: Date | null;
+
+  /**
+   * Last time we persisted a checkpoint or keepalive.
+   *
+   * This should stay fairly current while replicating.
+   */
+  last_keepalive_ts: Date | null;
+
+  /**
+   * If an error is stopping replication, it will be stored here.
+   */
+  last_fatal_error: string | null;
+
+  content: string;
+}
+
+export interface WriteCheckpointDocument {
+  _id: bson.ObjectId;
+  user_id: string;
+  lsns: Record<string, string>;
+  client_id: bigint;
+}
+
+export interface InstanceDocument {
+  // The instance UUID
+  _id: string;
+}

package/src/storage/mongo/util.ts

@@ -0,0 +1,88 @@
+import { SqliteJsonValue } from '@powersync/service-sync-rules';
+import * as bson from 'bson';
+import * as mongo from 'mongodb';
+import * as crypto from 'crypto';
+
+/**
+ * Lookup serialization must be number-agnostic. I.e. normalize numbers, instead of preserving numbers.
+ * @param lookup
+ */
+
+export function serializeLookup(lookup: SqliteJsonValue[]) {
+  const normalized = lookup.map((value) => {
+    if (typeof value == 'number' && Number.isInteger(value)) {
+      return BigInt(value);
+    } else {
+      return value;
+    }
+  });
+  return new bson.Binary(bson.serialize({ l: normalized }));
+}
+
+export function idPrefixFilter<T>(prefix: Partial<T>, rest: (keyof T)[]): mongo.Condition<T> {
+  let filter = {
+    $gte: {
+      ...prefix
+    } as any,
+    $lt: {
+      ...prefix
+    } as any
+  };
+
+  for (let key of rest) {
+    filter.$gte[key] = new bson.MinKey();
+    filter.$lt[key] = new bson.MaxKey();
+  }
+
+  return filter;
+}
+
+export function generateSlotName(prefix: string, sync_rules_id: number) {
+  const slot_suffix = crypto.randomBytes(2).toString('hex');
+  return `${prefix}${sync_rules_id}_${slot_suffix}`;
+}
+
+/**
+ * Read a single batch of data from a cursor, then close it.
+ *
+ * We do our best to avoid MongoDB fetching any more data than this single batch.
+ *
+ * This is similar to using `singleBatch: true` in find options.
+ * However, that makes `has_more` detection very difficult, since the cursor is always closed
+ * after the first batch. Instead, we do a workaround to only fetch a single batch below.
+ *
+ * For this to be effective, set batchSize = limit in the find command.
+ */
+export async function readSingleBatch<T>(cursor: mongo.FindCursor<T>): Promise<{ data: T[]; hasMore: boolean }> {
+  try {
+    let data: T[];
+    let hasMore = true;
+    // Let MongoDB load the first batch of data
+    const hasAny = await cursor.hasNext();
+    // Now it's in memory, and we can read it
+    data = cursor.readBufferedDocuments();
+    if (!hasAny || cursor.id?.isZero()) {
+      // A zero id means the cursor is exhaused.
+      // No results (hasAny == false) means even this batch doesn't have data.
+      // This should similar results as `await cursor.hasNext()`, but without
+      // actually fetching the next batch.
+      //
+      // Note that it is safe (but slightly inefficient) to return `hasMore: true`
+      // without there being more data, as long as the next batch
+      // will return `hasMore: false`.
+      hasMore = false;
+    }
+    return { data, hasMore };
+  } finally {
+    // Match the from the cursor iterator logic here:
+    // https://github.com/mongodb/node-mongodb-native/blob/e02534e7d1c627bf50b85ca39f5995dbf165ad44/src/cursor/abstract_cursor.ts#L327-L331
+    if (!cursor.closed) {
+      await cursor.close();
+    }
+  }
+}
+
+export const BSON_DESERIALIZE_OPTIONS: bson.DeserializeOptions = {
+  // use bigint instead of Long
+  useBigInt64: true
+};