@rocicorp/zero 1.11.0-canary.7 → 1.11.0-canary.9

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

@@ -7,113 +7,28 @@ import { mapValues } from "../../../../../shared/src/objects.js";
 import { equals, intersection, symmetricDifferences } from "../../../../../shared/src/set-utils.js";
 import { stringify } from "../../../../../shared/src/bigint-json.js";
 import { upstreamSchema } from "../../../types/shards.js";
-import { Database } from "../../../../../zqlite/src/db.js";
 import { UnsupportedColumnDefaultError, mapPostgresToLiteColumn } from "../../../db/pg-to-lite.js";
 import { getOrCreateCounter } from "../../../observability/metrics.js";
-import { getSubscriptionStateAndContext } from "../../replicator/schema/replication-state.js";
-import { StatementRunner } from "../../../db/statements.js";
-import { connectPgClient, isPostgresError, pgClient } from "../../../types/pg.js";
+import { isPostgresError, pgClient } from "../../../types/pg.js";
 import { majorVersionFromString, majorVersionToString } from "../../../types/state-version.js";
 import { fromBigInt, toBigInt, toStateVersionString } from "./lsn.js";
 import { getPublicationInfo, warnForSkippedIndexes } from "./schema/published.js";
 import { replicationEventSchema } from "./schema/ddl.js";
 import { validate } from "./schema/validation.js";
-import { dropShard, getActiveReplicas, getInternalShardConfig, getReplicaAtVersion, internalPublicationPrefix, replicaIdentitiesForTablesWithoutPrimaryKeys, replicationSlotPrefix } from "./schema/shard.js";
-import { AutoResetSignal } from "../../change-streamer/schema/tables.js";
+import { getInternalShardConfig, getReplicaAtVersion, replicaIdentitiesForTablesWithoutPrimaryKeys, replicationSlotPrefix } from "./schema/shard.js";
 import { mapIndexPredicateColumns } from "../../../db/index-predicate.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 { subscribe } from "./logical-replication/stream.js";
-import { dropOldReplicasAndSlots } from "./replication-slots.js";
-import { ensureShardSchema } from "./schema/init.js";
-import { initialSync } from "./initial-sync.js";
+import "./initial-sync.js";
 import { streamBackfill } from "./backfill-stream.js";
+import { ReplicationSlotCleanupMonitor } from "./replication-slot-cleanup-monitor.js";
 import { registerReplicationSlotHealthMetrics } from "./replication-slot-health.js";
 import { nanoid } from "nanoid";
 import postgres from "postgres";
 import { PG_ADMIN_SHUTDOWN, PG_INSUFFICIENT_PRIVILEGE } from "@drdgvhbh/postgres-error-codes";
 //#region ../zero-cache/src/services/change-source/pg/change-source.ts
