@powersync/service-core 1.23.0 → 1.23.2

package/src/replication/AbstractReplicator.ts

@@ -1,6 +1,7 @@
 import { container, ErrorCode, logger } from '@powersync/lib-services-framework';
 import { ReplicationMetric } from '@powersync/service-types';
 import { hrtime } from 'node:process';
+import { setTimeout as sleep } from 'node:timers/promises';
 import winston from 'winston';
 import { MetricsEngine } from '../metrics/MetricsEngine.js';
 import * as storage from '../storage/storage-index.js';
@@ -13,6 +14,12 @@ import { ConnectionTestResult } from './ReplicationModule.js';
 // 1 minute
 const PING_INTERVAL = 1_000_000_000n * 60n;
 
+// In the initial startup period, we use a short refresh interval. This helps to take over replication quickly in the case
+// of rolling deploys. After the initial period, we switch to a longer refresh interval to reduce load.
+const FAST_REFRESH_INTERVAL_MS = 500;
+const FAST_REFRESH_TIMEOUT_MS = 120_000;
+const REFRESH_INTERVAL_MS = 5_000;
+
 export interface CreateJobOptions {
   lock: storage.ReplicationLock;
   storage: storage.SyncRulesBucketStorage;
@@ -36,7 +43,7 @@ export interface AbstractReplicatorOptions {
  */
 export abstract class AbstractReplicator<T extends AbstractReplicationJob = AbstractReplicationJob> {
   protected logger: winston.Logger;
-  private lockAlerted: boolean = false;
+  private lastReplicationStreamInfoLogs = new Map<string, string>();
   /**
    *  Map of replication jobs by replication stream job id. Usually there is only one running job, but there could be two
    *  when transitioning to a new replication stream.
@@ -65,6 +72,11 @@ export abstract class AbstractReplicator<T extends AbstractReplicationJob = Abst
     this.logger = logger.child({ name: `Replicator:${options.id}` });
   }
 
+  /**
+   * Create the source-specific replication job for a locked storage stream.
+   *
+   * The job owns the per-stream snapshot and streaming work for this attempt.
+   */
   abstract createJob(options: CreateJobOptions): T;
 
   /**
@@ -131,17 +143,20 @@ export abstract class AbstractReplicator<T extends AbstractReplicationJob = Abst
   }
 
   private async runLoop() {
-    const syncRules = await this.syncRuleProvider.get();
+    const loadedSyncConfig = await this.syncRuleProvider.get();
 
     let configuredLock: storage.ReplicationLock | undefined = undefined;
-    if (syncRules != null) {
+    if (loadedSyncConfig != null) {
       this.logger.info('Loaded sync config');
       try {
         // Configure new sync config, if they have changed.
         // In that case, also immediately take out a lock, so that another process doesn't start replication on it.
+        // This upfront lock is mostly superseded by the replicationStreamLoadedSyncConfigMatch() check. However, that doesn't cover old
+        // versions before that check was added, so we keep the lock for now - for where te service version and sync config is updated at
+        // the same time.
 
         const { lock } = await this.storage.configureSyncRules(
-          storage.updateSyncRulesFromYaml(syncRules, { lock: true, validate: this.syncRuleProvider.exitOnError })
+          storage.updateSyncRulesFromYaml(loadedSyncConfig, { lock: true, validate: this.syncRuleProvider.exitOnError })
         );
         if (lock) {
           configuredLock = lock;
@@ -155,10 +170,15 @@ export abstract class AbstractReplicator<T extends AbstractReplicationJob = Abst
     } else {
       this.logger.info('No sync streams or rules configured - configure via API');
     }
+    let useFastRefresh = true;
+    const fastRefreshDeadline = Date.now() + FAST_REFRESH_TIMEOUT_MS;
     while (!this.stopped) {
       await container.probes.touch();
       try {
-        await this.refresh({ configured_lock: configuredLock });
+        const refreshResult = await this.refresh({ configuredLock, loadedSyncConfig });
+        if (refreshResult.replicationJobStarted || Date.now() >= fastRefreshDeadline) {
+          useFastRefresh = false;
+        }
         // The lock is only valid on the first refresh.
         configuredLock = undefined;
 
@@ -177,24 +197,40 @@ export abstract class AbstractReplicator<T extends AbstractReplicationJob = Abst
       } catch (e) {
         this.logger.error('Failed to refresh replication jobs', e);
       }
-      await new Promise((resolve) => setTimeout(resolve, 5000));
+      if (Date.now() >= fastRefreshDeadline) {
+        useFastRefresh = false;
+      }
+      await sleep(useFastRefresh ? FAST_REFRESH_INTERVAL_MS : REFRESH_INTERVAL_MS);
     }
   }
 
-  private async refresh(options?: { configured_lock?: storage.ReplicationLock }) {
+  private async refresh(options?: {
+    configuredLock?: storage.ReplicationLock;
+    loadedSyncConfig?: string;
+  }): Promise<{ replicationJobStarted: boolean }> {
     if (this.stopped) {
-      return;
+      return { replicationJobStarted: false };
     }
 
-    let configuredLock = options?.configured_lock;
+    let configuredLock = options?.configuredLock;
+    let replicationJobStarted = false;
 
     const existingJobs = new Map<string, T>(this.replicationJobs.entries());
     const replicatingStreams = await this.storage.getReplicatingReplicationStreams();
     const newJobs = new Map<string, T>();
+    const streamsToStart: storage.PersistedReplicationStream[] = [];
     let activeJob: T | undefined = undefined;
     for (let replicationStream of replicatingStreams) {
       const jobId = replicationStream.replicationJobId;
       const existingJob = existingJobs.get(jobId);
+      const syncConfigMismatchMessage = 'Ignoring replication stream for sync config not loaded by this process';
+      if (!this.shouldHandleReplicationStream(replicationStream, options?.loadedSyncConfig)) {
+        this.logReplicationStreamInfoOnce(replicationStream, 'sync-config-mismatch', () => {
+          replicationStream.logger.info(syncConfigMismatchMessage);
+        });
+        continue;
+      }
+      this.clearReplicationStreamInfoLog(replicationStream, 'sync-config-mismatch');
       if (replicationStream.state == storage.SyncRuleState.ACTIVE && activeJob == null) {
         activeJob = existingJob;
       }
@@ -208,47 +244,12 @@ export abstract class AbstractReplicator<T extends AbstractReplicationJob = Abst
         existingJobs.delete(jobId);
       } else {
         // New sync config was found (or resume after restart)
-        try {
-          let lock: storage.ReplicationLock;
-          if (configuredLock?.sync_rules_id == replicationStream.replicationStreamId) {
-            lock = configuredLock;
-          } else {
-            lock = await replicationStream.lock();
-          }
-          const syncRuleStorage = this.storage.getInstance(replicationStream);
-          const newJob = this.createJob({
-            lock: lock,
-            storage: syncRuleStorage
-          });
-
-          newJobs.set(jobId, newJob);
-          newJob.start();
-          if (replicationStream.state == storage.SyncRuleState.ACTIVE) {
-            activeJob = newJob;
-          }
-          this.lockAlerted = false;
-        } catch (e) {
-          if (e?.errorData?.code === ErrorCode.PSYNC_S1003) {
-            if (!this.lockAlerted) {
-              replicationStream.logger.info(`[${e.errorData.code}] ${e.errorData.description}`);
-              this.lockAlerted = true;
-            }
-          } else {
-            // Could be a sync config parse error,
-            // for example from stricter validation that was added.
-            // This will be retried every couple of seconds.
-            // When new (valid) sync config is deployed and processed, this one be disabled.
-            replicationStream.logger.error('Failed to start replication for new sync config', e);
-          }
-        }
+        streamsToStart.push(replicationStream);
       }
     }
 
-    this.replicationJobs = newJobs;
-    this.activeReplicationJob = activeJob;
-
-    // Stop any orphaned jobs that no longer have a replication stream.
-    // Termination happens below
+    // Stop any orphaned jobs that no longer have a replication stream before starting replacements.
+    // Termination happens below.
     for (let job of existingJobs.values()) {
       // Old - stop and clean up
       try {
@@ -259,6 +260,46 @@ export abstract class AbstractReplicator<T extends AbstractReplicationJob = Abst
       }
     }
 
+    for (let replicationStream of streamsToStart) {
+      const jobId = replicationStream.replicationJobId;
+      try {
+        let lock: storage.ReplicationLock;
+        if (configuredLock?.sync_rules_id == replicationStream.replicationStreamId) {
+          lock = configuredLock;
+        } else {
+          lock = await replicationStream.lock();
+        }
+        const syncRuleStorage = this.storage.getInstance(replicationStream);
+        const newJob = this.createJob({
+          lock: lock,
+          storage: syncRuleStorage
+        });
+
+        newJobs.set(jobId, newJob);
+        newJob.start();
+        replicationJobStarted = true;
+        if (replicationStream.state == storage.SyncRuleState.ACTIVE) {
+          activeJob = newJob;
+        }
+        this.lastReplicationStreamInfoLogs.delete(replicationStream.replicationStreamName);
+      } catch (e) {
+        if (e?.errorData?.code === ErrorCode.PSYNC_S1003) {
+          this.logReplicationStreamInfoOnce(replicationStream, 'replication-stream-locked', () => {
+            replicationStream.logger.info(`[${e.errorData.code}] ${e.errorData.description}`);
+          });
+        } else {
+          // Could be a sync config parse error,
+          // for example from stricter validation that was added.
+          // This will be retried every couple of seconds.
+          // When new (valid) sync config is deployed and processed, this one be disabled.
+          replicationStream.logger.error('Failed to start replication for new sync config', e);
+        }
+      }
+    }
+
+    this.replicationJobs = newJobs;
+    this.activeReplicationJob = activeJob;
+
     // Replication stream stopped previously, including by a different process.
     const stopped = await this.storage.getStoppedReplicationStreams();
     for (let replicationStream of stopped) {
@@ -280,6 +321,49 @@ export abstract class AbstractReplicator<T extends AbstractReplicationJob = Abst
         });
       this.clearingJobs.set(replicationStream.replicationStreamId, promise);
     }
+
+    return { replicationJobStarted };
+  }
+
+  private logReplicationStreamInfoOnce(
+    replicationStream: storage.PersistedReplicationStream,
+    category: string,
+    log: () => void
+  ) {
+    if (this.lastReplicationStreamInfoLogs.get(replicationStream.replicationStreamName) == category) {
+      return;
+    }
+
+    log();
+    this.lastReplicationStreamInfoLogs.set(replicationStream.replicationStreamName, category);
+  }
+
+  private clearReplicationStreamInfoLog(replicationStream: storage.PersistedReplicationStream, category: string) {
+    if (this.lastReplicationStreamInfoLogs.get(replicationStream.replicationStreamName) == category) {
+      this.lastReplicationStreamInfoLogs.delete(replicationStream.replicationStreamName);
+    }
+  }
+
+  /**
+   * When we load sync config from the filesystem/config source, we ignore any config loaded by other processes.
+   * The idea is that if a different process loads updated config, that process should process it, not this one.
+   *
+   * This specifically helps for cases with rolling deploys.
+   *
+   * This does not apply when sync config is loaded from the database/API, instead of in the config.
+   */
+  private shouldHandleReplicationStream(
+    replicationStream: storage.PersistedReplicationStream,
+    loadedSyncRules: string | undefined
+  ) {
+    if (loadedSyncRules == null) {
+      return true;
+    }
+
+    const processingConfig = replicationStream.syncConfigContent.find(
+      (syncConfig) => syncConfig.syncConfigState == storage.SyncRuleState.PROCESSING
+    );
+    return processingConfig == null || processingConfig.sync_rules_content == loadedSyncRules;
   }
 
   protected createJobId(syncRuleId: number) {
@@ -295,6 +379,9 @@ export abstract class AbstractReplicator<T extends AbstractReplicationJob = Abst
     syncRuleStorage.logger.info(`Successfully terminated replication stream`);
   }
 
+  /**
+   * Check source connectivity using the replicator's decoded connection config.
+   */
   abstract testConnection(): Promise<ConnectionTestResult>;
 
   /**

package/src/routes/RouterEngine.ts

@@ -11,6 +11,7 @@ export type RouterSetupResponse = {
 
 export type RouterEngineRoutes = {
   api_routes: RouteDefinition[];
+  checkpointing_routes: RouteDefinition[];
   stream_routes: RouteDefinition[];
   socket_routes: SocketRouteGenerator[];
 };
@@ -42,6 +43,7 @@ export class RouterEngine {
 
     this.routes = {
       api_routes: [],
+      checkpointing_routes: [],
       stream_routes: [],
       socket_routes: []
     };
@@ -49,13 +51,17 @@ export class RouterEngine {
 
   public registerRoutes(routes: Partial<RouterEngineRoutes>) {
     this.routes.api_routes.push(...(routes.api_routes ?? []));
+    this.routes.checkpointing_routes.push(...(routes.checkpointing_routes ?? []));
     this.routes.stream_routes.push(...(routes.stream_routes ?? []));
     this.routes.socket_routes.push(...(routes.socket_routes ?? []));
   }
 
   public get hasRoutes() {
     return (
-      this.routes.api_routes.length > 0 || this.routes.stream_routes.length > 0 || this.routes.socket_routes.length > 0
+      this.routes.api_routes.length > 0 ||
+      this.routes.checkpointing_routes.length > 0 ||
+      this.routes.stream_routes.length > 0 ||
+      this.routes.socket_routes.length > 0
     );
   }
 

package/src/routes/configure-fastify.ts

@@ -31,6 +31,7 @@ export type RouteRegistrationOptions = {
  */
 export type RouteDefinitions = {
   api?: Partial<RouteRegistrationOptions>;
+  checkpointing?: Partial<RouteRegistrationOptions>;
   sync_stream?: Partial<RouteRegistrationOptions>;
 };
 
@@ -41,12 +42,20 @@ export type FastifyServerConfig = {
 
 export const DEFAULT_ROUTE_OPTIONS = {
   api: {
-    routes: [...ADMIN_ROUTES, ...CHECKPOINT_ROUTES, ...SYNC_RULES_ROUTES, ...PROBES_ROUTES],
+    routes: [...ADMIN_ROUTES, ...SYNC_RULES_ROUTES, ...PROBES_ROUTES],
     queue_options: {
       concurrency: 10,
       max_queue_depth: 20
     }
   },
+  checkpointing: {
+    routes: [...CHECKPOINT_ROUTES],
+    queue_options: {
+      // Note that the values here has an effect on WriteCheckpointBatcher
+      concurrency: 100,
+      max_queue_depth: 100
+    }
+  },
   sync_stream: {
     routes: [...SYNC_STREAM_ROUTES],
     queue_options: {
@@ -72,6 +81,7 @@ export function configureFastifyServer(server: fastify.FastifyInstance, options:
 
   // Set on the outer server so both child scopes inherit.
   registerFastifyErrorHandler(server);
+  registerFastifyNotFoundHandler(server);
 
   /**
    * Fastify creates an encapsulated context for each `.register` call.
@@ -81,7 +91,6 @@ export function configureFastifyServer(server: fastify.FastifyInstance, options:
    */
   server.register(async function (childContext) {
     registerFastifyRoutes(childContext, generateContext, routes.api?.routes ?? DEFAULT_ROUTE_OPTIONS.api.routes);
-    registerFastifyNotFoundHandler(childContext);
 
     // Limit the active concurrent requests
     childContext.addHook(
@@ -90,6 +99,20 @@ export function configureFastifyServer(server: fastify.FastifyInstance, options:
     );
   });
 
+  server.register(async function (childContext) {
+    registerFastifyRoutes(
+      childContext,
+      generateContext,
+      routes.checkpointing?.routes ?? DEFAULT_ROUTE_OPTIONS.checkpointing.routes
+    );
+
+    // Limit the active concurrent requests
+    childContext.addHook(
+      'onRequest',
+      createRequestQueueHook(routes.checkpointing?.queue_options ?? DEFAULT_ROUTE_OPTIONS.checkpointing.queue_options)
+    );
+  });
+
   // Create a separate context for concurrency queueing
   server.register(async function (childContext) {
     registerFastifyRoutes(

package/src/routes/endpoints/admin.ts

@@ -241,7 +241,8 @@ export const validate = routeDefinition({
       replicationStreamName: '',
       storageVersion: storage.LEGACY_STORAGE_VERSION,
       sync_rules_content: content,
-      compiled_plan: null
+      compiled_plan: null,
+      syncConfigState: storage.SyncRuleState.PROCESSING
     });
 
     const connectionStatus = await apiHandler.getConnectionStatus();

package/src/routes/endpoints/checkpointing.ts

@@ -54,13 +54,10 @@ export const writeCheckpoint2 = routeDefinition({
   handler: async (payload) => {
     const { token_payload, service_context } = payload.context;
 
-    const apiHandler = service_context.routerEngine.getAPI();
-
     const { replicationHead, writeCheckpoint } = await util.createWriteCheckpoint({
       userId: token_payload!.userIdString,
       clientId: payload.params.client_id,
-      api: apiHandler,
-      storage: service_context.storageEngine.activeBucketStorage
+      batcher: service_context.writeCheckpointBatcher
     });
 
     logger.info(

package/src/storage/BucketStorageBatch.ts

@@ -12,6 +12,15 @@ export const DEFAULT_BUCKET_BATCH_COMMIT_OPTIONS: ResolvedBucketBatchCommitOptio
   oldestUncommittedChange: null
 };
 
+/**
+ * Storage writer for a source-specific unit of replication work.
+ *
+ * A batch evaluates source row changes against the hydrated sync config,
+ * persists bucket operations and parameter lookup rows, tracks source table and
+ * snapshot state, records the committed source position for visible
+ * checkpoints, and stores restart cursors for sources with PowerSync-managed
+ * resume state.
+ */
 export interface BucketStorageBatch extends ObserverClient<BucketBatchStorageListener>, AsyncDisposable {
   /**
    * Alias for [Symbol.asyncDispose]
@@ -29,21 +38,28 @@ export interface BucketStorageBatch extends ObserverClient<BucketBatchStorageLis
   readonly skipExistingRows: boolean;
 
   /**
-   * Save an op, and potentially flush.
+   * Queue one normalized source row operation and potentially flush.
    *
-   * This can be an insert, update or delete op.
+   * When flushed, the operation is evaluated against data and parameter
+   * queries. That may produce bucket operations, parameter lookup rows, and
+   * current-data updates. `save()` may trigger that flush automatically when
+   * the batch reaches its size limits, but a flush does not create a checkpoint
+   * by itself.
    */
   save(record: SaveOptions): Promise<FlushedResult | null>;
 
   /**
-   * Replicate a truncate op - deletes all data in the specified tables.
+   * Replicate a truncate operation by removing all currently replicated rows
+   * for the specified source tables.
    */
   truncate(sourceTables: SourceTable[]): Promise<FlushedResult | null>;
 
   /**
-   * Drop one or more tables.
+   * Drop one or more source table mappings.
    *
-   * This is the same as truncate, but additionally removes the SourceTable record.
+   * This performs the same logical row removal as `truncate()`, then removes
+   * the stored `SourceTable` record. It is used for table drops, renames,
+   * replica identity changes, and superseded mappings.
    */
   drop(sourceTables: SourceTable[]): Promise<FlushedResult | null>;
 
@@ -58,21 +74,26 @@ export interface BucketStorageBatch extends ObserverClient<BucketBatchStorageLis
   flush(options?: BatchBucketFlushOptions): Promise<FlushedResult | null>;
 
   /**
-   * Flush and commit any saved ops. This creates a new checkpoint by default.
+   * Flush saved ops and advance the committed source checkpoint position.
    *
-   * Only call this after a transaction.
+   * Only call this at a source transaction, page, snapshot, or other boundary
+   * where it is valid for clients to observe all flushed changes together. A
+   * single commit may include multiple complete source transactions, but must
+   * not split one source transaction across checkpoint boundaries.
    */
   commit(lsn: string, options?: BucketBatchCommitOptions): Promise<CheckpointResult>;
 
   /**
-   * Advance the checkpoint LSN position, without any associated op.
+   * Advance the committed source checkpoint position without associated row
+   * operations.
    *
-   * This must only be called when not inside a transaction.
+   * Use this for source heartbeats, marker-only events, and idle-source write
+   * checkpoint progress. This must only be called outside a source transaction.
    */
   keepalive(lsn: string): Promise<CheckpointResult>;
 
   /**
-   * Set the LSN that replication should resume from.
+   * Set the source position that replication should resume from after restart.
    *
    * This can be used for:
    * 1. Setting the LSN for a snapshot, before starting replication.
@@ -80,6 +101,10 @@ export interface BucketStorageBatch extends ObserverClient<BucketBatchStorageLis
    *
    * Not required if the source database keeps track of this, for example with
    * PostgreSQL logical replication slots.
+   *
+   * This is for restart safety only. If the same source position should become
+   * visible to clients, also call `commit()` or `keepalive()` at a valid source
+   * boundary.
    */
   setResumeLsn(lsn: string): Promise<void>;
 
@@ -90,7 +115,21 @@ export interface BucketStorageBatch extends ObserverClient<BucketBatchStorageLis
    */
   resumeFromLsn: string | null;
 
+  /**
+   * Mark table snapshots complete and record the source position that
+   * checkpoints must not precede for those tables.
+   */
   markTableSnapshotDone(tables: SourceTable[], no_checkpoint_before_lsn?: string): Promise<SourceTable[]>;
+
+  /**
+   * Mark a source table as requiring snapshot work again.
+   *
+   * Use this when a table that was previously considered snapshot-complete
+   * must be copied again before it can safely contribute to checkpoints. Common
+   * cases include detected schema or replica identity changes, interrupted
+   * re-snapshot work, or other source metadata changes that make the existing
+   * replicated rows no longer trustworthy.
+   */
   markTableSnapshotRequired(table: SourceTable): Promise<void>;
 
   /**
@@ -113,6 +152,10 @@ export interface BucketStorageBatch extends ObserverClient<BucketBatchStorageLis
    */
   markSnapshotDone(no_checkpoint_before_lsn: string, options?: { throwOnConflict?: boolean }): Promise<void>;
 
+  /**
+   * Persist resumable table snapshot progress such as row counts or source
+   * cursors.
+   */
   updateTableProgress(table: SourceTable, progress: Partial<TableSnapshotStatus>): Promise<SourceTable>;
 
   /**
@@ -120,10 +163,17 @@ export interface BucketStorageBatch extends ObserverClient<BucketBatchStorageLis
    */
   getSourceTableStatus(table: SourceTable): Promise<SourceTable | null>;
 
+  /**
+   * Resolve discovered source metadata against the active sync config and
+   * persisted source table state.
+   *
+   * @returns Current mappings that should receive replicated rows, plus
+   *          outdated mappings that the source connector should drop.
+   */
   resolveTables(options: ResolveTablesOptions): Promise<ResolveTablesResult>;
 
   /**
-   * Queues the creation of a custom Write Checkpoint. This will be persisted after operations are flushed.
+   * Queue a custom write checkpoint to be persisted after operations are flushed.
    */
   addCustomWriteCheckpoint(checkpoint: BatchedCustomWriteCheckpointOptions): void;
 }

package/src/storage/ChecksumCache.ts

@@ -5,7 +5,7 @@ import { BucketChecksum } from '../util/protocol-types.js';
 import { addBucketChecksums, ChecksumMap, InternalOpId, PartialChecksum } from '../util/utils.js';
 import { BucketChecksumRequest } from './SyncRulesBucketStorage.js';
 
-interface ChecksumFetchContext {
+interface ChecksumFetchContext<TContext> {
   fetch(bucket: string): Promise<BucketChecksum>;
   checkpoint: InternalOpId;
 }
@@ -20,26 +20,35 @@ export interface FetchPartialBucketChecksum {
 export type PartialOrFullChecksum = PartialChecksum | BucketChecksum;
 export type PartialChecksumMap = Map<string, PartialOrFullChecksum>;
 
-export type FetchChecksums = (batch: FetchPartialBucketChecksum[]) => Promise<PartialChecksumMap>;
+export type FetchChecksums<TContext = undefined> = (
+  batch: FetchPartialBucketChecksum[],
+  context: TContext
+) => Promise<PartialChecksumMap>;
 
-export interface ChecksumCacheOptions {
+export interface ChecksumCacheOptions<TContext = undefined> {
   /**
    * Upstream checksum implementation.
    *
    * This fetches a batch of either entire bucket checksums, or a partial range.
    */
-  fetchChecksums: FetchChecksums;
+  fetchChecksums: FetchChecksums<TContext>;
 
   /**
    * Maximum number of cached checksums.
    */
   maxSize?: number;
+
+  /**
+   * How long to keep cached checksums before refreshing them.
+   */
+  ttlMs?: number;
 }
 
 // Approximately 5MB of memory, if we assume 50 bytes per entry
 const DEFAULT_MAX_SIZE = 100_000;
 
-const TTL_MS = 3_600_000;
+// Default to 1h
+export const DEFAULT_CHECKSUM_CACHE_TTL_MS = 1 * 3_600_000;
 
 /**
  * Implement a LRU cache for checksum requests. Each (bucket, checkpoint) request is cached separately,
@@ -50,22 +59,22 @@ const TTL_MS = 3_600_000;
  *
  * We use the LRUCache fetchMethod to deduplicate in-progress requests.
  */
-export class ChecksumCache {
+export class ChecksumCache<TContext = undefined> {
   /**
    * The primary checksum cache, with key of `${checkpoint}/${bucket}`.
    */
-  private cache: LRUCache<string, BucketChecksum, ChecksumFetchContext>;
+  private cache: LRUCache<string, BucketChecksum, ChecksumFetchContext<TContext>>;
   /**
    * For each bucket, an ordered set of cached checkpoints.
    */
   private bucketCheckpoints = new Map<string, OrderedSet<bigint>>();
 
-  private fetchChecksums: FetchChecksums;
+  private fetchChecksums: FetchChecksums<TContext>;
 
-  constructor(options: ChecksumCacheOptions) {
+  constructor(options: ChecksumCacheOptions<TContext>) {
     this.fetchChecksums = options.fetchChecksums;
 
-    this.cache = new LRUCache<string, BucketChecksum, ChecksumFetchContext>({
+    this.cache = new LRUCache<string, BucketChecksum, ChecksumFetchContext<TContext>>({
       max: options.maxSize ?? DEFAULT_MAX_SIZE,
       fetchMethod: async (cacheKey, _staleValue, options) => {
         // Called when this checksum hasn't been cached yet.
@@ -105,7 +114,7 @@ export class ChecksumCache {
       // We use a TTL so that counts can eventually be refreshed
       // after a compact. This only has effect if the bucket has
       // not been checked in the meantime.
-      ttl: TTL_MS,
+      ttl: options.ttlMs ?? DEFAULT_CHECKSUM_CACHE_TTL_MS,
       ttlResolution: 1_000,
       allowStale: false
     });
@@ -129,7 +138,11 @@ export class ChecksumCache {
    *
    * @returns a Map with exactly one entry for each bucket requested
    */
-  async getChecksumMap(checkpoint: InternalOpId, buckets: BucketChecksumRequest[]): Promise<ChecksumMap> {
+  async getChecksumMap(
+    checkpoint: InternalOpId,
+    buckets: BucketChecksumRequest[],
+    fetchContext?: TContext
+  ): Promise<ChecksumMap> {
     // Buckets that don't have a cached checksum for this checkpoint yet
     let toFetch = new Set<string>();
 
@@ -147,7 +160,7 @@ export class ChecksumCache {
     // Accumulated results - both from cached checksums, and fetched checksums
     let finalResults = new Map<string, BucketChecksum>();
 
-    const context: ChecksumFetchContext = {
+    const context: ChecksumFetchContext<TContext> = {
       async fetch(bucket) {
         await fetchPromise;
         if (!toFetch.has(bucket)) {
@@ -260,7 +273,7 @@ export class ChecksumCache {
         }
 
         // Fetch partial checksums from upstream
-        const results = await this.fetchChecksums(bucketRequests);
+        const results = await this.fetchChecksums(bucketRequests, fetchContext as TContext);
 
         for (let bucket of toFetch) {
           const result = results.get(bucket);