@rocicorp/zero 1.10.0-canary.0 → 1.10.0-canary.2

package/out/zero-cache/src/services/change-source/pg/change-source.js

@@ -18,14 +18,15 @@ import { fromBigInt, toBigInt, toStateVersionString } from "./lsn.js";
 import { getPublicationInfo } from "./schema/published.js";
 import { replicationEventSchema } from "./schema/ddl.js";
 import { validate } from "./schema/validation.js";
-import { dropShard, getInternalShardConfig, getReplicaAtVersion, internalPublicationPrefix, replicaIdentitiesForTablesWithoutPrimaryKeys, replicationSlotPrefix } from "./schema/shard.js";
+import { dropShard, getActiveReplicas, getInternalShardConfig, getReplicaAtVersion, internalPublicationPrefix, replicaIdentitiesForTablesWithoutPrimaryKeys, replicationSlotPrefix } from "./schema/shard.js";
 import { AutoResetSignal } from "../../change-streamer/schema/tables.js";
+import { restoreReplica } from "../common/replica-restore.js";
 import { initReplica } from "../common/replica-schema.js";
 import "../../../types/pg-versions.js";
 import { BackfillManager } from "../common/backfill-manager.js";
 import { ChangeStreamMultiplexer } from "../common/change-stream-multiplexer.js";
 import { dropOldReplicasAndSlots } from "./replication-slots.js";
-import { updateShardSchema } from "./schema/init.js";
+import { ensureShardSchema } from "./schema/init.js";
 import { initialSync } from "./initial-sync.js";
 import { streamBackfill } from "./backfill-stream.js";
 import { subscribe } from "./logical-replication/stream.js";
@@ -40,25 +41,57 @@ var REPLICA_SLOT_CLEANUP_INTERVAL_MS = 3e4;
 * replica, before streaming changes from the corresponding logical replication
 * stream.
 */
