@clickup/pg-microsharding/src/internal/waitUntilEverythingIsNotLaggingMuch.ts

Microshards support for PostgreSQL

import readline from "readline";
import type { ClearablePromise } from "delay";
import type { PhysicalReplica } from "./getPhysicalReplicas";
import { getPhysicalReplicas } from "./getPhysicalReplicas";
import type { SlotInfo } from "./getSlotInfo";
import { getSlotInfo } from "./getSlotInfo";
import type { TableInfo } from "./getTablesInSchema";
import { log, progress } from "./logging";
import { subName } from "./names";
import { pollDelay } from "./pollDelay";
import { secToHuman } from "./secToHuman";
import { tookToHuman } from "./tookToHuman";

const FORCE_CONTINUE_KEY = {
  sequence: "S",
  prompt: "press Shift+S to force-continue: not recommended!",
};

/**
 * Waits until all physical replicas on a DSN are not lagging for too much, so
 * we can e.g. call waitUntilIncrementalCompletes() and then resultCommit().
 * Also, if fromDsn is passed, waits for the logical slot to lag not more than
 * maxReplicationLagSec time.
 *
 * For the logical slot replication lag estimation, in the beginning, we
 * "latch", i.e. remember the master's LSN, and then wait until the slot replay
 * is past that LSN. Once it happens, the amount of time passed is roughly the
 * logical replication slot lag in seconds. This is the only way to estimate,
 * how many seconds behind is a logical slot.
 *
 * For physical replicas, we query their replication lag in seconds directly.
 *
 * If waitForUserToContinue=true, then the function exits ONLY when the user
 * told it to do so explicitly by pressing the hotkey.
 */
export async function waitUntilEverythingIsNotLaggingMuch(
  {
    when,
    fromDsn,
    tables,
    toDsn,
    schema,
    maxReplicationLagSec: maxLagSec,
    waitForUserToContinue,
  }: {
    when: string;
    fromDsn?: string | null;
    tables: TableInfo[];
    toDsn: string;
    schema: string;
    maxReplicationLagSec: number;
    waitForUserToContinue: boolean;
  },
  throwIfAborted?: () => void,
): Promise<void> {
  if (tables.length === 0) {
    return;
  }

  const waitStartedAt = performance.now();

  const rl = readline.createInterface({
    input: process.stdin,
    output: process.stdout,
  });

  try {
    let physicalReplicas: PhysicalReplica[] | undefined = undefined;
    let slotInfo: SlotInfo | undefined = undefined;
    let stats: string[] = [];
    let delayPromise: ClearablePromise<void> | null = null;
    let userForcedContinue = false;

    readline.emitKeypressEvents(process.stdin, rl);
    rl.on("SIGINT", () => {
      throwIfAborted = () => {
        throw "Aborted due to SIGINT";
      };
    });
    process.stdin.setRawMode?.(true);
    process.stdin.on("keypress", (_str, key) => {
      if (key.sequence === FORCE_CONTINUE_KEY.sequence) {
        userForcedContinue = true;
        delayPromise?.clear();
      }
    });

    for (let i = 0, needWait = true; needWait && !userForcedContinue; i++) {
      throwIfAborted?.();

      physicalReplicas = await getPhysicalReplicas({
        dsn: toDsn,
        comment:
          i === 0
            ? `Waiting till ${fromDsn ? "the slot and " : ""}replicas are not lagging much (${when})`
            : undefined,
      });
      slotInfo = fromDsn
        ? await getSlotInfo(
            { dsn: fromDsn, slotName: subName(schema) },
            slotInfo,
          )
        : undefined;

      needWait = waitForUserToContinue;
      stats = [];

      if (slotInfo) {
        const { srcLsn, dstLsn, gap, gapTillLatch, lagSec } = slotInfo;
        const prefix = "Logical slot";
        const info = `src=${srcLsn} dst=${dstLsn} gap=${gap} gapTillLatch=${gapTillLatch}`;
        if (lagSec === null) {
          needWait = true;
          stats.push(
            `? ${prefix}: lag is unknown yet (likely LARGE; be patient, estimating); ${info}`,
          );
        } else if (lagSec > maxLagSec) {
          needWait = true;
          stats.push(
            `- ${prefix}: lag is ${secToHuman(lagSec)} > ${secToHuman(maxLagSec)}; ${info}`,
          );
        } else {
          stats.push(
            `+ ${prefix}: lag is ${secToHuman(lagSec)} <= ${secToHuman(maxLagSec)}; ${info}`,
          );
        }
      }

      for (const {
        clientAddr,
        userName,
        applicationName,
        lagSec,
        feedbackAgoSec,
        masterLsn,
        replicaLsn,
        gap,
      } of physicalReplicas) {
        const prefix = `Physical replica of the dst ${userName}@${clientAddr} (${applicationName})`;
        const info = `master=${masterLsn} replica=${replicaLsn} gap=${gap} feedbackAgoSec=${feedbackAgoSec}`;
        if (lagSec === null) {
          needWait = true;
          stats.push(`? ${prefix}: lag is unknown yet; ${info}`);
        } else if (lagSec > maxLagSec) {
          needWait = true;
          stats.push(
            `- ${prefix}: lag is ${secToHuman(lagSec)} > ${secToHuman(maxLagSec)}; ${info}`,
          );
        } else {
          stats.push(
            `+ ${prefix}: lag is ${secToHuman(lagSec)} <= ${secToHuman(maxLagSec)}; ${info}`,
          );
        }
      }

      progress(
        waitForUserToContinue
          ? `Waiting for the user action: ${FORCE_CONTINUE_KEY.prompt}`
          : `Waiting for the lag to stabilize (or ${FORCE_CONTINUE_KEY.prompt}):`,
        ...stats,
      );
      await (delayPromise = pollDelay());
    }

    progress.clear();
    log.shellCmd({
      comment: userForcedContinue
        ? `Force-continued in ${tookToHuman(waitStartedAt)}`
        : `Everything is caught up in ${tookToHuman(waitStartedAt)}`,
      cmd: "",
      input: stats.join("\n"),
    });
  } finally {
    rl.close();
    process.stdin.setRawMode?.(false);
    process.stdin.pause();
  }
}