import { resolver } from "@rocicorp/resolver";
import { getHeapStatistics } from "node:v8";
import "postgres";
import { AbortError } from "../../../../shared/src/abort-error.js";
import { assert } from "../../../../shared/src/asserts.js";
import { BigIntJSON } from "../../../../shared/src/bigint-json.js";
import { Queue } from "../../../../shared/src/queue.js";
import { promiseVoid } from "../../../../shared/src/resolved-promises.js";
import { parse } from "../../../../shared/src/valita.js";
import { READONLY, READ_COMMITTED } from "../../db/mode-enum.js";
import { runTx } from "../../db/run-transaction.js";
import { TransactionPool } from "../../db/transaction-pool.js";
import "../../types/pg.js";
import { cdcSchema } from "../../types/shards.js";
import "../change-source/protocol/current/control.js";
import { isDataChange, isSchemaChange } from "../change-source/protocol/current/data.js";
import "../change-source/protocol/current/downstream.js";
import "../change-source/protocol/current/json.js";
import "../change-source/protocol/current/status.js";
import { backfillRequestSchema } from "../change-source/protocol/current/upstream.js";
import "./change-streamer.js";
import { WatermarkTooOld } from "./error-type-enum.js";
import { AutoResetSignal, markResetRequired } from "./schema/tables.js";
import { array } from "@badrap/valita";
const backfillRequestsSchema = array(backfillRequestSchema);
class Storer {
  id = "storer";
  #lc;
  #shard;
  #taskID;
  #discoveryAddress;
  #discoveryProtocol;
  #db;
  #replicaVersion;
  #onConsumed;
  #onFatal;
  #queue = new Queue();
  #backPressureThresholdBytes;
  #approximateQueuedBytes = 0;
  #running = false;
  constructor(lc, shard, taskID, discoveryAddress, discoveryProtocol, db, replicaVersion, onConsumed, onFatal, backPressureLimitHeapProportion) {
    this.#lc = lc.withContext("component", "change-log");
    this.#shard = shard;
    this.#taskID = taskID;
    this.#discoveryAddress = discoveryAddress;
    this.#discoveryProtocol = discoveryProtocol;
    this.#db = db;
    this.#replicaVersion = replicaVersion;
    this.#onConsumed = onConsumed;
    this.#onFatal = onFatal;
    const heapStats = getHeapStatistics();
    this.#backPressureThresholdBytes = (heapStats.heap_size_limit - heapStats.used_heap_size) * backPressureLimitHeapProportion;
    this.#lc.info?.(
      `Using up to ${(this.#backPressureThresholdBytes / 1024 ** 2).toFixed(2)} MB of --max-old-space-size (~${(heapStats.heap_size_limit / 1024 ** 2).toFixed(2)} MB) to absorb upstream spikes`,
      { heapStats }
    );
  }
  // For readability in SQL statements.
  #cdc(table) {
    return this.#db(`${cdcSchema(this.#shard)}.${table}`);
  }
  async assumeOwnership() {
    const db = this.#db;
    const owner = this.#taskID;
    const ownerAddress = this.#discoveryAddress;
    const ownerProtocol = this.#discoveryProtocol;
    const addressWithProtocol = ownerProtocol === "ws" ? ownerAddress : `${ownerProtocol}://${ownerAddress}`;
    this.#lc.info?.(`assuming ownership at ${addressWithProtocol}`);
    const start = performance.now();
    await db`UPDATE ${this.#cdc("replicationState")} SET ${db({ owner, ownerAddress: addressWithProtocol })}`;
    const elapsed = (performance.now() - start).toFixed(2);
    this.#lc.info?.(
      `assumed ownership at ${addressWithProtocol} (${elapsed} ms)`
    );
  }
  async getStartStreamInitializationParameters() {
    const [[{ lastWatermark }], result] = await runTx(
      this.#db,
      (sql) => [
        sql`
        SELECT "lastWatermark" FROM ${this.#cdc("replicationState")}`,
        // Formats a BackfillRequest using json_object_agg() to construct the
        // `columns` object. It is LEFT JOIN'ed with the `tableMetadata` table
        // to make it optional and possibly `null`.
        sql`
        SELECT 
            json_build_object(
              'schema', b."schema",
              'name', b."table",
              'metadata', t."metadata"
            ) as "table",
            json_object_agg(b."column", b."backfill") 
              as "columns"
          FROM ${this.#cdc("backfilling")} as b
          LEFT JOIN ${this.#cdc("tableMetadata")} as t
          ON (b."schema" = t."schema" AND b."table" = t."table")
          GROUP BY b."schema", b."table", t."metadata"
        `
      ],
      { mode: READONLY }
    );
    return {
      lastWatermark,
      backfillRequests: parse(result, backfillRequestsSchema)
    };
  }
  async getMinWatermarkForCatchup() {
    const [{ minWatermark }] = await this.#db`
      SELECT min(watermark) as "minWatermark" FROM ${this.#cdc("changeLog")}`;
    return minWatermark;
  }
  purgeRecordsBefore(watermark) {
    return runTx(this.#db, async (sql) => {
      const [{ deleted }] = await sql`
        WITH purged AS (
          DELETE FROM ${this.#cdc("changeLog")} WHERE watermark < ${watermark} 
            RETURNING watermark, pos
        ) SELECT COUNT(*) as deleted FROM purged;`;
      const [{ owner }] = await sql`
        SELECT * FROM ${this.#cdc("replicationState")} FOR SHARE`;
      if (owner !== this.#taskID) {
        throw new AbortError(
          `aborting changeLog purge to ${watermark} because ownership has been taken by ${owner}`
        );
      }
      return Number(deleted);
    });
  }
  /**
   * @returns The size of the serialized entry, for memory / I/O estimations.
   */
  store(entry) {
    const [watermark, [_tag, change]] = entry;
    const json = BigIntJSON.stringify(change);
    this.#approximateQueuedBytes += json.length;
    this.#queue.enqueue([
      "change",
      watermark,
      json,
      isDataChange(change) ? null : change
      // drop DataChanges to save memory
    ]);
    return json.length;
  }
  abort() {
    this.#queue.enqueue(["abort"]);
  }
  status(s) {
    this.#queue.enqueue(s);
  }
  catchup(subscriber, mode) {
    this.#queue.enqueue(["subscriber", { subscriber, mode }]);
  }
  #readyForMore = null;
  readyForMore() {
    if (!this.#running) {
      return void 0;
    }
    if (this.#readyForMore === null && this.#approximateQueuedBytes > this.#backPressureThresholdBytes) {
      this.#lc.warn?.(
        `applying back pressure with ${this.#queue.size()} queued changes (~${(this.#approximateQueuedBytes / 1024 ** 2).toFixed(2)} MB)

To inspect changeLog backlog in your change DB:
  SELECT
    (change->'relation'->>'schema') || '.' || (change->'relation'->>'name') AS table_name,
    change->>'tag' AS operation,
    COUNT(*) AS count
  FROM "<app_id>/cdc"."changeLog"
  GROUP BY 1, 2
  ORDER BY 3 DESC
  LIMIT 20;`
      );
      this.#readyForMore = resolver();
    }
    return this.#readyForMore?.promise;
  }
  #maybeReleaseBackPressure() {
    if (this.#readyForMore !== null && // Wait for at least 20% of the threshold to free up.
    this.#approximateQueuedBytes < this.#backPressureThresholdBytes * 0.8) {
      this.#lc.info?.(
        `releasing back pressure with ${this.#queue.size()} queued changes (~${(this.#approximateQueuedBytes / 1024 ** 2).toFixed(2)} MB)`
      );
      this.#readyForMore.resolve();
      this.#readyForMore = null;
    }
  }
  #stopped = promiseVoid;
  /**
   * Runs the storer loop until {@link stop()} is called, or an error is thrown.
   * Once {@link run()} completes, it can be called again.
   */
  async run() {
    assert(!this.#running, `storer is already running`);
    const { promise: stopped, resolve: signalStopped } = resolver();
    this.#running = true;
    this.#stopped = stopped;
    this.#lc.info?.("starting storer");
    try {
      await this.#processQueue();
    } finally {
      if (this.#readyForMore !== null) {
        this.#readyForMore.resolve();
        this.#readyForMore = null;
      }
      const unprocessed = this.#queue.drain();
      if (unprocessed.length) {
        this.#lc.warn?.(
          `dropped ${unprocessed.length} entries from the changeLog queue`
        );
      }
      this.#running = false;
      signalStopped();
      this.#lc.info?.("storer stopped");
    }
  }
  async #processQueue() {
    let tx = null;
    let msg;
    const catchupQueue = [];
    while ((msg = await this.#queue.dequeue()) !== "stop") {
      const [msgType] = msg;
      switch (msgType) {
        case "ready": {
          const signalReady = msg[1];
          signalReady();
          continue;
        }
        case "subscriber": {
          const subscriber = msg[1];
          if (tx) {
            catchupQueue.push(subscriber);
          } else {
            await this.#startCatchup([subscriber]);
          }
          continue;
        }
        case "status":
          this.#onConsumed(msg);
          continue;
        case "abort": {
          if (tx) {
            tx.pool.abort();
            await tx.pool.done();
            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();
        tx = {
          pool: new TransactionPool(
            this.#lc.withContext("watermark", watermark),
            READ_COMMITTED
          ),
          preCommitWatermark: watermark,
          pos: 0,
          startingReplicationState: promise,
          ack: !change.skipAck
        };
        tx.pool.run(this.#db);
        void tx.pool.process((tx2) => {
          tx2`
          SELECT * 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: json
      };
      const processed = tx.pool.process((sql) => [
        sql`INSERT INTO ${this.#cdc("changeLog")} ${sql(entry)}`,
        ...change !== null && isSchemaChange(change) ? this.#trackBackfillMetadata(sql, change) : []
      ]);
      if (tx.pos % 100 === 0) {
        await processed;
      }
      this.#maybeReleaseBackPressure();
      if (tag === "commit") {
        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;
          void tx.pool.process((tx2) => [
            tx2`
            UPDATE ${this.#cdc("replicationState")} SET ${tx2({ lastWatermark })}`
          ]);
          tx.pool.setDone();
        }
        await tx.pool.done();
        if (tx.ack) {
          this.#onConsumed(["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));
      }
    }
  }
  async #startCatchup(subs) {
    if (subs.length === 0) {
      return;
    }
    const reader = new TransactionPool(
      this.#lc.withContext("pool", "catchup"),
      READONLY
    );
    reader.run(this.#db);
    const [{ lastWatermark }] = await reader.processReadTask(
      (sql) => sql`
        SELECT * FROM ${this.#cdc("replicationState")}
      `
    );
    void Promise.all(
      subs.map((sub) => this.#catchup(sub, lastWatermark, reader))
    ).finally(() => reader.setDone());
  }
  async #catchup({ subscriber: sub, mode }, lastWatermark, reader) {
    try {
      await reader.processReadTask(async (tx) => {
        const start = Date.now();
        let watermarkFound = sub.watermark === this.#replicaVersion;
        let count = 0;
        let lastBatchConsumed;
        for await (const entries of tx`
          SELECT watermark, change FROM ${this.#cdc("changeLog")}
           WHERE watermark >= ${sub.watermark}
             AND watermark <= ${lastWatermark}
           ORDER BY watermark, pos`.cursor(2e3)) {
          const start2 = performance.now();
          await lastBatchConsumed;
          const elapsed = performance.now() - start2;
          if (lastBatchConsumed) {
            (elapsed > 100 ? this.#lc.info : this.#lc.debug)?.(
              `waited ${elapsed.toFixed(3)} ms for ${sub.id} to consume last batch of catchup entries`
            );
          }
          for (const entry of entries) {
            if (entry.watermark === sub.watermark) {
              watermarkFound = true;
            } else if (watermarkFound) {
              lastBatchConsumed = sub.catchup(toDownstream(entry));
              count++;
            } else if (mode === "backup") {
              throw new AutoResetSignal(
                `backup replica at watermark ${sub.watermark} is behind change db: ${entry.watermark})`
              );
            } else {
              this.#lc.warn?.(
                `rejecting subscriber at watermark ${sub.watermark} (earliest watermark: ${entry.watermark})`
              );
              sub.close(
                WatermarkTooOld,
                `earliest supported watermark is ${entry.watermark} (requested ${sub.watermark})`
              );
              return;
            }
          }
        }
        if (watermarkFound) {
          await lastBatchConsumed;
          this.#lc.info?.(
            `caught up ${sub.id} with ${count} changes (${Date.now() - start} ms)`
          );
        } else {
          this.#lc.warn?.(
            `subscriber at watermark ${sub.watermark} is ahead of latest watermark`
          );
        }
        sub.setCaughtUp();
      });
    } catch (err) {
      this.#lc.error?.(`error while catching up subscriber ${sub.id}`, err);
      if (err instanceof AutoResetSignal) {
        await markResetRequired(this.#db, this.#shard);
        this.#onFatal(err);
      }
      sub.fail(err);
    }
  }
  /**
   * Returns the db statements necessary to track backfill and table metadata
   * presented in the `change`, if any.
   */
  #trackBackfillMetadata(sql, change) {
    const stmts = [];
    switch (change.tag) {
      case "update-table-metadata": {
        const { table, new: metadata } = change;
        stmts.push(this.#upsertTableMetadataStmt(sql, table, metadata));
        break;
      }
      case "create-table": {
        const { spec, metadata, backfill } = change;
        if (metadata) {
          stmts.push(this.#upsertTableMetadataStmt(sql, spec, metadata));
        }
        if (backfill) {
          Object.entries(backfill).forEach(([col, backfill2]) => {
            stmts.push(
              this.#upsertColumnBackfillStmt(sql, spec, col, backfill2)
            );
          });
        }
        break;
      }
      case "rename-table": {
        const { old } = change;
        const row = { schema: change.new.schema, table: change.new.name };
        stmts.push(
          sql`UPDATE ${this.#cdc("tableMetadata")} SET ${sql(row)}
                WHERE "schema" = ${old.schema} AND "table" = ${old.name}`,
          sql`UPDATE ${this.#cdc("backfilling")} SET ${sql(row)}
                WHERE "schema" = ${old.schema} AND "table" = ${old.name}`
        );
        break;
      }
      case "drop-table": {
        const {
          id: { schema, name }
        } = change;
        stmts.push(
          sql`DELETE FROM ${this.#cdc("tableMetadata")}
                WHERE "schema" = ${schema} AND "table" = ${name}`,
          sql`DELETE FROM ${this.#cdc("backfilling")}
                WHERE "schema" = ${schema} AND "table" = ${name}`
        );
        break;
      }
      case "add-column": {
        const { table, tableMetadata, column, backfill } = change;
        if (tableMetadata) {
          stmts.push(this.#upsertTableMetadataStmt(sql, table, tableMetadata));
        }
        if (backfill) {
          stmts.push(
            this.#upsertColumnBackfillStmt(sql, table, column.name, backfill)
          );
        }
        break;
      }
      case "update-column": {
        const {
          table: { schema, name: table },
          old: { name: oldName },
          new: { name: newName }
        } = change;
        if (oldName !== newName) {
          stmts.push(
            sql`UPDATE ${this.#cdc("backfilling")} SET "column" = ${newName}
                WHERE "schema" = ${schema} AND "table" = ${table} AND "column" = ${oldName}`
          );
        }
        break;
      }
      case "drop-column": {
        const {
          table: { schema, name },
          column
        } = change;
        stmts.push(
          sql`DELETE FROM ${this.#cdc("backfilling")}
                WHERE "schema" = ${schema} AND "table" = ${name} AND "column" = ${column}`
        );
        break;
      }
      case "backfill-completed": {
        const {
          relation: { schema, name: table, rowKey },
          columns
        } = change;
        const cols = [...rowKey.columns, ...columns];
        stmts.push(
          sql`DELETE FROM ${this.#cdc("backfilling")}
                WHERE "schema" = ${schema} AND "table" = ${table} AND "column" IN ${sql(cols)}`
        );
      }
    }
    return stmts;
  }
  #upsertTableMetadataStmt(sql, { schema, name: table }, metadata) {
    const row = { schema, table, metadata };
    return sql`
        INSERT INTO ${this.#cdc("tableMetadata")} ${sql(row)}
          ON CONFLICT ("schema", "table") 
          DO UPDATE SET ${sql(row)};
    `;
  }
  #upsertColumnBackfillStmt(sql, { schema, name: table }, column, backfill) {
    const row = { schema, table, column, backfill };
    return sql`
        INSERT INTO ${this.#cdc("backfilling")} ${sql(row)}
          ON CONFLICT ("schema", "table", "column") 
          DO UPDATE SET ${sql(row)};
    `;
  }
  /**
   * Waits until all currently queued entries have been processed.
   * This is only used in tests.
   */
  async allProcessed() {
    if (this.#running) {
      const { promise, resolve } = resolver();
      this.#queue.enqueue(["ready", resolve]);
      await promise;
    }
  }
  stop() {
    if (this.#running) {
      this.#lc.info?.(`draining ${this.#queue.size()} changeLog entries`);
      this.#queue.enqueue("stop");
    }
    return this.#stopped;
  }
}
function toDownstream(entry) {
  const { watermark, change } = entry;
  switch (change.tag) {
    case "begin":
      return [watermark, ["begin", change, { commitWatermark: watermark }]];
    case "commit":
      return [watermark, ["commit", change, { watermark }]];
    case "rollback":
      return [watermark, ["rollback", change]];
    default:
      return [watermark, ["data", change]];
  }
}
export {
  Storer
};
//# sourceMappingURL=storer.js.map
