@rocicorp/zero 1.11.0-canary.3 → 1.11.0-canary.30

package/out/zero-cache/src/services/change-streamer/storer.js

@@ -2,6 +2,7 @@ import { assert, unreachable } from "../../../../shared/src/asserts.js";
 import { promiseVoid } from "../../../../shared/src/resolved-promises.js";
 import { AbortError } from "../../../../shared/src/abort-error.js";
 import { parse, valita_exports } from "../../../../shared/src/valita.js";
+import { must } from "../../../../shared/src/must.js";
 import { BigIntJSON } from "../../../../shared/src/bigint-json.js";
 import { cdcSchema } from "../../types/shards.js";
 import { isDataChange, isSchemaChange } from "../change-source/protocol/current/data.js";
@@ -13,10 +14,10 @@ import "../../types/pg.js";
 import { READONLY, READ_COMMITTED } from "../../db/mode-enum.js";
 import { runTx } from "../../db/run-transaction.js";
 import { Queue } from "../../../../shared/src/queue.js";
-import { AutoResetSignal, markResetRequired } from "./schema/tables.js";
+import { AutoResetSignal, markResetRequired, reinitializeChangeLog } from "./schema/tables.js";
 import { TransactionPool, sharedSnapshot } from "../../db/transaction-pool.js";
 import { orTimeout } from "../../types/timeout.js";
-import { extractChangeSubstring, reconstructWatermarkedChange, serializeChangeStreamData } from "./change-log-codec.js";
+import { reconstructWatermarkedChange, serializeChangeStreamDataWithChange } from "./change-log-codec.js";
 import { resolver } from "@rocicorp/resolver";
 import "postgres";
 import { getHeapStatistics } from "node:v8";
@@ -66,6 +67,11 @@ var Storer = class {
 	#replicaVersion;
 	#onCommitted;
 	#onFatal;
+	/**
+	* Set while waiting to assume ownership of the change log at its head
+	* (see {@link assumeOwnershipAtHead}).
+	*/
+	#pendingOwnership = null;
 	#queue = new Queue();
 	#backPressureThresholdBytes;
 	#statementTimeoutMs;
@@ -120,12 +126,128 @@ var Storer = class {
 		}
 		return result;
 	}
