kuber/lib/rollback.ts

import type {
  V1Deployment,
  V1PodTemplateSpec,
  V1ReplicaSet,
} from "@kubernetes/client-node";
import { LABELS } from "../const";
import { ctx } from "./context";
import { apps } from "./k8s";
import { listManagedDeployments, waitForDeploymentRollout } from "./shared";

const RevisionAnnotation = "deployment.kubernetes.io/revision";
const PodTemplateHashLabel = "pod-template-hash";
const ManagedBySelector = `app.kubernetes.io/managed-by=${LABELS["app.kubernetes.io/managed-by"]}`;

export type RollbackCandidate = {
  name: string;
  currentRevision: number;
  previousRevision: number;
  image: string;
};

function getRevision(rs: V1ReplicaSet): number | undefined {
  const raw = rs.metadata?.annotations?.[RevisionAnnotation];
  if (!raw) return undefined;
  const revision = Number(raw);
  return Number.isInteger(revision) && revision > 0 ? revision : undefined;
}

export function stripPodTemplateHash(template: V1PodTemplateSpec): V1PodTemplateSpec {
  const labels = { ...template.metadata?.labels };
  delete labels[PodTemplateHashLabel];
  return {
    ...template,
    metadata: {
      ...template.metadata,
      labels: Object.keys(labels).length > 0 ? labels : undefined,
    },
  };
}

export function templateMatches(
  template: V1PodTemplateSpec | undefined,
  previous: V1PodTemplateSpec | undefined,
): boolean {
  return (
    JSON.stringify(stripPodTemplateHash(template ?? {})) ===
    JSON.stringify(stripPodTemplateHash(previous ?? {}))
  );
}

export async function listDeploymentReplicaSets(
  deployment: V1Deployment,
): Promise<V1ReplicaSet[]> {
  const name = deployment.metadata?.name;
  const ownerUid = deployment.metadata?.uid;
  if (!name || !ownerUid) return [];
  const { items } = await apps.listNamespacedReplicaSet({
    namespace: ctx().project,
    labelSelector: ManagedBySelector,
  });
  return items.filter((rs) =>
    rs.metadata?.ownerReferences?.some(
      (ref) => ref.uid === ownerUid && ref.kind === "Deployment",
    ),
  );
}

export async function findRollbackCandidates(
  deployment: V1Deployment,
): Promise<RollbackCandidate | undefined> {
  const name = deployment.metadata?.name;
  if (!name) return undefined;
  const replicaSets = await listDeploymentReplicaSets(deployment);
  const revisions = replicaSets
    .map((rs) => ({ rs, revision: getRevision(rs) }))
    .filter(
      (entry): entry is { rs: V1ReplicaSet; revision: number } =>
        entry.revision !== undefined,
    )
    .sort((a, b) => b.revision - a.revision);

  if (revisions.length < 2) return undefined;

  const [current, ...rest] = revisions;
  if (!current) return undefined;
  const previous =
    rest.find((entry) =>
      !templateMatches(entry.rs.spec?.template, deployment.spec?.template),
    ) ?? rest[0];
  if (!previous) return undefined;

  const image = previous.rs.spec?.template?.spec?.containers?.[0]?.image;
  if (!image) return undefined;

  return {
    name,
    currentRevision: current.revision,
    previousRevision: previous.revision,
    image,
  };
}

export async function planRollback(
  names?: string[],
): Promise<RollbackCandidate[]> {
  const deployments = await listManagedDeployments();
  const targets = names
    ? deployments.filter((deployment) =>
        names.includes(deployment.metadata?.name ?? ""),
      )
    : deployments;

  if (names) {
    const found = new Set(targets.map((deployment) => deployment.metadata?.name));
    for (const name of names) {
      if (!found.has(name)) {
        throw new Error(`No managed deployment named ${name} in ${ctx().project}`);
      }
    }
  }

  const candidates: RollbackCandidate[] = [];
  for (const deployment of targets) {
    const candidate = await findRollbackCandidates(deployment);
    if (candidate) candidates.push(candidate);
  }
  return candidates;
}

export async function rollbackDeployment(
  candidate: RollbackCandidate,
): Promise<V1Deployment> {
  const { name, previousRevision } = candidate;
  const deployment = await apps.readNamespacedDeployment({
    namespace: ctx().project,
    name,
  });
  const replicaSets = await listDeploymentReplicaSets(deployment);
  const previous = replicaSets.find(
    (rs) => getRevision(rs) === previousRevision,
  );
  if (!previous?.spec?.template) {
    throw new Error(
      `Deployment ${name} has no ReplicaSet for revision ${previousRevision}`,
    );
  }

  const template = stripPodTemplateHash(previous.spec.template);

  return apps.patchNamespacedDeployment({
    namespace: ctx().project,
    name,
    body: [
      {
        op: "replace",
        path: "/spec/template",
        value: template,
      },
    ],
  });
}

export async function rollbackAll(candidates: RollbackCandidate[], timeoutMs = 300000) {
  for (const candidate of candidates) {
    await rollbackDeployment(candidate);
  }
  for (const candidate of candidates) {
    await waitForDeploymentRollout(candidate.name, timeoutMs);
  }
}