@powersync/service-core 0.0.0-dev-20240620165206

package/src/storage/BucketStorage.ts

@@ -0,0 +1,396 @@
+import {
+  EvaluatedParameters,
+  EvaluatedRow,
+  SqlSyncRules,
+  SqliteJsonRow,
+  SqliteJsonValue,
+  SqliteRow,
+  ToastableSqliteRow
+} from '@powersync/service-sync-rules';
+
+import * as replication from '../replication/replication-index.js';
+import * as util from '../util/util-index.js';
+import { SourceTable } from './SourceTable.js';
+
+export interface BucketStorageFactory {
+  /**
+   * Update sync rules from configuration, if changed.
+   */
+  configureSyncRules(
+    sync_rules: string,
+    options?: { lock?: boolean }
+  ): Promise<{ updated: boolean; persisted_sync_rules?: PersistedSyncRulesContent; lock?: ReplicationLock }>;
+
+  /**
+   * Get a storage instance to query sync data for specific sync rules.
+   */
+  getInstance(options: PersistedSyncRules): SyncRulesBucketStorage;
+
+  /**
+   * Deploy new sync rules.
+   */
+  updateSyncRules(options: UpdateSyncRulesOptions): Promise<PersistedSyncRulesContent>;
+
+  /**
+   * Indicate that a slot was removed, and we should re-sync by creating
+   * a new sync rules instance.
+   *
+   * This is roughly the same as deploying a new version of the current sync
+   * rules, but also accounts for cases where the current sync rules are not
+   * the latest ones.
+   *
+   * Replication should be restarted after this.
+   *
+   * @param slot_name The removed slot
+   */
+  slotRemoved(slot_name: string): Promise<void>;
+
+  /**
+   * Get the sync rules used for querying.
+   */
+  getActiveSyncRules(): Promise<PersistedSyncRules | null>;
+
+  /**
+   * Get the sync rules used for querying.
+   */
+  getActiveSyncRulesContent(): Promise<PersistedSyncRulesContent | null>;
+
+  /**
+   * Get the sync rules that will be active next once done with initial replicatino.
+   */
+  getNextSyncRules(): Promise<PersistedSyncRules | null>;
+
+  /**
+   * Get the sync rules that will be active next once done with initial replicatino.
+   */
+  getNextSyncRulesContent(): Promise<PersistedSyncRulesContent | null>;
+
+  /**
+   * Get all sync rules currently replicating. Typically this is the "active" and "next" sync rules.
+   */
+  getReplicatingSyncRules(): Promise<PersistedSyncRulesContent[]>;
+
+  /**
+   * Get all sync rules stopped but not terminated yet.
+   */
+  getStoppedSyncRules(): Promise<PersistedSyncRulesContent[]>;
+
+  /**
+   * Same as:
+   * getInstance(await getActiveSyncRules()).getCheckpoint().
+   */
+  getActiveCheckpoint(): Promise<ActiveCheckpoint>;
+
+  createWriteCheckpoint(user_id: string, lsns: Record<string, string>): Promise<bigint>;
+
+  lastWriteCheckpoint(user_id: string, lsn: string): Promise<bigint | null>;
+
+  watchWriteCheckpoint(user_id: string, signal: AbortSignal): AsyncIterable<WriteCheckpoint>;
+
+  /**
+   * Get storage size of active sync rules.
+   */
+  getStorageMetrics(): Promise<StorageMetrics>;
+
+  /**
+   * Get the unique identifier for this instance of Powersync
+   */
+  getPowerSyncInstanceId(): Promise<string>;
+}
+
+export interface WriteCheckpoint {
+  base: ActiveCheckpoint;
+  writeCheckpoint: bigint | null;
+}
+
+export interface ActiveCheckpoint {
+  readonly checkpoint: util.OpId;
+  readonly lsn: string;
+
+  hasSyncRules(): boolean;
+
+  getBucketStorage(): Promise<SyncRulesBucketStorage | null>;
+}
+
+export interface StorageMetrics {
+  /**
+   * Size of operations (bucket_data)
+   */
+  operations_size_bytes: number;
+
+  /**
+   * Size of parameter storage.
+   *
+   * Replication storage -> raw data as received from Postgres.
+   */
+  parameters_size_bytes: number;
+
+  /**
+   * Size of current_data.
+   */
+  replication_size_bytes: number;
+}
+
+export interface PersistedSyncRulesContent {
+  readonly id: number;
+  readonly sync_rules_content: string;
+  readonly slot_name: string;
+
+  readonly last_fatal_error?: string | null;
+  readonly last_keepalive_ts?: Date | null;
+  readonly last_checkpoint_ts?: Date | null;
+
+  parsed(): PersistedSyncRules;
+
+  lock(): Promise<ReplicationLock>;
+}
+
+export interface ReplicationLock {
+  sync_rules_id: number;
+
+  release(): Promise<void>;
+}
+
+export interface PersistedSyncRules {
+  readonly id: number;
+  readonly sync_rules: SqlSyncRules;
+  readonly slot_name: string;
+}
+
+export class DefaultPersistedSyncRules implements PersistedSyncRules {
+  public readonly checkpoint_lsn: string | null;
+
+  constructor(public readonly id: number, public readonly sync_rules: SqlSyncRules, checkpoint_lsn: string | null) {
+    this.checkpoint_lsn = checkpoint_lsn;
+  }
+
+  get slot_name(): string {
+    return `powersync_${this.id}`;
+  }
+}
+
+export interface UpdateSyncRulesOptions {
+  content: string;
+  lock?: boolean;
+}
+
+export interface SyncRulesBucketStorageOptions {
+  sync_rules: SqlSyncRules;
+  group_id: number;
+}
+
+export const DEFAULT_DOCUMENT_BATCH_LIMIT = 1000;
+export const DEFAULT_DOCUMENT_CHUNK_LIMIT_BYTES = 1 * 1024 * 1024;
+
+export interface BucketDataBatchOptions {
+  /** Limit number of documents returned. Defaults to 1000. */
+  limit?: number;
+
+  /**
+   * Limit size of chunks returned. Defaults to 1MB.
+   *
+   * This is a lower bound, not an upper bound. As soon as the chunk size goes over this limit,
+   * it is returned.
+   *
+   * Note that an individual data row can be close to 16MB in size, so this does not help in
+   * extreme cases.
+   */
+  chunkLimitBytes?: number;
+}
+
+export interface SyncRulesBucketStorage {
+  readonly sync_rules: SqlSyncRules;
+  readonly group_id: number;
+  readonly slot_name: string;
+
+  readonly factory: BucketStorageFactory;
+
+  resolveTable(options: ResolveTableOptions): Promise<ResolveTableResult>;
+
+  startBatch(options: {}, callback: (batch: BucketStorageBatch) => Promise<void>): Promise<FlushedResult | null>;
+
+  getCheckpoint(): Promise<{ checkpoint: util.OpId; lsn: string }>;
+
+  getParameterSets(checkpoint: util.OpId, lookups: SqliteJsonValue[][]): Promise<SqliteJsonRow[]>;
+
+  /**
+   * Get a "batch" of data for a checkpoint.
+   *
+   * The results will be split into separate SyncBucketData chunks to:
+   * 1. Separate buckets.
+   * 2. Limit the size of each individual chunk according to options.batchSizeLimitBytes.
+   *
+   * @param checkpoint the checkpoint
+   * @param dataBuckets current bucket states
+   * @param options batch size options
+   */
+  getBucketDataBatch(
+    checkpoint: util.OpId,
+    dataBuckets: Map<string, string>,
+    options?: BucketDataBatchOptions
+  ): AsyncIterable<util.SyncBucketData>;
+
+  /**
+   * Compute checksums for a given list of buckets.
+   *
+   * Returns zero checksums for any buckets not found.
+   */
+  getChecksums(checkpoint: util.OpId, buckets: string[]): Promise<util.ChecksumMap>;
+
+  /**
+   * Terminate the sync rules.
+   *
+   * This clears the storage, and sets state to TERMINATED.
+   *
+   * Must only be called on stopped sync rules.
+   */
+  terminate(): Promise<void>;
+
+  getStatus(): Promise<SyncRuleStatus>;
+
+  /**
+   * Clear the storage, without changing state.
+   */
+  clear(): Promise<void>;
+
+  setSnapshotDone(lsn: string): Promise<void>;
+
+  autoActivate(): Promise<void>;
+
+  /**
+   * Record a replication error.
+   *
+   * This could be a recoverable error (e.g. temporary network failure),
+   * or a permanent error (e.g. missing toast data).
+   *
+   * Errors are cleared on commit.
+   */
+  reportError(e: any): Promise<void>;
+}
+
+export interface SyncRuleStatus {
+  checkpoint_lsn: string | null;
+  active: boolean;
+  snapshot_done: boolean;
+}
+export interface ResolveTableOptions {
+  group_id: number;
+  connection_id: number;
+  connection_tag: string;
+  relation: replication.PgRelation;
+
+  sync_rules: SqlSyncRules;
+}
+
+export interface ResolveTableResult {
+  table: SourceTable;
+  dropTables: SourceTable[];
+}
+
+export interface FlushedResult {
+  flushed_op: string;
+}
+
+export interface BucketStorageBatch {
+  /**
+   * Save an op, and potentially flush.
+   *
+   * This can be an insert, update or delete op.
+   */
+  save(record: SaveOptions): Promise<FlushedResult | null>;
+
+  /**
+   * Replicate a truncate op - deletes all data in the specified tables.
+   */
+  truncate(sourceTables: SourceTable[]): Promise<FlushedResult | null>;
+
+  /**
+   * Drop one or more tables.
+   *
+   * This is the same as truncate, but additionally removes the SourceTable record.
+   */
+  drop(sourceTables: SourceTable[]): Promise<FlushedResult | null>;
+
+  /**
+   * Explicitly flush all pending changes in the batch.
+   *
+   * This does not create a new checkpoint until `commit()` is called. This means it's
+   * safe to flush multiple times in the middle of a large transaction.
+   *
+   * @returns null if there are no changes to flush.
+   */
+  flush(): Promise<FlushedResult | null>;
+
+  /**
+   * Flush and commit any saved ops. This creates a new checkpoint.
+   *
+   * Only call this after a transaction.
+   */
+  commit(lsn: string): Promise<boolean>;
+
+  /**
+   * Advance the checkpoint LSN position, without any associated op.
+   *
+   * This must only be called when not inside a transaction.
+   *
+   * @returns true if the checkpoint was advanced, false if this was a no-op
+   */
+  keepalive(lsn: string): Promise<boolean>;
+
+  markSnapshotDone(tables: SourceTable[], no_checkpoint_before_lsn: string): Promise<SourceTable[]>;
+}
+
+export interface SaveParameterData {
+  sourceTable: SourceTable;
+  /** UUID */
+  sourceKey: string;
+  evaluated: EvaluatedParameters[];
+}
+
+export interface SaveBucketData {
+  sourceTable: SourceTable;
+  /** UUID */
+  sourceKey: string;
+
+  evaluated: EvaluatedRow[];
+}
+
+export type SaveOptions = SaveInsert | SaveUpdate | SaveDelete;
+
+export interface SaveInsert {
+  tag: 'insert';
+  sourceTable: SourceTable;
+  before?: undefined;
+  after: SqliteRow;
+}
+
+export interface SaveUpdate {
+  tag: 'update';
+  sourceTable: SourceTable;
+  before?: SqliteRow;
+  /**
+   * A null value means null column.
+   *
+   * An undefined value means it's a TOAST value - must be copied from another record.
+   */
+  after: ToastableSqliteRow;
+}
+
+export interface SaveDelete {
+  tag: 'delete';
+  sourceTable: SourceTable;
+  before: SqliteRow;
+  after?: undefined;
+}
+
+export function mergeToast(record: ToastableSqliteRow, persisted: ToastableSqliteRow): ToastableSqliteRow {
+  const newRecord: ToastableSqliteRow = {};
+  for (let key in record) {
+    if (typeof record[key] == 'undefined') {
+      newRecord[key] = persisted[key];
+    } else {
+      newRecord[key] = record[key];
+    }
+  }
+  return newRecord;
+}