-	async assumeOwnership(purgeLock) {
-		const db = this.#db;
-		const owner = this.#taskID;
+	#ownerAddress() {
 		const ownerAddress = this.#discoveryAddress;
 		const ownerProtocol = this.#discoveryProtocol;
-		const addressWithProtocol = ownerProtocol === "ws" ? ownerAddress : `${ownerProtocol}://${ownerAddress}`;
+		return ownerProtocol === "ws" ? ownerAddress : `${ownerProtocol}://${ownerAddress}`;
+	}
+	/**
+	* Re-initializes a change log that is no longer maintained from the
+	* replica's state, taking ownership of it in the process (which is then
+	* assumed again at its new head, see {@link assumeOwnershipAtHead}).
+	* See {@link reinitializeChangeLog}.
+	*/
+	reinitializeFromReplica(params) {
+		return reinitializeChangeLog(this.#lc, this.#db, this.#shard, {
+			...params,
+			owner: this.#taskID,
+			ownerAddress: this.#ownerAddress()
+		}, this.#pendingOwnership?.purgeLock ?? void 0);
+	}
+	/**
+	* Prepares to assume ownership of the change log at its head, while it may
+	* still be owned and maintained by another change-streamer. This allows the
+	* previous owner to continue to serve subscribers while this change-source's
+	* upstream decodes the WAL to catch up to the head of the change log, which
+	* may be many minutes worth of transactions for large DBs.
+	*
+	* Once the change stream is started from the head of the change log:
+	*
+	* * Transactions that the change log already contains (i.e. stored by the
+	*   previous owner) are not stored again, but are otherwise processed
+	*   normally (e.g. for upstream acks, since they have been stored).
+	* * Ownership is assumed in the transaction that stores the first
+	*   transaction past the head of the change log. The previous owner's
+	*   (subsequent) transactions fail its ownership check, which makes this
+	*   transaction the one immediately after the previous owner's last one.
+	* * Ownership is also assumed upon a {@link status} at or past the head of
+	*   the change log (outside of a transaction), which covers an idle
+	*   upstream: all transactions up to the status position have been
+	*   received, and any past the head would have been stored.
+	*
+	* Subscribers are caught up from the change log as usual, as the stream
+	* (which is forwarded to them) starts from its head.
+	*
+	* Backfill transactions are synthesized by each change-streamer's change
+	* source, and thus never correspond to those stored by the previous owner.
+	* Until ownership is assumed, they are not stored, and no backfills are
+	* requested from the change log (see
+	* {@link getStartStreamInitializationParameters}). Once ownership is
+	* assumed, a restart of the change stream is requested if the change log has
+	* pending backfills or if a backfill transaction was skipped.
+	*/
+	assumeOwnershipAtHead(takeover) {
+		this.#pendingOwnership = {
+			...takeover,
+			head: "",
+			skippedBackfill: false,
+			start: performance.now()
+		};
+	}
+	/**
+	* Whether ownership of the change log is pending, i.e. after
+	* {@link assumeOwnershipAtHead} and before ownership is assumed.
+	*/
+	get ownershipPending() {
+		return this.#pendingOwnership !== null;
+	}
+	/**
+	* Relays a status message received outside of a transaction, which is
+	* used to assume ownership of the change log on an idle upstream. A no-op
+	* unless ownership is pending.
+	*/
+	status(watermark) {
+		if (this.#pendingOwnership) this.#queue.enqueue(["status", watermark]);
+	}
+	/**
+	* Assumes ownership if the change log contains nothing past `watermark`.
+	* Otherwise, updates the known head of the change log.
+	*
+	* @param ifBefore When set, ownership is only assumed if the head of the
+	*     change log is before (rather than at or before) the `watermark`,
+	*     i.e. when the watermark is that of a transaction to be stored.
+	* @returns The state of the change log at which ownership is assumed (in
+	*     `tx`), or `null` if ownership was not assumed.
+	*/
+	async #tryAssumeOwnership(tx, watermark, ifBefore) {
+		const pending = must(this.#pendingOwnership);
+		const [state] = await tx.processReadTask((sql) => sql`
+        SELECT "lastWatermark" FROM ${this.#cdc("replicationState")} FOR UPDATE`);
+		const head = state.lastWatermark;
+		pending.head = head > pending.head ? head : pending.head;
+		if (ifBefore ? head >= watermark : head > watermark) return null;
+		const owner = this.#taskID;
+		const ownerAddress = this.#ownerAddress();
+		tx.process((sql) => [sql`UPDATE ${this.#cdc("replicationState")} SET ${sql({
+			owner,
+			ownerAddress
+		})}`]);
+		const [{ backfilling }] = await tx.processReadTask((sql) => sql`
+        SELECT EXISTS (SELECT 1 FROM ${this.#cdc("backfilling")}) AS "backfilling"`);
+		this.#lc.info?.(`assuming ownership at ${ownerAddress} after change log head ${head}`);
+		return {
+			head,
+			backfilling
+		};
+	}
+	#ownershipAssumed({ head, backfilling }) {
+		const { purgeLock, requestRestart, skippedBackfill, start } = must(this.#pendingOwnership);
+		this.#pendingOwnership = null;
+		this.#lc.info?.(`assumed ownership of the change log at ${head} after ${(performance.now() - start).toFixed(2)} ms`);
+		purgeLock?.release();
+		if (backfilling || skippedBackfill) requestRestart(`resuming backfills after assuming ownership`);
+	}
+	/**
+	* Assumes ownership of the change log immediately, regardless of its head.
+	* This is used for in tests.
+	*
+	* A production change-streamer instead takes over the change log at its head
+	* (see {@link assumeOwnershipAtHead}).
+	*/
+	async assumeOwnership() {
+		const db = this.#db;
+		const owner = this.#taskID;
+		const addressWithProtocol = this.#ownerAddress();
 		this.#lc.info?.(`assuming ownership at ${addressWithProtocol}`);
 		const start = performance.now();
 		await this.#withTimeout("assume-ownership", db`UPDATE ${this.#cdc("replicationState")} SET ${db({
@@ -134,7 +256,6 @@ var Storer = class {
 		})}`);
 		const elapsed = (performance.now() - start).toFixed(2);
 		this.#lc.info?.(`assumed ownership at ${addressWithProtocol} (${elapsed} ms)`);
-		if (purgeLock) purgeLock.release();
 	}
 	/**
 	* The whole of what a stream connection starts from: where to resume, what to
@@ -164,7 +285,7 @@ var Storer = class {
               'name', b."table",
               'metadata', t."metadata"
             ) as "table",
-            json_object_agg(b."column", b."backfill")
+            json_object_agg(b."column", json_build_object('id', b."backfill"))
               as "columns"
           FROM ${this.#cdc("backfilling")} as b
           LEFT JOIN ${this.#cdc("tableMetadata")} as t
@@ -180,7 +301,7 @@ var Storer = class {
 		], { mode: READONLY }));
 		return {
 			lastWatermark,
-			backfillRequests: parse(result, backfillRequestsSchema),
+			backfillRequests: this.#pendingOwnership ? [] : parse(result, backfillRequestsSchema),
 			cookies: {
 				tableMetadata: [...tableMetadata],
 				backfilling: [...backfilling]
@@ -263,14 +384,15 @@ var Storer = class {
 	/**
 	* @returns The JSON stringified stream message to be sent downstream.
 	*/
-	store(watermark, data) {
-		const json = serializeChangeStreamData(data);
+	store(watermark, data, serialized = serializeChangeStreamDataWithChange(data)) {
+		const { json, change: storedChange } = serialized;
 		this.#approximateQueuedBytes += json.length;
 		const change = data[1];
 		this.#queue.enqueue([
 			"change",
 			watermark,
 			json,
+			storedChange,
 			isDataChange(change) ? null : change
 		]);
 		return json;
@@ -403,88 +525,152 @@ var Storer = class {
 						else await this.#startCatchup([subscriber]);
 						continue;
 					}
+					case "status": {
+						const pending = this.#pendingOwnership;
+						if (pending && !tx && msg[1] >= pending.head) {
+							const pool = new TransactionPool(this.#lc, { mode: READ_COMMITTED });
+							pool.run(this.#inserter);
+							const assumed = await this.#tryAssumeOwnership(pool, msg[1], false);
+							if (assumed === null) {
+								pool.abort();
+								await pool.done();
+							} else {
+								pool.setDone();
+								await pool.done();
+								this.#ownershipAssumed(assumed);
+							}
+						}
+						continue;
+					}
 					case "abort":
-						if (tx) {
+						if (tx && !tx.skipped) {
 							tx.pool.abort();
 							await tx.pool.done();
-							tx = null;
 						}
+						tx = null;
 						continue;
 				}
-				const [_, watermark, json, change] = msg;
-				const tag = change?.tag;
-				this.#approximateQueuedBytes -= json.length;
-				if (tag === "begin") {
-					assert(!tx, "received BEGIN in the middle of a transaction");
-					const { promise, resolve, reject } = resolver();
-					promise.catch(() => {});
-					tx = {
-						pool: new TransactionPool(this.#lc.withContext("watermark", watermark), { mode: READ_COMMITTED }),
-						preCommitWatermark: watermark,
-						pos: 0,
-						startingReplicationState: promise,
-						ack: !change.skipAck,
-						batch: [],
-						lastFlush: void 0
-					};
-					tx.pool.run(this.#inserter);
-					tx.pool.process((tx) => {
-						tx`
-          SELECT "owner" FROM ${this.#cdc("replicationState")} FOR UPDATE`.then(([result]) => resolve(result), reject);
-						return [];
-					});
-				} else {
-					assert(tx, () => `received change outside of transaction: ${json}`);
-					tx.pos++;
-				}
-				const entry = {
-					watermark: tag === "commit" ? watermark : tx.preCommitWatermark,
-					precommit: tag === "commit" ? tx.preCommitWatermark : null,
-					pos: tx.pos,
-					change: extractChangeSubstring(json, tag)
-				};
-				if (change !== null && isSchemaChange(change)) {
-					await this.#flushChangeLog(tx);
-					tx.lastFlush = tx.pool.process((sql) => [sql`INSERT INTO ${this.#cdc("changeLog")} ${sql(entry)}`, ...this.#trackBackfillMetadata(sql, change)]);
-				} else {
-					tx.batch.push(entry);
-					if (tx.batch.length >= this.#changeLogBatchSize) {
-						const prevFlush = tx.lastFlush;
-						this.#flushChangeLog(tx);
-						await prevFlush;
-					}
-				}
+				this.#approximateQueuedBytes -= msg[2].length;
+				tx = await this.#processChange(tx, msg);
 				this.#maybeReleaseBackPressure();
-				if (tag === "commit") {
-					this.#flushChangeLog(tx);
-					const { owner } = await tx.startingReplicationState;
-					if (owner !== this.#taskID) tx.pool.fail(new AbortError(`changeLog ownership has been assumed by ${owner}`));
-					else {
-						const lastWatermark = watermark;
-						tx.pool.process((tx) => [tx`
-            UPDATE ${this.#cdc("replicationState")} SET ${tx({ lastWatermark })}`]);
-						tx.pool.setDone();
-					}
-					await tx.pool.done();
-					if (tx.ack) this.#onCommitted([
-						"commit",
-						change,
-						{ watermark }
-					]);
-					tx = null;
-					await this.#startCatchup(catchupQueue.splice(0));
-				} else if (tag === "rollback") {
-					tx.pool.abort();
-					await tx.pool.done();
-					tx = null;
-					await this.#startCatchup(catchupQueue.splice(0));
-				}
+				if (tx === null) await this.#startCatchup(catchupQueue.splice(0));
 			}
