@rocicorp/zero 1.0.0 → 1.1.0

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

@@ -1,8 +1,8 @@
+import { assert } from "../../../../../shared/src/asserts.js";
 import { deepEqual } from "../../../../../shared/src/json.js";
-import { promiseVoid } from "../../../../../shared/src/resolved-promises.js";
 import { AbortError } from "../../../../../shared/src/abort-error.js";
 import { sleep } from "../../../../../shared/src/sleep.js";
-import { parse } from "../../../../../shared/src/valita.js";
+import { parse, valita_exports } from "../../../../../shared/src/valita.js";
 import { must } from "../../../../../shared/src/must.js";
 import { mapValues } from "../../../../../shared/src/objects.js";
 import { equals, intersection, symmetricDifferences } from "../../../../../shared/src/set-utils.js";
@@ -12,7 +12,7 @@ import { upstreamSchema } from "../../../types/shards.js";
 import { StatementRunner } from "../../../db/statements.js";
 import { pgClient } from "../../../types/pg.js";
 import { majorVersionFromString, majorVersionToString } from "../../../types/state-version.js";
-import { fromBigInt, toStateVersionString } from "./lsn.js";
+import { fromBigInt, toBigInt, toStateVersionString } from "./lsn.js";
 import { UnsupportedColumnDefaultError, mapPostgresToLiteColumn } from "../../../db/pg-to-lite.js";
 import { getSubscriptionStateAndContext } from "../../replicator/schema/replication-state.js";
 import { runTx } from "../../../db/run-transaction.js";
@@ -29,6 +29,7 @@ import { initialSync } from "./initial-sync.js";
 import { streamBackfill } from "./backfill-stream.js";
 import { subscribe } from "./logical-replication/stream.js";
 import postgres from "postgres";
+import { nanoid } from "nanoid";
 import { PG_ADMIN_SHUTDOWN, PG_OBJECT_IN_USE } from "@drdgvhbh/postgres-error-codes";
 //#region ../zero-cache/src/services/change-source/pg/change-source.ts
 /**
@@ -36,7 +37,7 @@ import { PG_ADMIN_SHUTDOWN, PG_OBJECT_IN_USE } from "@drdgvhbh/postgres-error-co
 * replica, before streaming changes from the corresponding logical replication
 * stream.
 */
-async function initializePostgresChangeSource(lc, upstreamURI, shard, replicaDbFile, syncOptions, context) {
+async function initializePostgresChangeSource(lc, upstreamURI, shard, replicaDbFile, syncOptions, context, lagReportIntervalMs) {
 	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));
