@powersync/service-module-postgres-storage 0.16.3 → 0.17.0

package/src/storage/batch/PostgresBucketBatch.ts

@@ -60,7 +60,6 @@ export interface PostgresBucketBatchOptions {
  * via the Postgres NOTIFY protocol.
  */
 const StatefulCheckpoint = models.ActiveCheckpoint.and(t.object({ state: t.Enum(storage.SyncRuleState) }));
-type StatefulCheckpointDecoded = t.Decoded<typeof StatefulCheckpoint>;
 
 const CheckpointWithStatus = StatefulCheckpoint.and(
   t.object({
@@ -465,12 +464,11 @@ export class PostgresBucketBatch
         }
         const { flushedAny } = await persistedBatch.flush(db);
         clearedError = flushedAny && !this.clearedError;
-        if (clearedError) {
-          // No need to clear an error more than once per batch, since an error would always result in restarting the batch.
-          await this.clearError(db);
-        }
       });
       if (clearedError) {
+        // No need to clear an error more than once per batch, since an error would always result in restarting the batch.
+        // Cleared outside the replication transaction - see flushInner.
+        await this.clearError();
         this.clearedError = true;
       }
     }
@@ -537,6 +535,10 @@ export class PostgresBucketBatch
     });
 
     if (clearedError) {
+      // Clear the error outside the replication transaction (plain autocommit update,
+      // like the keepalive), to avoid serialization conflicts on the sync_rules row
+      // when multiple writers flush concurrently.
+      await this.clearError();
       this.clearedError = true;
     }
 
@@ -562,136 +564,144 @@ export class PostgresBucketBatch
 
     const persisted_op = this.persisted_op ?? null;
 
