Refactor write checkpoint creation.

Author: rkistner

modules/module-mongodb/src/api/MongoRouteAPIAdapter.ts

@@ -1,6 +1,6 @@
 import * as lib_mongo from '@powersync/lib-service-mongodb';
 import { mongo } from '@powersync/lib-service-mongodb';
-import { api, ParseSyncRulesOptions, SourceTable } from '@powersync/service-core';
+import { api, ParseSyncRulesOptions, ReplicationHeadCallback, SourceTable } from '@powersync/service-core';
 import * as sync_rules from '@powersync/service-sync-rules';
 import * as service_types from '@powersync/service-types';
 
@@ -9,6 +9,8 @@ import { constructAfterRecord, createCheckpoint } from '../replication/MongoRela
 import { CHECKPOINTS_COLLECTION } from '../replication/replication-utils.js';
 import * as types from '../types/types.js';
 import { escapeRegExp } from '../utils.js';
+import { ServiceAssertionError } from '@powersync/lib-services-framework';
+import { MongoLSN } from '../common/MongoLSN.js';
 
 export class MongoRouteAPIAdapter implements api.RouteAPI {
   protected client: mongo.MongoClient;
@@ -208,6 +210,44 @@ export class MongoRouteAPIAdapter implements api.RouteAPI {
     return createCheckpoint(this.client, this.db);
   }
 
+  async createReplicationHead<T>(callback: ReplicationHeadCallback<T>): Promise<T> {
+    const session = this.client.startSession();
+    try {
+      await this.db.command({ hello: 1 }, { session });
+      const head = session.clusterTime?.clusterTime;
+      if (head == null) {
+        throw new ServiceAssertionError(`clusterTime not available for write checkpoint`);
+      }
+
+      const r = await callback(new MongoLSN({ timestamp: head }).comparable);
+
+      // Trigger a change on the changestream.
+      await this.db.collection(CHECKPOINTS_COLLECTION).findOneAndUpdate(
+        {
+          _id: 'checkpoint' as any
+        },
+        {
+          $inc: { i: 1 }
+        },
+        {
+          upsert: true,
+          returnDocument: 'after',
+          session
+        }
+      );
+      const time = session.operationTime!;
+      if (time == null) {
+        throw new ServiceAssertionError(`operationTime not available for write checkpoint`);
+      } else if (time.lt(head)) {
+        throw new ServiceAssertionError(`operationTime must be > clusterTime`);
+      }
+
+      return r;
+    } finally {
+      await session.endSession();
+    }
+  }
+
   async getConnectionSchema(): Promise<service_types.DatabaseSchema[]> {
     const sampleSize = 50;
 

modules/module-mysql/src/api/MySQLRouteAPIAdapter.ts

@@ -1,4 +1,4 @@
-import { api, ParseSyncRulesOptions, storage } from '@powersync/service-core';
+import { api, ParseSyncRulesOptions, ReplicationHeadCallback, storage } from '@powersync/service-core';
 
 import * as sync_rules from '@powersync/service-sync-rules';
 import * as service_types from '@powersync/service-types';
@@ -275,6 +275,15 @@ export class MySQLRouteAPIAdapter implements api.RouteAPI {
     return result.comparable;
   }
 
+  async createReplicationHead<T>(callback: ReplicationHeadCallback<T>): Promise<T> {
+    const head = await this.getReplicationHead();
+    const r = await callback(head);
+
+    // TODO: make sure another message is replicated
+
+    return r;
+  }
+
   async getConnectionSchema(): Promise<service_types.DatabaseSchema[]> {
     const [results] = await this.retriedQuery({
       query: `

modules/module-postgres/src/api/PostgresRouteAPIAdapter.ts

@@ -1,6 +1,6 @@
 import * as lib_postgres from '@powersync/lib-service-postgres';
 import { ErrorCode, ServiceError } from '@powersync/lib-services-framework';
-import { api, ParseSyncRulesOptions } from '@powersync/service-core';
+import { api, ParseSyncRulesOptions, ReplicationHeadCallback } from '@powersync/service-core';
 import * as pgwire from '@powersync/service-jpgwire';
 import * as sync_rules from '@powersync/service-sync-rules';
 import * as service_types from '@powersync/service-types';
@@ -241,17 +241,22 @@ FROM pg_replication_slots WHERE slot_name = $1 LIMIT 1;`,
     // However, on Aurora (Postgres compatible), it can return an entirely different LSN,
     // causing the write checkpoints to never be replicated back to the client.
     // For those, we need to use pg_current_wal_lsn() instead.
-    const { results } = await lib_postgres.retriedQuery(
-      this.pool,
-      { statement: `SELECT pg_current_wal_lsn() as lsn` },
-      KEEPALIVE_STATEMENT
-    );
+    const { results } = await lib_postgres.retriedQuery(this.pool, `SELECT pg_current_wal_lsn() as lsn`);
 
-    // Specifically use the lsn from the first statement, not the second one.
     const lsn = results[0].rows[0][0];
     return String(lsn);
   }
 
+  async createReplicationHead<T>(callback: ReplicationHeadCallback<T>): Promise<T> {
+    const currentLsn = await this.getReplicationHead();
+
+    const r = await callback(currentLsn);
+
+    await lib_postgres.retriedQuery(this.pool, KEEPALIVE_STATEMENT);
+
+    return r;
+  }
+
   async getConnectionSchema(): Promise<service_types.DatabaseSchema[]> {
     // https://github.com/Borvik/vscode-postgres/blob/88ec5ed061a0c9bced6c5d4ec122d0759c3f3247/src/language/server.ts
     const results = await lib_postgres.retriedQuery(

packages/service-core/src/api/RouteAPI.ts

@@ -54,6 +54,14 @@ export interface RouteAPI {
    */
   getReplicationHead(): Promise<string>;
 
+  /**
+   * Get the current LSN or equivalent replication HEAD position identifier.
+   *
+   * The position is provided to the callback. After the callback returns,
+   * the replication head or a greater one will be streamed on the replication stream.
+   */
+  createReplicationHead<T>(callback: ReplicationHeadCallback<T>): Promise<T>;
+
   /**
    * @returns The schema for tables inside the connected database. This is typically
    *          used to validate sync rules.
@@ -76,3 +84,5 @@ export interface RouteAPI {
    */
   getParseSyncRulesOptions(): ParseSyncRulesOptions;
 }
+
+export type ReplicationHeadCallback<T> = (head: string) => Promise<T>;

packages/service-core/src/util/checkpointing.ts

@@ -11,18 +11,21 @@ export interface CreateWriteCheckpointOptions {
 export async function createWriteCheckpoint(options: CreateWriteCheckpointOptions) {
   const full_user_id = checkpointUserId(options.userId, options.clientId);
 
-  const currentCheckpoint = await options.api.getReplicationHead();
-
   const activeSyncRules = await options.storage.getActiveSyncRulesContent();
   if (!activeSyncRules) {
     throw new ServiceError(ErrorCode.PSYNC_S2302, `Cannot create Write Checkpoint since no sync rules are active.`);
   }
 
   using syncBucketStorage = options.storage.getInstance(activeSyncRules);
-  const writeCheckpoint = await syncBucketStorage.createManagedWriteCheckpoint({
-    user_id: full_user_id,
-    heads: { '1': currentCheckpoint }
+
+  const { writeCheckpoint, currentCheckpoint } = await options.api.createReplicationHead(async (currentCheckpoint) => {
+    const writeCheckpoint = await syncBucketStorage.createManagedWriteCheckpoint({
+      user_id: full_user_id,
+      heads: { '1': currentCheckpoint }
+    });
+    return { writeCheckpoint, currentCheckpoint };
   });
+
   return {
     writeCheckpoint: String(writeCheckpoint),
     replicationHead: currentCheckpoint