@@ -45,7 +46,7 @@ async function initializePostgresChangeSource(lc, upstreamURI, shard, replicaDbF
 	try {
 		return {
 			subscriptionState,
-			changeSource: new PostgresChangeSource(lc, upstreamURI, shard, await checkAndUpdateUpstream(lc, db, shard, subscriptionState), context)
+			changeSource: new PostgresChangeSource(lc, upstreamURI, shard, await checkAndUpdateUpstream(lc, db, shard, subscriptionState), context, lagReportIntervalMs ?? null)
 		};
 	} finally {
 		await db.end();
@@ -83,50 +84,69 @@ var MAX_LOW_PRIORITY_DELAY_MS = 1e3;
 */
 var PostgresChangeSource = class {
 	#lc;
+	#db;
 	#upstreamUri;
 	#shard;
 	#replica;
 	#context;
-	constructor(lc, upstreamUri, shard, replica, context) {
+	#lagReporter;
+	constructor(lc, upstreamUri, shard, replica, context, lagReportIntervalMs) {
 		this.#lc = lc.withContext("component", "change-source");
+		this.#db = pgClient(lc, upstreamUri, {
+			["idle_timeout"]: 60,
+			connection: { ["application_name"]: "zero-replication-monitor" }
+		});
 		this.#upstreamUri = upstreamUri;
 		this.#shard = shard;
 		this.#replica = replica;
 		this.#context = context;
+		this.#lagReporter = lagReportIntervalMs ? new LagReporter(lc.withContext("component", "lag-reporter"), shard, this.#db, lagReportIntervalMs) : null;
+	}
+	startLagReporter() {
+		return this.#lagReporter ? this.#lagReporter.initiateLagReport(true) : null;
 	}
 	async startStream(clientWatermark, backfillRequests = []) {
-		const db = pgClient(this.#lc, this.#upstreamUri);
 		const { slot } = this.#replica;
-		let cleanup = promiseVoid;
-		try {
-			({cleanup} = await this.#stopExistingReplicationSlotSubscribers(db, slot));
-			const config = await getInternalShardConfig(db, this.#shard);
-			this.#lc.info?.(`starting replication stream@${slot}`);
-			return await this.#startStream(db, slot, clientWatermark, config, backfillRequests);
-		} finally {
-			cleanup.then(() => db.end());
-		}
+		await this.#stopExistingReplicationSlotSubscribers(slot);
+		const config = await getInternalShardConfig(this.#db, this.#shard);
+		this.#lc.info?.(`starting replication stream@${slot}`);
+		return this.#startStream(slot, clientWatermark, config, backfillRequests);
 	}
-	async #startStream(db, slot, clientWatermark, shardConfig, backfillRequests) {
+	async #startStream(slot, clientWatermark, shardConfig, backfillRequests) {
 		const clientStart = majorVersionFromString(clientWatermark) + 1n;
-		const { messages, acks } = await subscribe(this.#lc, db, slot, [...shardConfig.publications], clientStart);
+		const { messages, acks } = await subscribe(this.#lc, this.#db, slot, [...shardConfig.publications], clientStart);
 		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));
 		changes.addProducers(messages, backfillManager).addListeners(backfillManager, acker);
 		backfillManager.run(clientWatermark, backfillRequests);
-		const changeMaker = new ChangeMaker(this.#lc, this.#shard, shardConfig, this.#replica.initialSchema, this.#upstreamUri);
+		const changeMaker = new ChangeMaker(this.#lc, this.#shard, shardConfig, this.#db, this.#replica.initialSchema);
+		/**
+		* Determines if the incoming message is transactional, otherwise handling
+		* non-transactional messages with a downstream status message.
+		*/
+		const isTransactionalMessage = (lsn, msg) => {
+			if (msg.tag === "message" && msg.prefix === this.#lagReporter?.messagePrefix) {
+				changes.pushStatus(this.#lagReporter.processLagReport(msg));
+				return false;
+			}
+			this.#lagReporter?.checkCurrentLSN(lsn);
+			if (msg.tag === "keepalive") {
+				changes.pushStatus([
+					"status",
+					{ ack: msg.shouldRespond },
+					{ watermark: majorVersionToString(lsn) }
+				]);
+				return false;
+			}
+			return true;
+		};
 		(async () => {
 			try {
 				let reservation = null;
 				let inTransaction = false;
 				for await (const [lsn, msg] of messages) {
-					if (msg.tag === "keepalive") {
-						changes.pushStatus([
-							"status",
-							{ ack: msg.shouldRespond },
-							{ watermark: majorVersionToString(lsn) }
-						]);
+					if (!isTransactionalMessage(lsn, msg)) {
 						if (!inTransaction && reservation?.lastWatermark) {
 							changes.release(reservation.lastWatermark);
 							reservation = null;
@@ -170,14 +190,11 @@ var PostgresChangeSource = class {
 		};
 	}
 	async #logCurrentReplicaInfo() {
-		const db = pgClient(this.#lc, this.#upstreamUri);
 		try {
-			const replica = await getReplicaAtVersion(this.#lc, db, this.#shard, this.#replica.version);
+			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)}`);
 		} catch (e) {
 			this.#lc.warn?.(`error logging replica info`, e);
-		} finally {
-			await db.end();
 		}
 	}
 	/**
@@ -189,10 +206,10 @@ var PostgresChangeSource = class {
 	* the timestamp suffix) are preserved, as those are newly syncing replicas
 	* that will soon take over the slot.
 	*/
-	async #stopExistingReplicationSlotSubscribers(db, slotToKeep) {
+	async #stopExistingReplicationSlotSubscribers(slotToKeep) {
 		const slotExpression = replicationSlotExpression(this.#shard);
 		const legacySlotName = legacyReplicationSlot(this.#shard);
-		const result = await runTx(db, async (sql) => {
+		const result = await runTx(this.#db, async (sql) => {
 			const result = await sql`
       SELECT slot_name as slot, pg_terminate_backend(active_pid) as terminated, active_pid as pid
         FROM pg_replication_slots 
@@ -230,10 +247,11 @@ var PostgresChangeSource = class {
 		const pids = result.filter(({ pid }) => pid !== null).map(({ pid }) => pid);
 		if (pids.length) this.#lc.info?.(`signaled subscriber ${pids} to shut down`);
 		const otherSlots = result.filter(({ slot }) => slot !== slotToKeep).map(({ slot }) => slot);
-		return { cleanup: otherSlots.length ? this.#dropReplicationSlots(db, otherSlots) : promiseVoid };
+		if (otherSlots.length) this.#dropReplicationSlots(otherSlots).catch((e) => this.#lc.warn?.(`error dropping replication slots`, e));
 	}
-	async #dropReplicationSlots(sql, slots) {
+	async #dropReplicationSlots(slots) {
 		this.#lc.info?.(`dropping other replication slot(s) ${slots}`);
+		const sql = this.#db;
 		for (let i = 0; i < 5; i++) try {
 			await sql`
           SELECT pg_drop_replication_slot(slot_name) FROM pg_replication_slots
@@ -282,24 +300,132 @@ var Acker = class {
 		this.#acks.push(lsn);
 	}
 };
+var lagReportSchema = valita_exports.object({
+	id: valita_exports.string(),
+	sendTimeMs: valita_exports.number(),
+	commitTimeMs: valita_exports.number()
+});
+var LagReporter = class LagReporter {
+	static MESSAGE_SUFFIX = "/lag-report/v1";
+	#lc;
+	messagePrefix;
+	#db;
+	#lagIntervalMs;
+	#pgVersion;
+	#expectingLagReport = null;
+	#timer;
+	constructor(lc, shard, db, lagIntervalMs) {
+		this.#lc = lc;
+		this.messagePrefix = `${shard.appID}/${shard.shardNum}${LagReporter.MESSAGE_SUFFIX}`;
+		this.#db = db;
+		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;
+	}
+	async initiateLagReport(log = false) {
+		const pgVersion = this.#pgVersion ?? await this.#getPgVersion();
+		const now = Date.now();
+		const id = nanoid();
+		const lagReport = {
+			id,
+			lsn: 0n
+		};
+		this.#expectingLagReport = lagReport;
+		let lsn;
+		if (pgVersion >= 17e4) [{lsn}] = await this.#db`
+        SELECT pg_logical_emit_message(
+          false,
+          ${this.messagePrefix},
+          json_build_object(
+            'id', ${id}::text,
+            'sendTimeMs', ${now}::int8,
+            'commitTimeMs', extract(epoch from now()) * 1000
+          )::text,
+          true
+        ) as lsn;
+    `;
+		else [{lsn}] = await this.#db`
+        SELECT pg_logical_emit_message(
+          false,
+          ${this.messagePrefix},
+          json_build_object(
+            'id', ${id}::text,
+            'sendTimeMs', ${now}::int8,
+            'commitTimeMs', extract(epoch from now()) * 1000
+          )::text
+        ) as lsn;
+    `;
+		lagReport.lsn = toBigInt(lsn);
+		if (log) this.#lc.info?.(`initiated lag report at lsn ${lsn}`, {
+			id,
+			lsn
+		});
+		return { nextSendTimeMs: now };
+	}
+	checkCurrentLSN(lsn) {
+		if (this.#expectingLagReport?.lsn && lsn > this.#expectingLagReport.lsn) {
+			this.#lc.warn?.(`LSN ${fromBigInt(lsn)} is passed expected lag report ${fromBigInt(this.#expectingLagReport.lsn)}. Initiating new report.`);
+			this.#scheduleNextReport(0);
+		}
+	}
+	#scheduleNextReport(delayMs) {
+		this.#expectingLagReport = null;
+		clearTimeout(this.#timer);
+		this.#timer = setTimeout(async () => {
+			try {
+				await this.initiateLagReport();
+			} catch (e) {
+				this.#lc.warn?.(`error initiating lag report`, e);
+				this.#scheduleNextReport(this.#lagIntervalMs);
+			}
+		}, delayMs);
+	}
+	processLagReport(msg) {
+		assert(msg.prefix === this.messagePrefix, `unexpected message prefix: ${msg.prefix}`);
+		const report = parseLogicalMessageContent(msg, lagReportSchema);
+		const now = Date.now();
+		const nextSendTimeMs = Math.max(now, report.sendTimeMs + this.#lagIntervalMs);
+		if (report.id === this.#expectingLagReport?.id) this.#scheduleNextReport(nextSendTimeMs - now);
+		else this.#lc.debug?.(`received extraneous lag report`, { report });
+		const { sendTimeMs, commitTimeMs } = report;
+		return [
+			"status",
+			{
+				ack: false,
+				lagReport: {
+					lastTimings: {
+						sendTimeMs,
+						commitTimeMs,
+						receiveTimeMs: now
+					},
+					nextSendTimeMs
+				}
+			},
+			{ watermark: toStateVersionString(msg.messageLsn ?? "0/0") }
+		];
+	}
+};
 var SET_REPLICA_IDENTITY_DELAY_MS = 50;
 var ChangeMaker = class {
 	#lc;
 	#shardPrefix;
 	#shardConfig;
 	#initialSchema;
-	#upstreamDB;
+	#db;
 	#replicaIdentityTimer;
 	#error;
-	constructor(lc, { appID, shardNum }, shardConfig, initialSchema, upstreamURI) {
+	constructor(lc, { appID, shardNum }, shardConfig, db, initialSchema) {
 		this.#lc = lc;
 		this.#shardPrefix = `${appID}/${shardNum}`;
 		this.#shardConfig = shardConfig;
 		this.#initialSchema = initialSchema;
-		this.#upstreamDB = pgClient(lc, upstreamURI, {
-			["idle_timeout"]: 10,
-			connection: { ["application_name"]: "zero-schema-change-detector" }
-		});
+		this.#db = db;
 	}
 	async makeChanges(lsn, msg) {
 		if (this.#error) {
@@ -398,7 +524,7 @@ var ChangeMaker = class {
 	#preSchema;
 	#lastSnapshotInTx;
 	#handleDdlMessage(msg) {
-		const event = this.#parseReplicationEvent(msg.content);
+		const event = parseLogicalMessageContent(msg, replicationEventSchema);
 		clearTimeout(this.#replicaIdentityTimer);
 		let previousSchema;
 		const { type } = event;
@@ -427,7 +553,7 @@ var ChangeMaker = class {
 		const replicaIdentities = replicaIdentitiesForTablesWithoutPrimaryKeys(event.schema);
 		if (replicaIdentities) this.#replicaIdentityTimer = setTimeout(async () => {
 			try {
-				await replicaIdentities.apply(this.#lc, this.#upstreamDB);
+				await replicaIdentities.apply(this.#lc, this.#db);
 			} catch (err) {
 				this.#lc.warn?.(`error setting replica identities`, err);
 			}
@@ -602,10 +728,6 @@ var ChangeMaker = class {
 		}
 		return changes;
 	}
-	#parseReplicationEvent(content) {
-		const str = content instanceof Buffer ? content.toString("utf-8") : new TextDecoder().decode(content);
-		return parse(JSON.parse(str), replicationEventSchema, "passthrough");
-	}
 	/**
 	* If `ddlDetection === true`, relation messages are irrelevant,
 	* as schema changes are detected by event triggers that
@@ -627,7 +749,7 @@ var ChangeMaker = class {
 	async #handleRelation(rel) {
 		const { publications, ddlDetection } = this.#shardConfig;
 		if (ddlDetection) return [];
-		const currentSchema = await getPublicationInfo(this.#upstreamDB, publications);
+		const currentSchema = await getPublicationInfo(this.#db, publications);
 		const difference = getSchemaDifference(this.#initialSchema, currentSchema);
 		if (difference !== null) throw new MissingEventTriggerSupport(difference);
 		const orel = this.#initialSchema.tables.find((t) => t.oid === rel.relationOid);
@@ -752,6 +874,10 @@ var ShutdownSignal = class extends AbortError {
 		super("shutdown signal received (e.g. another zero-cache taking over the replication stream)", { cause });
 	}
 };
+function parseLogicalMessageContent({ content }, schema) {
+	const str = content instanceof Buffer ? content.toString("utf-8") : new TextDecoder().decode(content);
+	return parse(JSON.parse(str), schema, "passthrough");
+}
 //#endregion
 export { initializePostgresChangeSource };