-    const result = await this.db.sql`
-      WITH
-        selected AS (
-          SELECT
-            id,
-            state,
-            last_checkpoint,
-            last_checkpoint_lsn,
-            snapshot_done,
-            no_checkpoint_before,
-            keepalive_op,
-            (
-              snapshot_done = TRUE
-              AND (
-                last_checkpoint_lsn IS NULL
-                OR last_checkpoint_lsn <= ${{ type: 'varchar', value: lsn }}
-              )
+    // Transaction failures restart replication from the last durable checkpoint.
+    const result = await this.db.transaction(async (db) => {
+      const checkpoint = await db.sql`
+        WITH
+          selected AS (
+            SELECT
+              id,
+              state,
+              last_checkpoint,
+              last_checkpoint_lsn,
+              snapshot_done,
+              no_checkpoint_before,
+              keepalive_op,
+              (
+                snapshot_done = TRUE
+                AND (
+                  last_checkpoint_lsn IS NULL
+                  OR last_checkpoint_lsn <= ${{ type: 'varchar', value: lsn }}
+                )
+                AND (
+                  no_checkpoint_before IS NULL
+                  OR no_checkpoint_before <= ${{ type: 'varchar', value: lsn }}
+                )
+              ) AS can_checkpoint
+            FROM
+              sync_rules
+            WHERE
+              id = ${{ type: 'int4', value: this.group_id }}
+            FOR UPDATE
+          ),
+          computed AS (
+            SELECT
+              selected.*,
+              CASE
+                WHEN selected.can_checkpoint THEN GREATEST(
+                  selected.last_checkpoint,
+                  ${{ type: 'int8', value: persisted_op }},
+                  selected.keepalive_op,
+                  0
+                )
+                ELSE selected.last_checkpoint
+              END AS new_last_checkpoint,
+              CASE
+                WHEN selected.can_checkpoint THEN NULL
+                ELSE GREATEST(
+                  selected.keepalive_op,
+                  ${{ type: 'int8', value: persisted_op }},
+                  0
+                )
+              END AS new_keepalive_op
+            FROM
+              selected
+          ),
+          updated AS (
+            UPDATE sync_rules AS sr
+            SET
+              last_checkpoint_lsn = CASE
+                WHEN computed.can_checkpoint THEN ${{ type: 'varchar', value: lsn }}
+                ELSE sr.last_checkpoint_lsn
+              END,
+              last_checkpoint_ts = CASE
+                WHEN computed.can_checkpoint THEN ${{ type: 1184, value: now }}
+                ELSE sr.last_checkpoint_ts
+              END,
+              last_keepalive_ts = ${{ type: 1184, value: now }},
+              last_fatal_error = CASE
+                WHEN computed.can_checkpoint THEN NULL
+                ELSE sr.last_fatal_error
+              END,
+              keepalive_op = computed.new_keepalive_op,
+              last_checkpoint = computed.new_last_checkpoint,
+              snapshot_lsn = CASE
+                WHEN computed.can_checkpoint THEN NULL
+                ELSE sr.snapshot_lsn
+              END
+            FROM
+              computed
+            WHERE
+              sr.id = computed.id
               AND (
-                no_checkpoint_before IS NULL
-                OR no_checkpoint_before <= ${{ type: 'varchar', value: lsn }}
-              )
-            ) AS can_checkpoint
-          FROM
-            sync_rules
-          WHERE
-            id = ${{ type: 'int4', value: this.group_id }}
-          FOR UPDATE
-        ),
-        computed AS (
-          SELECT
-            selected.*,
-            CASE
-              WHEN selected.can_checkpoint THEN GREATEST(
-                selected.last_checkpoint,
-                ${{ type: 'int8', value: persisted_op }},
-                selected.keepalive_op,
-                0
+                sr.keepalive_op IS DISTINCT FROM computed.new_keepalive_op
+                OR sr.last_checkpoint IS DISTINCT FROM computed.new_last_checkpoint
+                OR ${{ type: 'bool', value: createEmptyCheckpoints }}
               )
-              ELSE selected.last_checkpoint
-            END AS new_last_checkpoint,
-            CASE
-              WHEN selected.can_checkpoint THEN NULL
-              ELSE GREATEST(
-                selected.keepalive_op,
-                ${{ type: 'int8', value: persisted_op }},
-                0
-              )
-            END AS new_keepalive_op
-          FROM
-            selected
-        ),
-        updated AS (
-          UPDATE sync_rules AS sr
-          SET
-            last_checkpoint_lsn = CASE
-              WHEN computed.can_checkpoint THEN ${{ type: 'varchar', value: lsn }}
-              ELSE sr.last_checkpoint_lsn
-            END,
-            last_checkpoint_ts = CASE
-              WHEN computed.can_checkpoint THEN ${{ type: 1184, value: now }}
-              ELSE sr.last_checkpoint_ts
-            END,
-            last_keepalive_ts = ${{ type: 1184, value: now }},
-            last_fatal_error = CASE
-              WHEN computed.can_checkpoint THEN NULL
-              ELSE sr.last_fatal_error
-            END,
-            keepalive_op = computed.new_keepalive_op,
-            last_checkpoint = computed.new_last_checkpoint,
-            snapshot_lsn = CASE
-              WHEN computed.can_checkpoint THEN NULL
-              ELSE sr.snapshot_lsn
-            END
-          FROM
-            computed
-          WHERE
-            sr.id = computed.id
-            AND (
-              sr.keepalive_op IS DISTINCT FROM computed.new_keepalive_op
-              OR sr.last_checkpoint IS DISTINCT FROM computed.new_last_checkpoint
-              OR ${{ type: 'bool', value: createEmptyCheckpoints }}
-            )
-          RETURNING
-            sr.id,
-            sr.state,
-            sr.last_checkpoint,
-            sr.last_checkpoint_lsn,
-            sr.snapshot_done,
-            sr.no_checkpoint_before,
-            computed.can_checkpoint,
-            computed.keepalive_op,
-            computed.new_last_checkpoint
-        )
-      SELECT
-        id,
-        state,
-        last_checkpoint,
-        last_checkpoint_lsn,
-        snapshot_done,
-        no_checkpoint_before,
-        can_checkpoint,
-        keepalive_op,
-        new_last_checkpoint,
-        TRUE AS created_checkpoint
-      FROM
-        updated
-      UNION ALL
-      SELECT
-        id,
-        state,
-        new_last_checkpoint AS last_checkpoint,
-        last_checkpoint_lsn,
-        snapshot_done,
-        no_checkpoint_before,
-        can_checkpoint,
-        keepalive_op,
-        new_last_checkpoint,
-        FALSE AS created_checkpoint
-      FROM
-        computed
-      WHERE
-        NOT EXISTS (
-          SELECT
-            1
-          FROM
-            updated
-        )
-    `
-      .decoded(CheckpointWithStatus)
-      .first();
+            RETURNING
+              sr.id,
+              sr.state,
+              sr.last_checkpoint,
+              sr.last_checkpoint_lsn,
+              sr.snapshot_done,
+              sr.no_checkpoint_before,
+              computed.can_checkpoint,
+              computed.keepalive_op,
+              computed.new_last_checkpoint
+          )
+        SELECT
+          id,
+          state,
+          last_checkpoint,
+          last_checkpoint_lsn,
+          snapshot_done,
+          no_checkpoint_before,
+          can_checkpoint,
+          keepalive_op,
+          new_last_checkpoint,
+          TRUE AS created_checkpoint
+        FROM
+          updated
+        UNION ALL
+        SELECT
+          id,
+          state,
+          new_last_checkpoint AS last_checkpoint,
+          last_checkpoint_lsn,
+          snapshot_done,
+          no_checkpoint_before,
+          can_checkpoint,
+          keepalive_op,
+          new_last_checkpoint,
+          FALSE AS created_checkpoint
+        FROM
+          computed
+        WHERE
+          NOT EXISTS (
+            SELECT
+              1
+            FROM
+              updated
+          )
+      `
+        .decoded(CheckpointWithStatus)
+        .first();
+
+      if (checkpoint?.can_checkpoint && checkpoint.state == storage.SyncRuleState.ACTIVE) {
+        await notifySyncRulesUpdate(db, checkpoint);
+      }
+      return checkpoint;
+    });
 
     if (result == null) {
       throw new ReplicationAssertionError('Failed to update sync_rules during checkpoint');
@@ -723,10 +733,8 @@ export class PostgresBucketBatch
         });
       }
     }