-async function initializePostgresChangeSource(lc, upstreamURI, shard, replicaDbFile, syncOptions, context, lagReportIntervalMs = 0) {
-	await initReplica(lc, `replica-${shard.appID}-${shard.shardNum}`, replicaDbFile, (log, tx) => initialSync(log, shard, tx, upstreamURI, syncOptions, context));
-	const replica = new Database(lc, replicaDbFile);
-	const subscriptionState = getSubscriptionStateAndContext(new StatementRunner(replica));
-	replica.close();
+async function initializePostgresChangeSource(lc, upstreamURI, shard, replicaDbFile, syncOptions, context, lagReportIntervalMs = 0, restoreOptions = {}, { backupV5 } = { backupV5: true }, purgeLock, streamInboundTimeoutMs) {
 	const db = await connectPgClient(lc, upstreamURI, "change-source-init");
 	try {
+		await ensureShardSchema(lc, db, shard);
+		const restoredReplica = await selectAndRestoreReplica(lc, db, shard, replicaDbFile, restoreOptions);
+		let initialSyncedReplica;
+		await initReplica(lc, `replica-${shard.appID}-${shard.shardNum}`, replicaDbFile, async (log, tx) => {
+			purgeLock?.release();
+			initialSyncedReplica = await initialSync(log, shard, tx, upstreamURI, syncOptions, context, { backupV5 });
+		});
+		const replica = new Database(lc, replicaDbFile);
+		const subscriptionState = getSubscriptionStateAndContext(new StatementRunner(replica));
+		replica.close();
+		const upstreamReplica = await checkAndUpdateUpstream(lc, db, shard, subscriptionState, (initialSyncedReplica ?? restoredReplica)?.id);
+		const backupPath = initialSyncedReplica ? initialSyncedReplica.backupPath : backupV5 ? String(Date.now()) : restoredReplica?.backupPath ?? null;
 		return {
 			subscriptionState,
-			changeSource: new PostgresChangeSource(lc, upstreamURI, shard, await checkAndUpdateUpstream(lc, db, shard, subscriptionState), context, lagReportIntervalMs, syncOptions.textCopy)
+			changeSource: new PostgresChangeSource(lc, upstreamURI, shard, upstreamReplica, {
+				backupPath,
+				backupV5
+			}, context, lagReportIntervalMs, syncOptions.textCopy, streamInboundTimeoutMs),
+			destinationBackupURL: backupPath && restoreOptions.litestream?.backupURL ? new URL(backupPath, restoreOptions.litestream.backupURL).toString() : restoreOptions.litestream?.backupURL,
+			replicaID: upstreamReplica.id,
+			waitForBackupBeforeServing: initialSyncedReplica !== void 0 || backupPath !== (restoredReplica?.backupPath ?? null)
 		};
 	} finally {
 		await db.end();
 	}
 }
-async function checkAndUpdateUpstream(lc, sql, shard, { replicaVersion, publications: subscribed, initialSyncContext }) {
-	await updateShardSchema(lc, sql, shard, replicaVersion);
-	const upstreamReplica = await getReplicaAtVersion(lc, sql, shard, replicaVersion, initialSyncContext);
-	if (!upstreamReplica) throw new AutoResetSignal(`No replication slot for replica at version ${replicaVersion}`);
+async function selectAndRestoreReplica(lc, sql, shard, replicaFile, { litestream, constraints }) {
+	const replicas = (await getActiveReplicas(lc, sql, shard)).filter(({ generation }) => generation === (constraints?.replicaVersion ?? generation));
+	if (replicas.length === 0) {
+		lc.info?.(`no suitable replicas to restore from`, { replicas });
+		return;
+	}
+	const [replica] = replicas;
+	if (litestream?.backupURL) {
+		const { backupURL: backupBaseURL } = litestream;
+		const { slot, backupPath, confirmedFlushLsn } = replica;
+		const backupURL = new URL(backupPath ?? "", backupBaseURL).toString();
+		lc.info?.(`restoring replica from ${backupURL} (${slot}@${confirmedFlushLsn})`, { replicas });
+		await restoreReplica(lc, {
+			...litestream,
+			backupURL
+		}, replicaFile, constraints);
+	}
+	return replica;
+}
+async function checkAndUpdateUpstream(lc, sql, shard, { replicaVersion, publications: subscribed, initialSyncContext }, replicaID) {
+	const upstreamReplica = await getReplicaAtVersion(lc, sql, shard, replicaVersion, replicaID, initialSyncContext);
+	if (!upstreamReplica) throw new AutoResetSignal(`No replication slot for replica at version ${replicaVersion} and id ${replicaID}}`);
 	const requested = shard.publications.toSorted();
 	const replicated = upstreamReplica.publications.filter((p) => !p.startsWith(internalPublicationPrefix(shard))).sort();
 	if (!deepEqual(requested, replicated)) {
@@ -91,11 +124,13 @@ var PostgresChangeSource = class {
 	#upstreamUri;
 	#shard;
 	#replica;
+	#backupOptions;
 	#context;
 	#lagReporter;
 	#textCopy;
+	#streamInboundTimeoutMs;
 	#stopped = false;
-	constructor(lc, upstreamUri, shard, replica, context, lagReportIntervalMs, textCopy) {
+	constructor(lc, upstreamUri, shard, replica, backupOptions, context, lagReportIntervalMs, textCopy, streamInboundTimeoutMs) {
 		this.#lc = lc.withContext("component", "change-source");
 		this.#db = pgClient(lc, upstreamUri, "replication-monitor", {
 			max: 1,
@@ -104,8 +139,10 @@ var PostgresChangeSource = class {
 		this.#upstreamUri = upstreamUri;
 		this.#shard = shard;
 		this.#replica = replica;
+		this.#backupOptions = backupOptions;
 		this.#context = context;
 		this.#textCopy = textCopy ?? false;
+		this.#streamInboundTimeoutMs = streamInboundTimeoutMs;
 		this.#lagReporter = lagReportIntervalMs > 0 ? new LagReporter(lc.withContext("component", "lag-reporter"), shard, this.#db, lagReportIntervalMs) : null;
 		registerReplicationSlotHealthMetrics(this.#lc, this.#db, replica.slot, () => this.#stopped);
 	}
@@ -127,7 +164,7 @@ var PostgresChangeSource = class {
 		return null;
 	}
 	async startStream(clientWatermark, backfillRequests = []) {
-		await this.#stopExistingReplicationSlotSubscriber();
+		await this.#takeoverReplicationSlot();
 		const config = await getInternalShardConfig(this.#db, this.#shard);
 		const { slot } = this.#replica;
 		this.#lc.info?.(`starting replication stream@${slot}`);
@@ -135,7 +172,7 @@ var PostgresChangeSource = class {
 	}
 	async #startStream(slot, clientWatermark, shardConfig, backfillRequests) {
 		const clientStart = majorVersionFromString(clientWatermark) + 1n;
-		const { messages, acks } = await subscribe(this.#lc, this.#db, slot, [...shardConfig.publications], clientStart);
+		const { messages, acks } = await subscribe(this.#lc, this.#db, slot, [...shardConfig.publications], clientStart, void 0, void 0, this.#streamInboundTimeoutMs);
 		const acker = new Acker(acks);
 		const changes = new ChangeStreamMultiplexer(this.#lc, clientWatermark);
 		const backfillManager = new BackfillManager(this.#lc, changes, (req) => streamBackfill(this.#lc, this.#upstreamUri, this.#replica, req, { textCopy: this.#textCopy }));
@@ -205,7 +242,7 @@ var PostgresChangeSource = class {
 				changes.fail(err);
 			}
 		})();
-		this.#lc.info?.(`started replication stream@${slot} from ${clientWatermark} (replicaVersion: ${this.#replica.version})`);
+		this.#lc.info?.(`started replication stream@${slot} from ${clientWatermark} (replicaVersion: ${this.#replica.generation})`);
 		return {
 			changes: changes.asSource(),
 			acks: { push: (status) => acker.ack(status[2].watermark) }
@@ -213,17 +250,29 @@ var PostgresChangeSource = class {
 	}
 	async #logCurrentReplicaInfo() {
 		try {
-			const replica = await getReplicaAtVersion(this.#lc, this.#db, this.#shard, this.#replica.version);
-			if (replica) this.#lc.info?.(`Shutdown signal from replica@${this.#replica.version}: ${stringify(replica.subscriberContext)}`);
+			const replica = await getReplicaAtVersion(this.#lc, this.#db, this.#shard, this.#replica.generation, this.#replica.id);
+			if (replica) this.#lc.info?.(`Shutdown signal from replica@${this.#replica.generation}: ${stringify(replica.subscriberContext)}`);
 		} catch (e) {
 			this.#lc.warn?.(`error logging replica info`, e);
 		}
 	}
 	/**
-	* Stops replication slots associated with this shard, and asynchronously
-	* runs a cleanup task that drops older replicas and slots.
+	* In RMv1, a single replication slot is taken over by the next
+	* replication-manager, signaling the old one to shut down.
+	*
+	* With litestream v5, the two replication-managers must replicate to unique
+	* backupPaths, as is required by the LTX backup schema.
+	*
+	* Thus, in addition to taking over the replication slot, the
+	* replication-manager also updates the `replicas` table with its
+	* `backupPath` so that future RM's restore from that path instead
+	* of the vestigial path of the previous replication-manager.
+	*
+	* In RMv2, each replication-manager will have its own slot and
+	* row in the `replicas` table, so this "takeover" will be unnecessary
+	* (but a harmless no-op).
 	*/
-	async #stopExistingReplicationSlotSubscriber() {
+	async #takeoverReplicationSlot() {
 		const sql = this.#db;
 		const { id: replicaID, slot } = this.#replica;
 		const replicasTable = `${upstreamSchema(this.#shard)}.replicas`;
@@ -235,7 +284,7 @@ var PostgresChangeSource = class {
 		if (result.length === 0) {
 			const slotExpression = replicationSlotPrefix(this.#shard);
 			const replicas = await sql`
-        SELECT id, rank, slot, version, "initialSyncContext", "subscriberContext" 
+        SELECT id, rank, slot, generation, "initialSyncContext", "subscriberContext" 
           FROM ${sql(replicasTable)} ORDER BY rank DESC`;
 			const slots = await sql`
         SELECT slot_name as slot, active, active_pid as pid
@@ -251,7 +300,9 @@ var PostgresChangeSource = class {
 		this.#lc.info?.(`terminated replication slots: ${JSON.stringify(result)}`);
 		await sql`
       UPDATE ${sql(replicasTable)} 
-        SET "subscriberContext" = ${this.#context}
+        SET "subscriberContext" = ${this.#context},
+            "backupPath" = ${this.#backupOptions.backupPath},
+            "backupV5" = ${this.#backupOptions.backupV5}
         WHERE id = ${replicaID}`;
 		this.#cleanUpOlderReplicasAndSlots();
 	}
@@ -391,7 +442,10 @@ var LagReporter = class LagReporter {
 			sendTimeMs: now,
 			commitTimeMs
 		});
-		return { nextSendTimeMs: now };
+		return {
+			firstCommitTimeMs: commitTimeMs,
+			nextSendTimeMs: now
+		};
 	}
 	/**
 	* In Postgres < 17, the pg_logical_emit_message lacks an immediate "flush"
@@ -594,7 +648,10 @@ var ChangeMaker = class {
 				}
 			case "commit": return [[
 				"commit",
-				msg,
+				{
+					...msg,
+					commitTimeMs: Number(msg.commitTime / 1000n)
+				},
 				{ watermark: toStateVersionString(must(msg.commitLsn)) }
 			]];
 			case "relation": return await this.#handleRelation(msg);