package/src/storage/ChecksumCache.ts

@@ -0,0 +1,294 @@
+import { BucketChecksum, OpId } from '../util/protocol-types.js';
+import { ChecksumMap, addBucketChecksums } from '../util/utils.js';
+import { LRUCache } from 'lru-cache/min';
+import { OrderedSet } from '@js-sdsl/ordered-set';
+
+interface ChecksumFetchContext {
+  fetch(bucket: string): Promise<BucketChecksum>;
+  checkpoint: bigint;
+}
+
+export interface FetchPartialBucketChecksum {
+  bucket: string;
+  start?: OpId;
+  end: OpId;
+}
+
+export type FetchChecksums = (batch: FetchPartialBucketChecksum[]) => Promise<ChecksumMap>;
+
+export interface ChecksumCacheOptions {
+  /**
+   * Upstream checksum implementation.
+   *
+   * This fetches a batch of either entire bucket checksums, or a partial range.
+   */
+  fetchChecksums: FetchChecksums;
+
+  /**
+   * Maximum number of cached checksums.
+   */
+  maxSize?: number;
+}
+
+// Approximately 5MB of memory, if we assume 50 bytes per entry
+const DEFAULT_MAX_SIZE = 100_000;
+
+/**
+ * Implement a LRU cache for checksum requests. Each (bucket, checkpoint) request is cached separately,
+ * while the lookups occur in batches.
+ *
+ * For each bucket, we keep a separate OrderedSet of cached checkpoints.
+ * This allows us to do incrementally update checksums by using the last cached checksum for the same bucket.
+ *
+ * We use the LRUCache fetchMethod to deduplicate in-progress requests.
+ */
+export class ChecksumCache {
+  /**
+   * The primary checksum cache, with key of `${checkpoint}/${bucket}`.
+   */
+  private cache: LRUCache<string, BucketChecksum, ChecksumFetchContext>;
+  /**
+   * For each bucket, an ordered set of cached checkpoints.
+   */
+  private bucketCheckpoints = new Map<string, OrderedSet<bigint>>();
+
+  private fetchChecksums: FetchChecksums;
+
+  constructor(options: ChecksumCacheOptions) {
+    this.fetchChecksums = options.fetchChecksums;
+
+    this.cache = new LRUCache<string, BucketChecksum, ChecksumFetchContext>({
+      max: options.maxSize ?? DEFAULT_MAX_SIZE,
+      fetchMethod: async (cacheKey, _staleValue, options) => {
+        // Called when this checksum hasn't been cached yet.
+        // Pass the call back to the request, which implements batch fetching.
+        const { bucket } = parseCacheKey(cacheKey);
+        const result = await options.context.fetch(bucket);
+
+        // Add to the set of cached checkpoints for the bucket.
+        let checkpointSet = this.bucketCheckpoints.get(bucket);
+        if (checkpointSet == null) {
+          checkpointSet = new OrderedSet();
+          this.bucketCheckpoints.set(bucket, checkpointSet);
+        }
+        checkpointSet.insert(options.context.checkpoint);
+        return result;
+      },
+
+      dispose: (value, key) => {
+        // Remove from the set of cached checkpoints for the bucket
+        const { checkpointString } = parseCacheKey(key);
+        const checkpoint = BigInt(checkpointString);
+        const checkpointSet = this.bucketCheckpoints.get(value.bucket);
+        if (checkpointSet == null) {
+          return;
+        }
+        checkpointSet.eraseElementByKey(checkpoint);
+        if (checkpointSet.length == 0) {
+          this.bucketCheckpoints.delete(value.bucket);
+        }
+      },
+
+      noDisposeOnSet: true,
+
+      // When we have more fetches than the cache size, complete the fetches instead
+      // of failing with Error('evicted').
+      ignoreFetchAbort: true
+    });
+  }
+
+  async getChecksums(checkpoint: OpId, buckets: string[]): Promise<BucketChecksum[]> {
+    const checksums = await this.getChecksumMap(checkpoint, buckets);
+    // Return results in the same order as the request
+    return buckets.map((bucket) => checksums.get(bucket)!);
+  }
+
+  /**
+   * Get bucket checksums for a checkpoint.
+   *
+   * Any checksums not found upstream are returned as zero checksums.
+   *
+   * @returns a Map with exactly one entry for each bucket requested
+   */
+  async getChecksumMap(checkpoint: OpId, buckets: string[]): Promise<ChecksumMap> {
+    // Buckets that don't have a cached checksum for this checkpoint yet
+    let toFetch = new Set<string>();
+
+    // Newly fetched results
+    let fetchResults = new Map<string, BucketChecksum>();
+
+    // Promise for the bactch new fetch requests
+    let resolveFetch!: () => void;
+    let rejectFetch!: (err: any) => void;
+    let fetchPromise = new Promise<void>((resolve, reject) => {
+      resolveFetch = resolve;
+      rejectFetch = reject;
+    });
+
+    // Accumulated results - both from cached checksums, and fetched checksums
+    let finalResults = new Map<string, BucketChecksum>();
+
+    const context: ChecksumFetchContext = {
+      async fetch(bucket) {
+        await fetchPromise;
+        if (!toFetch.has(bucket)) {
+          // Should never happen
+          throw new Error(`Expected to fetch ${bucket}`);
+        }
+        const checksum = fetchResults.get(bucket);
+        if (checksum == null) {
+          // Should never happen
+          throw new Error(`Failed to fetch checksum for bucket ${bucket}`);
+        }
+        return checksum;
+      },
+      checkpoint: BigInt(checkpoint)
+    };
+
+    // One promise to await to ensure all fetch requests completed.
+    let settledPromise: Promise<PromiseSettledResult<void>[]> | null = null;
+
+    try {
+      // Individual cache fetch promises
+      let cacheFetchPromises: Promise<void>[] = [];
+
+      for (let bucket of buckets) {
+        const cacheKey = makeCacheKey(checkpoint, bucket);
+        let status: LRUCache.Status<BucketChecksum> = {};
+        const p = this.cache.fetch(cacheKey, { context: context, status: status }).then((checksums) => {
+          if (checksums == null) {
+            // Should never happen
+            throw new Error(`Failed to get checksums for ${cacheKey}`);
+          }
+          finalResults.set(bucket, checksums);
+        });
+        cacheFetchPromises.push(p);
+        if (status.fetch == 'hit' || status.fetch == 'inflight') {
+          // The checksums is either cached already (hit), or another request is busy
+          // fetching (inflight).
+          // In either case, we don't need to fetch a new checksum.
+        } else {
+          // We need a new request for this checksum.
+          toFetch.add(bucket);
+        }
+      }
+      // We do this directly after creating the promises, otherwise
+      // we could end up with weird uncaught rejection errors.
+      settledPromise = Promise.allSettled(cacheFetchPromises);
+
+      if (toFetch.size == 0) {
+        // Nothing to fetch, but resolve in case
+        resolveFetch();
+      } else {
+        let bucketRequests: FetchPartialBucketChecksum[] = [];
+        // Partial checksum (previously cached) to add to the partial fetch
+        let add = new Map<string, BucketChecksum>();
+
+        for (let bucket of toFetch) {
+          let bucketRequest: FetchPartialBucketChecksum | null = null;
+          const checkpointSet = this.bucketCheckpoints.get(bucket);
+          if (checkpointSet != null) {
+            // Find smaller checkpoints, sorted in descending order
+            let iter = checkpointSet.reverseUpperBound(context.checkpoint);
+            const begin = checkpointSet.begin();
+            while (iter.isAccessible()) {
+              const cp = iter.pointer;
+              const cacheKey = makeCacheKey(cp, bucket);
+              // peek to avoid refreshing the key
+              const cached = this.cache.peek(cacheKey);
+              // As long as dispose() works correctly, the checkpointset should
+              // match up with the cache, and `cached` should also have a value here.
+              // However, we handle caces where it's not present either way.
+              // Test by disabling the `dispose()` callback.
+              if (cached != null) {
+                // Partial checksum found - make a partial checksum request
+                bucketRequest = {
+                  bucket,
+                  start: cp.toString(),
+                  end: checkpoint
+                };
+                add.set(bucket, cached);
+                break;
+              }
+
+              if (iter.equals(begin)) {
+                // Cannot iterate further
+                break;
+              }
+              // Iterate backwards
+              iter = iter.pre();
+            }
+          }
+
+          if (bucketRequest == null) {
+            // No partial checksum found - make a new full checksum request
+            bucketRequest = {
+              bucket,
+              end: checkpoint
+            };
+            add.set(bucket, {
+              bucket,
+              checksum: 0,
+              count: 0
+            });
+          }
+          bucketRequests.push(bucketRequest);
+        }
+
+        // Fetch partial checksums from upstream
+        const results = await this.fetchChecksums(bucketRequests);
+
+        for (let bucket of toFetch) {
+          const result = results.get(bucket);
+          const toAdd = add.get(bucket);
+          if (toAdd == null) {
+            // Should never happen
+            throw new Error(`toAdd null for ${bucket}`);
+          }
+          // Compute the full checksum from the two partials.
+          // No results returned are treated the same as a zero result.
+          const added = addBucketChecksums(toAdd, result ?? null);
+          fetchResults.set(bucket, added);
+        }
+
+        // fetchResults is fully populated, so we resolve the Promise
+        resolveFetch();
+      }
+    } catch (e) {
+      // Failure when fetching checksums - reject the Promise.
+      // This will reject all individual cache fetch requests, and each will be retried
+      // on the next request.
+      rejectFetch(e);
+
+      // Wait for the above rejection to propagate, otherwise we end up with "uncaught" errors.
+      // This promise never throws.
+      await settledPromise;
+
+      throw e;
+    }
+
+    // Wait for all cache fetch reqeusts to complete
+    const settledResults = (await settledPromise) ?? [];
+    // Check if any of them failed
+    for (let result of settledResults) {
+      if (result.status == 'rejected') {
+        throw result.reason;
+      }
+    }
+
+    if (finalResults.size != buckets.length) {
+      // Should not happen
+      throw new Error(`Bucket results mismatch: ${finalResults.size} != ${buckets.length}`);
+    }
+    return finalResults;
+  }
+}
+
+function makeCacheKey(checkpoint: bigint | string, bucket: string) {
+  return `${checkpoint}/${bucket}`;
+}
+
+function parseCacheKey(key: string) {
+  const index = key.indexOf('/');
+  return { checkpointString: key.substring(0, index), bucket: key.substring(index + 1) };
+}