-    await this.autoActivate(lsn);
-    await notifySyncRulesUpdate(this.db, {
+    await this.autoActivate(lsn, {
       id: result.id,
-      state: result.state,
       last_checkpoint: result.last_checkpoint,
       last_checkpoint_lsn: result.last_checkpoint_lsn
     });
@@ -1033,10 +1041,8 @@ export class PostgresBucketBatch
       }
     }
 
+    // The error itself is cleared by the caller, outside the replication transaction - see flushInner.
     const clearedError = didFlush && !this.clearedError;
-    if (clearedError) {
-      await this.clearError(db);
-    }
 
     // Don't return empty batches
     return {
@@ -1305,14 +1311,13 @@ export class PostgresBucketBatch
    *
    * Called on new commits.
    */
-  private async autoActivate(lsn: string): Promise<void> {
+  private async autoActivate(lsn: string, checkpoint: models.ActiveCheckpointDecoded): Promise<void> {
     if (!this.needsActivation) {
       // Already activated
       return;
     }
 
-    let didActivate = false;
-    await this.db.transaction(async (db) => {
+    const activationResult = await this.db.transaction(async (db) => {
       const syncRulesRow = await db.sql`
         SELECT
           state,
@@ -1346,13 +1351,18 @@ export class PostgresBucketBatch
             )
             AND id != ${{ type: 'int4', value: this.group_id }}
         `.execute();
-        didActivate = true;
-        this.needsActivation = false;
+        await notifySyncRulesUpdate(db, checkpoint);
+        return 'activated';
       } else if (syncRulesRow?.state != storage.SyncRuleState.PROCESSING) {
-        this.needsActivation = false;
+        return 'not-processing';
       }
+      return 'pending';
     });
-    if (didActivate) {
+
+    if (activationResult != 'pending') {
+      this.needsActivation = false;
+    }
+    if (activationResult == 'activated') {
       this.logger.info(`Activated new replication stream at ${lsn}`);
     }
   }
@@ -1438,12 +1448,17 @@ export class PostgresBucketBatch
  * Uses Postgres' NOTIFY functionality to update different processes when the
  * active checkpoint has been updated.
  */
-export const notifySyncRulesUpdate = async (db: lib_postgres.DatabaseClient, update: StatefulCheckpointDecoded) => {
-  if (update.state != storage.SyncRuleState.ACTIVE) {
-    return;
-  }
-
-  await db.query({
-    statement: `NOTIFY ${NOTIFICATION_CHANNEL}, '${models.ActiveCheckpointNotification.encode({ active_checkpoint: update })}'`
-  });
+export const notifySyncRulesUpdate = async (
+  db: lib_postgres.AbstractPostgresConnection,
+  update: models.ActiveCheckpointDecoded
+) => {
+  const payload = models.ActiveCheckpointNotification.encode({ active_checkpoint: update });
+
+  await db.sql`
+    SELECT
+      pg_notify (
+        ${{ type: 'varchar', value: NOTIFICATION_CHANNEL }},
+        ${{ type: 'varchar', value: payload }}
+      )
+  `.execute();
 };

package/src/storage/checkpoints/PostgresWriteCheckpointAPI.ts

@@ -34,52 +34,170 @@ export class PostgresWriteCheckpointAPI implements storage.WriteCheckpointAPI {
 
   async createManagedWriteCheckpoints(
     checkpoints: storage.ManagedWriteCheckpointOptions[]
-  ): Promise<Map<string, bigint>> {
+  ): Promise<storage.CreateManagedWriteCheckpointsResult> {
     if (this.writeCheckpointMode !== storage.WriteCheckpointMode.MANAGED) {
       throw new framework.errors.ValidationError(
         `Attempting to create a managed Write Checkpoint when the current Write Checkpoint mode is set to "${this.writeCheckpointMode}"`
       );
     }
 
-    const uniqueCheckpoints = [...new Map(checkpoints.map((checkpoint) => [checkpoint.user_id, checkpoint])).values()];
+    const uniqueCheckpoints = storage.uniqueManagedWriteCheckpoints(checkpoints);
     if (uniqueCheckpoints.length == 0) {
-      return new Map();
+      return { writeCheckpoints: new Map(), shouldAdvance: false };
     }
 
-    const mappedCheckpoints = uniqueCheckpoints.map((checkpoint) => ({
-      user_id: checkpoint.user_id,
-      lsns: checkpoint.heads
-    }));
+    const writeCheckpoints = new Map<string, bigint>();
+    const generatedCheckpoints = uniqueCheckpoints.filter((checkpoint) => checkpoint.checkpoint_request_id == null);
+    const suppliedCheckpoints = uniqueCheckpoints.filter((checkpoint) => checkpoint.checkpoint_request_id != null);
 
-    const rows = await this.db.sql`
-      WITH
-        json_data AS (
+    if (generatedCheckpoints.length > 0) {
+      const mappedCheckpoints = generatedCheckpoints.map((checkpoint) => ({
+        user_id: checkpoint.user_id,
+        lsns: checkpoint.heads
+      }));
+
+      const generatedRows = await this.db.sql`
+        WITH
+          json_data AS (
+            SELECT
+            CHECKPOINT ->> 'user_id' AS user_id,
+            CHECKPOINT -> 'lsns' AS lsns
+            FROM
+              jsonb_array_elements(${{ type: 'jsonb', value: mappedCheckpoints }}) AS
+            CHECKPOINT
+          )
+        INSERT INTO
+          write_checkpoints (
+            user_id,
+            lsns,
+            write_checkpoint,
+            checkpoint_requested_at
+          )
+        SELECT
+          user_id,
+          lsns,
+          1,
+          NULL
+        FROM
+          json_data
+        ON CONFLICT (user_id) DO UPDATE
+        SET
+          write_checkpoint = write_checkpoints.write_checkpoint + 1,
+          lsns = EXCLUDED.lsns,
+          checkpoint_requested_at = NULL
+        RETURNING
+          *;
+      `
+        .decoded(models.WriteCheckpoint)
+        .rows();
+
+      for (const row of generatedRows) {
+        writeCheckpoints.set(row.user_id, row.write_checkpoint);
+      }
+    }
+
+    if (suppliedCheckpoints.length > 0) {
+      // Supplied request ids are monotonic. Greater values update the checkpoint
+      // id and heads, while equal retries only refresh the retention timestamp.
+      // Stale requests return the stored id without changing the row.
+      const mappedCheckpoints = suppliedCheckpoints.map((checkpoint) => ({
+        user_id: checkpoint.user_id,
+        lsns: checkpoint.heads,
+        checkpoint_request_id: String(checkpoint.checkpoint_request_id)
+      }));
+
+      const suppliedRows = await this.db.sql`
+        WITH
+          json_data AS (
+            SELECT
+            CHECKPOINT ->> 'user_id' AS user_id,
+            CHECKPOINT -> 'lsns' AS lsns,
+            (
+              CHECKPOINT ->> 'checkpoint_request_id'
+            )::int8 AS checkpoint_request_id
+            FROM
+              jsonb_array_elements(${{ type: 'jsonb', value: mappedCheckpoints }}) AS
+            CHECKPOINT
+          )
+        INSERT INTO
+          write_checkpoints (
+            user_id,
+            lsns,
+            write_checkpoint,
+            checkpoint_requested_at
+          )
+        SELECT
+          user_id,
+          lsns,
+          checkpoint_request_id,
+          NOW()
+        FROM
+          json_data
+        ON CONFLICT (user_id) DO UPDATE
+        SET
+          write_checkpoint = CASE
+            WHEN EXCLUDED.write_checkpoint > write_checkpoints.write_checkpoint THEN EXCLUDED.write_checkpoint
+            ELSE write_checkpoints.write_checkpoint
+          END,
+          lsns = CASE
+            WHEN EXCLUDED.write_checkpoint > write_checkpoints.write_checkpoint THEN EXCLUDED.lsns
+            ELSE write_checkpoints.lsns
+          END,
+          checkpoint_requested_at = NOW()
+        WHERE
+          EXCLUDED.write_checkpoint >= write_checkpoints.write_checkpoint
+        RETURNING
+          *;
+      `
+        .decoded(models.WriteCheckpoint)
+        .rows();
+
+      for (const row of suppliedRows) {
+        writeCheckpoints.set(row.user_id, row.write_checkpoint);
+      }
+
+      // RETURNING only excludes stale requests rejected by the monotonic
+      // conflict condition. Those requests still need the stored id.
+      const returnedSuppliedUserIds = new Set(suppliedRows.map((row) => row.user_id));
+      const unchangedUserIds = suppliedCheckpoints
+        .map((checkpoint) => checkpoint.user_id)
+        .filter((userId) => !returnedSuppliedUserIds.has(userId));
+
+      if (unchangedUserIds.length > 0) {
+        const mappedUserIds = unchangedUserIds.map((user_id) => ({ user_id }));
+        const unchangedRows = await this.db.sql`
+          WITH
+            json_data AS (
+              SELECT
+              CHECKPOINT ->> 'user_id' AS user_id
+              FROM
+                jsonb_array_elements(${{ type: 'jsonb', value: mappedUserIds }}) AS
+              CHECKPOINT
+            )
           SELECT
-          CHECKPOINT ->> 'user_id' AS user_id,
-          CHECKPOINT -> 'lsns' AS lsns
+            write_checkpoints.*
           FROM
-            jsonb_array_elements(${{ type: 'jsonb', value: mappedCheckpoints }}) AS
-          CHECKPOINT
-        )
-      INSERT INTO
-        write_checkpoints (user_id, lsns, write_checkpoint)
-      SELECT
-        user_id,
-        lsns,
-        1
-      FROM
-        json_data
-      ON CONFLICT (user_id) DO UPDATE
-      SET
-        write_checkpoint = write_checkpoints.write_checkpoint + 1,
-        lsns = EXCLUDED.lsns
-      RETURNING
-        *;
-    `
-      .decoded(models.WriteCheckpoint)
-      .rows();
+            write_checkpoints
+            JOIN json_data ON write_checkpoints.user_id = json_data.user_id;
+        `
+          .decoded(models.WriteCheckpoint)
+          .rows();
+
+        for (const row of unchangedRows) {
+          writeCheckpoints.set(row.user_id, row.write_checkpoint);
+        }
+      }
+    }
 
-    return new Map(rows.map((row) => [row.user_id, row.write_checkpoint]));
+    // Postgres storage does not track a per-row processed indicator: a write
+    // checkpoint is considered processed at read time by comparing its lsns
+    // against the replicated head (see lastManagedWriteCheckpoint). We therefore
+    // force the source marker whenever any checkpoint was matched, which also
+    // covers stale or duplicate requests whose stored checkpoint may still be
+    // pending. Forcing a marker for an already-processed checkpoint is wasteful
+    // but harmless.
+    const shouldAdvance = generatedCheckpoints.length > 0 || suppliedCheckpoints.length > 0;
+    return { writeCheckpoints, shouldAdvance };
   }
 
   async lastWriteCheckpoint(filters: storage.LastWriteCheckpointFilters): Promise<bigint | null> {
@@ -155,7 +273,8 @@ export async function batchCreateCustomWriteCheckpoints(
       // Cannot encode bigint directly using JSON.stringify.
       // The ::int8 in the query below will take care of casting back to a number
       checkpoint: String(cp.checkpoint),
-      sync_rules_id: cp.sync_rules_id
+      sync_rules_id: cp.sync_rules_id,
+      checkpoint_requested_at: cp.checkpoint_requested_at?.toISOString() ?? null
     };
   });
 
@@ -167,7 +286,12 @@ export async function batchCreateCustomWriteCheckpoints(
         CHECKPOINT
       )
     INSERT INTO
-      custom_write_checkpoints (user_id, write_checkpoint, sync_rules_id)
+      custom_write_checkpoints (
+        user_id,
+        write_checkpoint,
+        sync_rules_id,
+        checkpoint_requested_at
+      )
     SELECT
     CHECKPOINT ->> 'user_id'::varchar,
     (
@@ -175,11 +299,15 @@ export async function batchCreateCustomWriteCheckpoints(
     )::int8,
     (
       CHECKPOINT ->> 'sync_rules_id'
-    )::int4
+    )::int4,
+    (
+      CHECKPOINT ->> 'checkpoint_requested_at'
+    )::timestamptz
     FROM
       json_data
     ON CONFLICT (user_id, sync_rules_id) DO UPDATE
     SET
-      write_checkpoint = EXCLUDED.write_checkpoint;
+      write_checkpoint = EXCLUDED.write_checkpoint,
+      checkpoint_requested_at = EXCLUDED.checkpoint_requested_at;
   `.execute();
 }