-var REPLICA_SLOT_CLEANUP_INTERVAL_MS = 3e4;
-/**
-* Initializes a Postgres change source, including the initial sync of the
-* replica, before streaming changes from the corresponding logical replication
-* stream.
-*/
-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, syncOptions.installPartialIndexTriggers);
-		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, 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 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)) {
-		lc.warn?.(`Dropping shard to change publications to: [${requested}]`);
-		await sql.unsafe(dropShard(shard.appID, shard.shardNum));
-		throw new AutoResetSignal(`Requested publications [${requested}] do not match configured publications: [${replicated}]`);
-	}
-	if (!deepEqual(upstreamReplica.publications, subscribed)) throw new AutoResetSignal(`Upstream publications [${upstreamReplica.publications}] do not match subscribed publications [${subscribed}]`);
-	const exists = await sql`
-    SELECT pubname FROM pg_publication WHERE pubname IN ${sql(subscribed)};
-  `.values();
-	if (exists.length !== subscribed.length) throw new AutoResetSignal(`Upstream publications [${exists.flat()}] do not contain all subscribed publications [${subscribed}]`);
-	const { slot } = upstreamReplica;
-	const result = await sql`
-    SELECT restart_lsn as "restartLSN", wal_status as "walStatus" FROM pg_replication_slots
-      WHERE slot_name = ${slot}`;
-	if (result.length === 0) throw new AutoResetSignal(`replication slot ${slot} is missing`);
-	const [{ restartLSN, walStatus }] = result;
-	if (restartLSN === null || walStatus === "lost") throw new AutoResetSignal(`replication slot ${slot} has been invalidated for exceeding the max_slot_wal_keep_size`);
-	return upstreamReplica;
-}
 var MAX_LOW_PRIORITY_DELAY_MS = 1e3;
 /**
 * Postgres implementation of a {@link ChangeSource} backed by a logical
@@ -125,6 +40,7 @@ var PostgresChangeSource = class {
 	#upstreamUri;
 	#shard;
 	#replica;
+	#slotCleanupMonitor;
 	#backupOptions;
 	#context;
 	#lagReporter;
@@ -133,28 +49,28 @@ var PostgresChangeSource = class {
 	#subscribe;
 	#streamBackfill;
 	#stopped = false;
-	constructor(lc, upstreamUri, shard, replica, backupOptions, context, lagReportIntervalMs, textCopy, streamInboundTimeoutMs, deps = {}) {
+	constructor(lc, upstreamUri, shard, replica, pgVersion, backupOptions, context, lagReportIntervalMs, textCopy, streamInboundTimeoutMs, deps = {}) {
 		this.#lc = lc.withContext("component", "change-source");
 		this.#subscribe = deps.subscribe ?? subscribe;
 		this.#streamBackfill = deps.streamBackfill ?? streamBackfill;
 		this.#db = pgClient(lc, upstreamUri, "replication-monitor", {
-			max: 1,
-			["idle_timeout"]: 60
+			max: 3,
+			idle_timeout: 60
 		});
 		this.#upstreamUri = upstreamUri;
 		this.#shard = shard;
 		this.#replica = replica;
+		this.#slotCleanupMonitor = new ReplicationSlotCleanupMonitor(lc.withContext("component", "replication-slot-monitor"), shard, this.#db, pgVersion, replica.slot);
 		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;
+		this.#lagReporter = lagReportIntervalMs > 0 ? new LagReporter(lc.withContext("component", "lag-reporter"), shard, this.#db, pgVersion, lagReportIntervalMs) : null;
 		registerReplicationSlotHealthMetrics(this.#lc, this.#db, replica.slot, () => this.#stopped);
 	}
 	async stop() {
 		this.#stopped = true;
 		this.#lagReporter?.stop();
-		clearTimeout(this.#cleanupTimer);
 		await this.#db.end();
 	}
 	async startLagReporter() {
@@ -173,12 +89,18 @@ var PostgresChangeSource = class {
 		const config = await getInternalShardConfig(this.#db, this.#shard);
 		const { slot } = this.#replica;
 		this.#lc.info?.(`starting replication stream@${slot}`);
-		return this.startStreamInternal(slot, clientWatermark, config, backfillRequests);
+		const changeStream = await this.startStreamInternal(slot, clientWatermark, config, backfillRequests);
+		const { signal } = changeStream.changes;
+		if (!signal.aborted) {
+			this.#slotCleanupMonitor.start();
+			signal.addEventListener("abort", () => this.#slotCleanupMonitor.stop());
+		}
+		return changeStream;
 	}
 	async startStreamInternal(slot, clientWatermark, shardConfig, backfillRequests) {
 		const clientStart = majorVersionFromString(clientWatermark) + 1n;
 		const { messages, acks } = await this.#subscribe(this.#lc, this.#db, slot, [...shardConfig.publications], clientStart, void 0, void 0, this.#streamInboundTimeoutMs);
-		const acker = new Acker(acks);
+		const acker = new Acker(acks, clientWatermark);
 		const changes = new ChangeStreamMultiplexer(this.#lc, clientWatermark);
 		const backfillManager = new BackfillManager(this.#lc, changes, (req) => this.#streamBackfill(this.#lc, this.#upstreamUri, this.#replica, req, { textCopy: this.#textCopy }));
 		changes.addProducers(messages, backfillManager).addListeners(backfillManager, acker);
@@ -312,29 +234,22 @@ var PostgresChangeSource = class {
             "backupPath" = ${this.#backupOptions.backupPath},
             "backupV5" = ${this.#backupOptions.backupV5}
         WHERE id = ${replicaID}`;
-		this.#cleanUpOlderReplicasAndSlots();
-	}
-	#cleanupTimer;
-	async #cleanUpOlderReplicasAndSlots() {
-		clearTimeout(this.#cleanupTimer);
-		try {
-			const result = await dropOldReplicasAndSlots(this.#lc, this.#db, this.#shard, this.#replica.rank);
-			if (result.draining === 0) {
-				this.#lc.info?.(`finished cleaning up replicas and slots`, { result });
-				return;
-			}
-			this.#lc.info?.(`old slots still draining`, { result });
-		} catch (e) {
-			this.#lc.warn?.(`error dropping replication slots`, e);
-		}
-		this.#cleanupTimer = setTimeout(() => this.#cleanUpOlderReplicasAndSlots(), REPLICA_SLOT_CLEANUP_INTERVAL_MS);
 	}
 };
 var Acker = class {
 	#acks;
-	#waitingForDownstreamAck = null;
-	constructor(acks) {
+	#waitingForDownstreamAck;
+	/**
+	* @param resumeWatermark the watermark the stream resumes after. What came
+	*     before it was received on an earlier connection, and only the
+	*     change-streamer knows whether it has persisted it, so keepalives are
+	*     not acked until the change-streamer has acked it. A change-streamer
+	*     whose SQLite change log is ahead of its backup would otherwise have
+	*     the slot moved past transactions that only that log holds.
+	*/
+	constructor(acks, resumeWatermark) {
 		this.#acks = acks;
+		this.#waitingForDownstreamAck = resumeWatermark;
 	}
 	onChange(change) {
 		switch (change[0]) {
@@ -381,25 +296,17 @@ var LagReporter = class LagReporter {
 	#pgVersion;
 	#expectingLagReport = null;
 	#timer;
-	constructor(lc, shard, db, lagIntervalMs) {
+	constructor(lc, shard, db, pgVersion, lagIntervalMs) {
 		this.#lc = lc;
 		this.messagePrefix = `${shard.appID}/${shard.shardNum}${LagReporter.MESSAGE_SUFFIX}`;
 		this.#db = db;
+		this.#pgVersion = pgVersion;
 		this.#lagIntervalMs = lagIntervalMs;
 	}
-	async #getPgVersion() {
-		if (this.#pgVersion === void 0) {
-			const [{ pgVersion }] = await this.#db`
-        SELECT current_setting('server_version_num')::int as "pgVersion"`;
-			this.#pgVersion = pgVersion;
-		}
-		return this.#pgVersion;
-	}
 	get pgVersion() {
 		return this.#pgVersion;
 	}
 	async initiateLagReport(log = false) {
-		const pgVersion = this.#pgVersion ?? await this.#getPgVersion();
 		const now = Date.now();
 		const id = nanoid();
 		const lagReport = {
@@ -412,7 +319,7 @@ var LagReporter = class LagReporter {
 		let commitTimeMs;
 		let lsn;
 		try {
-			if (pgVersion >= 17e4) [{commitTimeMs, lsn}] = await this.#db`
+			if (this.#pgVersion >= 17e4) [{commitTimeMs, lsn}] = await this.#db`
           WITH CTE AS (SELECT extract(epoch from now()) * 1000 AS "commitTimeMs")
           SELECT "commitTimeMs", pg_logical_emit_message(
             false,
@@ -1071,6 +978,6 @@ function parseLogicalMessageContent(lc, msg, schema) {
 	}
 }
 //#endregion
-export { initializePostgresChangeSource };
+export { PostgresChangeSource };
 
 //# sourceMappingURL=change-source.js.map