@clickup/pg-microsharding/src/actions/actionRebalance.ts

Microshards support for PostgreSQL

import { resolve } from "path";
import chalk from "chalk";
import compact from "lodash/compact";
import first from "lodash/first";
import flatten from "lodash/flatten";
import sumBy from "lodash/sumBy";
import prompts from "prompts";
import { rebalance } from "../api/rebalance";
import { chooseDsns } from "../internal/chooseDsns";
import { chunkIntoN } from "../internal/chunkIntoN";
import { isTrueValue } from "../internal/isTrueValue";
import { print } from "../internal/logging";
import { dsnFromToShort, dsnShort } from "../internal/names";
import { normalizeDsns } from "../internal/normalizeDsns";
import type { Args } from "../internal/parseArgs";
import { pluralize } from "../internal/pluralize";
import { runInTmux, tryReattachToTmuxSession } from "../internal/runInTmux";
import { calcIslandWeights } from "./actionList";

/**
 * Runs a series of shard modes, concurrently, using TMUX if available.
 */
export async function actionRebalance(args: Args): Promise<boolean> {
  await tryReattachToTmuxSession();

  const dsns = await normalizeDsns(args["dsns"] || args["dsn"]);
  if (dsns.length === 0) {
    throw "Please provide --dsns, DB DSNs to rebalance microshards against";
  }

  const activateOnDestination = isTrueValue(args["activate-on-destination"]);
  if (activateOnDestination === undefined) {
    throw "Please provide --activate-on-destination=yes or --activate-on-destination=no";
  }

  const weightSql = args["weight-sql"] || undefined;

  const [decommissionDsns] = chooseDsns({
    patterns: compact(
      (args["decommission"] || args["decomission"] || "").split(/[,\s]+/s),
    ),
    existingDsns: dsns,
    requireMatch: true,
  });
  const maxReplicationLagSec =
    Number(args["max-replication-lag-sec"] ?? "") || undefined;
  const parallelism = Number(args["parallelism"] ?? "") || 3;
  const deactivateSQL = String(args["deactivate-sql"] || "") || undefined;
  const wait = args["wait"];

  const { islandNosToDsn, islands } = await calcIslandWeights({
    dsns,
    weightSql,
  });

  // Always exclude an island if it only has one global shard and no other
  // shards. We never move it while doing rebalance, since we don't know the
  // best destination fot it. The global shard can only be moved explicitly.
  const specialIslands = new Map<number, string>();
  for (const [no, shards] of islands.entries()) {
    if (shards.length === 1 && shards[0].no === 0) {
      if (decommissionDsns.includes(islandNosToDsn.get(no)!)) {
        throw (
          `Cannot decommission island ${dsnShort(islandNosToDsn.get(no)!)}, ` +
          "because it has only the global microshard and no other microshards.\n" +
          "Please use the move action for such an island."
        );
      }

      islands.delete(no);
      specialIslands.set(
        no,
        "an island with only the global microshard; won't be touched",
      );
    }
  }

  print.section("CURRENT islands weights:");
  print(inspectIslandsWeights(islandNosToDsn, islands, specialIslands));

  const dsnsToIslandNo = new Map([...dsns.entries()].map(([k, v]) => [v, k]));
  const decommissionIslandNos = decommissionDsns.map((dsn) => {
    const islandNo = dsnsToIslandNo.get(dsn);
    if (islandNo === undefined) {
      throw `No such island in the list: "${dsnShort(dsn)}"`;
    } else {
      return islandNo;
    }
  });

  const groupedMoves = rebalance(islands, decommissionIslandNos).map(
    ({ from, to, shards }) => ({
      fromDsn: islandNosToDsn.get(from)!,
      toDsn: islandNosToDsn.get(to)!,
      shards,
    }),
  );
  const moves = flatten(
    groupedMoves.map(({ fromDsn, toDsn, shards }) =>
      shards.map(({ no }) => ({ fromDsn, toDsn, no })),
    ),
  );

  if (moves.length === 0) {
    print.section(
      `Microshards are already fairly distributed among ${pluralize(islandNosToDsn.size, "island")}.`,
    );
    return true;
  }

  print.section(`Proposed moves (${moves.length}):`);
  print(
    groupedMoves
      .map(
        ({ fromDsn, toDsn, shards }) =>
          `- ${dsnFromToShort(fromDsn, toDsn)}: microshards [` +
          shards.map((shard) => shard.no).join(", ") +
          "]\n",
      )
      .join(""),
  );

  print.section("FUTURE weights after the above moves:");
  print(inspectIslandsWeights(islandNosToDsn, islands, specialIslands));

  const validResponse = "rebalance";
  const response = await prompts({
    type: "text",
    name: "value",
    message: `Type "${validResponse}" to proceed with the moves:`,
    validate: (value: string) =>
      value !== validResponse
        ? `Enter "${validResponse}" to confirm, ^C to skip`
        : true,
  });
  if (response.value !== validResponse) {
    return false;
  }

  const panes = chunkIntoN(moves, parallelism).map((sequence) =>
    sequence.map(({ fromDsn, toDsn, no }) => ({
      key: no.toString(),
      title: `microshard ${no} | ${dsnFromToShort(fromDsn, toDsn)}`,
      command: [
        "node",
        resolve(`${__dirname}/../cli`),
        "move",
        "--shard",
        no.toString(),
        "--from",
        fromDsn,
        "--to",
        toDsn,
        "--activate-on-destination",
        activateOnDestination ? "yes" : "no",
        ...(maxReplicationLagSec
          ? ["--max-replication-lag-sec", maxReplicationLagSec.toString()]
          : []),
        ...(deactivateSQL ? ["--deactivate-sql", deactivateSQL] : []),
        ...(wait ? ["--wait"] : []),
        "--skip-config",
      ],
    })),
  );

  await runInTmux(panes);
  return true;
}

/**
 * Prints weights of islands and their shards.
 */
function inspectIslandsWeights(
  islandNosToDsn: Map<number, string>,
  islands: Map<
    number,
    Array<{
      weight: number;
      appliedFactor: number;
      unit: string | undefined;
      no: number;
    }>
  >,
  specialIslands: Map<number, string>,
): string {
  const totalWeight =
    sumBy([...islands.values()], (shards) =>
      sumBy(shards, ({ weight }) => weight),
    ) || 1;
  const unit = first(
    compact([...islands.values()].map((shards) => first(shards)?.unit)),
  );
  return [...islands.entries(), ...specialIslands.entries()]
    .sort(([noA], [noB]) => noA - noB)
    .map(([no, shardsOrText]) => {
      const prefix = `- ${dsnShort(islandNosToDsn.get(no)!)}: `;
      if (typeof shardsOrText === "string") {
        return chalk.yellow(prefix + shardsOrText) + "\n";
      } else {
        const islandWeight = sumBy(shardsOrText, ({ weight }) => weight);
        const shardsWithAppliedFactor = sumBy(shardsOrText, (shard) =>
          shard.appliedFactor !== 1 ? 1 : 0,
        );
        return (
          prefix +
          `${Math.round((islandWeight / totalWeight) * 100)}% ` +
          `(${Math.round(islandWeight)}${unit ? ` ${unit}` : ""}${shardsWithAppliedFactor ? " adjusted by factor" : ""}), ` +
          `${pluralize(shardsOrText.length, "microshard")}${shardsWithAppliedFactor ? ` (${shardsWithAppliedFactor} with factor)` : ""}\n`
        );
      }
    })
    .join("");
}