+			this.#cancelQueueEntries(catchupQueue.splice(0).map((sub) => ["subscriber", sub]), void 0);
 		} catch (e) {
 			catchupQueue.forEach(({ subscriber }) => subscriber.fail(e));
 			throw e;
 		}
 	}
+	/**
+	* Processes a change from the queue, i.e. stores it in the change log as
+	* part of the current transaction (or skips it, see
+	* {@link assumeOwnershipAtHead}).
+	*
+	* @returns The transaction that is open after the change, or `null` if the
+	*     change ended (or was outside of) a transaction.
+	*/
+	async #processChange(tx, msg) {
+		const [_, watermark, json, storedChange, change] = msg;
+		const tag = change?.tag;
+		if (tag === "begin") {
+			assert(!tx, "received BEGIN in the middle of a transaction");
+			const pending = this.#pendingOwnership;
+			if (pending) {
+				if (change.skipAck) {
+					pending.skippedBackfill = true;
+					return {
+						skipped: true,
+						ack: false
+					};
+				}
+				if (watermark <= pending.head) return {
+					skipped: true,
+					ack: true
+				};
+			}
+			const pool = new TransactionPool(this.#lc.withContext("watermark", watermark), { mode: READ_COMMITTED });
+			pool.run(this.#inserter);
+			const { promise, resolve, reject } = resolver();
+			promise.catch(() => {});
+			const assumed = pending ? await this.#tryAssumeOwnership(pool, watermark, true) : void 0;
+			if (assumed === null) {
+				pool.abort();
+				await pool.done();
+				return {
+					skipped: true,
+					ack: true
+				};
+			}
+			if (pending) resolve({ owner: this.#taskID });
+			else pool.process((tx) => {
+				tx`
+            SELECT "owner" FROM ${this.#cdc("replicationState")} FOR UPDATE`.then(([result]) => resolve(result), reject);
+				return [];
+			});
+			tx = {
+				pool,
+				preCommitWatermark: watermark,
+				pos: 0,
+				startingReplicationState: promise,
+				ack: !change.skipAck,
+				batch: [],
+				lastFlush: void 0,
+				ownershipAssumed: assumed
+			};
+		} else if (tx?.skipped) {
+			if (tag === "commit" && tx.ack) this.#onCommitted([
+				"commit",
+				change,
+				{ watermark }
+			]);
+			return tag === "commit" || tag === "rollback" ? null : tx;
+		} else {
+			assert(tx, () => `received change outside of transaction: ${json}`);
+			tx.pos++;
+		}
+		const entry = {
+			watermark: tag === "commit" ? watermark : tx.preCommitWatermark,
+			precommit: tag === "commit" ? tx.preCommitWatermark : null,
+			pos: tx.pos,
+			change: storedChange
+		};
+		if (change !== null && isSchemaChange(change)) {
+			await this.#flushChangeLog(tx);
+			tx.lastFlush = tx.pool.process((sql) => [sql`INSERT INTO ${this.#cdc("changeLog")} ${sql(entry)}`, ...this.#trackBackfillMetadata(sql, change)]);
+		} else {
+			tx.batch.push(entry);
+			if (tx.batch.length >= this.#changeLogBatchSize) {
+				const prevFlush = tx.lastFlush;
+				this.#flushChangeLog(tx);
+				await prevFlush;
+			}
+		}
+		if (tag === "commit") {
+			this.#flushChangeLog(tx);
+			const { owner } = await tx.startingReplicationState;
+			if (owner !== this.#taskID) tx.pool.fail(new AbortError(`changeLog ownership has been assumed by ${owner}`));
+			else {
+				const lastWatermark = watermark;
+				tx.pool.process((tx) => [tx`
+        UPDATE ${this.#cdc("replicationState")} SET ${tx({ lastWatermark })}`]);
+				tx.pool.setDone();
+			}
+			await tx.pool.done();
+			const { ownershipAssumed, ack } = tx;
+			tx = null;
+			if (ownershipAssumed) this.#ownershipAssumed(ownershipAssumed);
+			if (ack) this.#onCommitted([
+				"commit",
+				change,
+				{ watermark }
+			]);
+		} else if (tag === "rollback") {
+			tx.pool.abort();
+			await tx.pool.done();
+			tx = null;
+		}
+		return tx;
+	}
 	async #startCatchup(subs) {
 		const numCatchups = subs.length;
 		if (numCatchups === 0) return;
@@ -510,6 +696,8 @@ var Storer = class {
 			lc.info?.(`snapshotted db at ${lastWatermark} (${await snapshotID}) to catchup ${numCatchups} subscriber(s)`);
 		} catch (e) {
 			subs.map(({ subscriber }) => subscriber.fail(e));
+			reader.abort();
+			catchupConns.end().catch(() => {});
 			throw e;
 		} finally {
 			catchupSnapshotted();
@@ -690,11 +878,18 @@ var PurgeLock = class {
 	#tx;
 	replicaVersion;
 	minWatermark;
-	constructor(lc, tx, replicaVersion, watermark) {
+	/**
+	* The head of the change log (i.e. its `lastWatermark`) when the lock was
+	* acquired. The change log is resumed from its head, which only advances,
+	* so a stream cannot (correctly) be started from before it.
+	*/
+	headWatermark;
+	constructor(lc, tx, replicaVersion, watermark, headWatermark) {
 		this.#lc = lc;
 		this.#tx = tx;
 		this.replicaVersion = replicaVersion;
 		this.minWatermark = watermark;
+		this.headWatermark = headWatermark;
 	}
 	#released = false;
 	async release() {
@@ -717,8 +912,23 @@ var PurgeLocker = class {
 	#cdc(table) {
 		return this.#db(`${cdcSchema(this.#shard)}.${table}`);
 	}
-	async acquire() {
+	async acquire(slotWatermark) {
 		const tx = new TransactionPool(this.#lc, { mode: READ_COMMITTED }).run(this.#db);
+		const [state] = await tx.processReadTask((sql) => sql`
+      SELECT "lastWatermark" FROM ${this.#cdc("replicationState")}`);
+		const lastWatermark = state?.lastWatermark;
+		if (slotWatermark !== void 0 && (lastWatermark === void 0 || lastWatermark < slotWatermark)) {
+			this.#lc.info?.(`changeLog@${lastWatermark ?? "(uninitialized)"} is behind the replication slot@${slotWatermark}`);
+			tx.setDone();
+			await tx.done();
+			return "behind-slot";
+		}
+		if (lastWatermark === void 0) {
+			this.#lc.info?.(`changeLog is uninitialized. No rows to purge-lock.`);
+			tx.setDone();
+			await tx.done();
+			return null;
+		}
 		const row = await tx.processReadTask((sql) => sql`
       SELECT watermark FROM ${this.#cdc("changeLog")}
         ORDER BY watermark, pos LIMIT 1
@@ -735,7 +945,7 @@ var PurgeLocker = class {
         SELECT "replicaVersion" FROM ${this.#cdc("replicationConfig")}
       `);
 		this.#lc.info?.(`locked watermark ${watermark} from being purged from replica@${replicaVersion}`);
-		return new PurgeLock(this.#lc, tx, replicaVersion, watermark);
+		return new PurgeLock(this.#lc, tx, replicaVersion, watermark, lastWatermark);
 	}
 };
 /**