kuber/server/build-store.ts

import { randomUUID } from "node:crypto";
import type {
  BuildEvent,
  BuildRequest,
  BuildStatus,
  Sha256Digest,
} from "../shared/build-protocol";

export const BUILD_RECORD_API_VERSION = "kuber.astrxl.dev/v1" as const;

export interface BuildRecord {
  apiVersion: typeof BUILD_RECORD_API_VERSION;
  kind: "BuildRecord";
  metadata: {
    name: string;
    resourceVersion: string;
    creationTimestamp: string;
    labels: Record<string, string>;
  };
  spec: {
    request: BuildRequest;
    imageKey: string;
    jobName: string;
    workspaceSubPath: string;
  };
  status: BuildStatus & {
    cancelled?: boolean;
    jobCreated?: boolean;
    logBytes: number;
    logOffset: number;
    nextSequence: number;
    events: BuildEvent[];
    reconcileLease?: BuildReconciliationLease;
  };
}

/** A record-local fencing token for distributed build reconciliation. */
export interface BuildReconciliationLease {
  holder: string;
  token: string;
  expiresAt: string;
}

export interface UploadRecord {
  apiVersion: typeof BUILD_RECORD_API_VERSION;
  kind: "BuildUpload";
  metadata: {
    name: string;
    resourceVersion: string;
    creationTimestamp: string;
  };
  spec: { digest: Sha256Digest; size: number };
  status: { offset: number; data: Uint8Array };
}

export type CreateBuildResult = {
  record: BuildRecord;
  created: boolean;
  superseded?: BuildRecord[];
};

/** Implementations must make createBuild and replace operations atomic. */
export interface BuildStore {
  createBuild(record: BuildRecord): Promise<CreateBuildResult>;
  getBuild(id: string): Promise<BuildRecord | undefined>;
  listBuilds(): Promise<BuildRecord[]>;
  replaceBuild(
    record: BuildRecord,
    expectedResourceVersion: string,
    reconcileLeaseToken?: string,
  ): Promise<void>;
  acquireBuildReconciliationLease(
    id: string,
    holder: string,
    ttlMs: number,
  ): Promise<BuildReconciliationLease | undefined>;
  renewBuildReconciliationLease(
    id: string,
    token: string,
    ttlMs: number,
  ): Promise<boolean>;
  releaseBuildReconciliationLease(id: string, token: string): Promise<void>;
  ownsBuild?(imageKey: string, buildId: string): Promise<boolean>;
  getUpload(digest: Sha256Digest): Promise<UploadRecord | undefined>;
  createUpload(record: UploadRecord): Promise<UploadRecord>;
  replaceUpload(
    record: UploadRecord,
    expectedResourceVersion: string,
  ): Promise<void>;
  deleteUpload(digest: Sha256Digest): Promise<void>;
}

export class BuildStoreConflictError extends Error {
  readonly code = "BUILD_STORE_CONFLICT";
}

function clone<T>(value: T): T {
  return structuredClone(value);
}

function terminal(record: BuildRecord): boolean {
  return (
    record.status.state === "succeeded" || record.status.state === "failed"
  );
}

function sameSpec(left: BuildRecord, right: BuildRecord): boolean {
  return JSON.stringify(left.spec) === JSON.stringify(right.spec);
}

export function buildTimestamp(record: BuildRecord): string {
  const creationTimestamp = record.metadata?.creationTimestamp;
  if (typeof creationTimestamp === "string") return creationTimestamp;
  const createdAt = record.status?.createdAt;
  return typeof createdAt === "string" ? createdAt : "";
}

export function compareBuildRecords(
  left: BuildRecord,
  right: BuildRecord,
): number {
  const timestamp = buildTimestamp(left).localeCompare(buildTimestamp(right));
  if (timestamp !== 0) return timestamp;
  return left.metadata.name.localeCompare(right.metadata.name);
}

function newer(left: BuildRecord, right: BuildRecord): boolean {
  return compareBuildRecords(left, right) > 0;
}

function supersede(record: BuildRecord, finishedAt: string): BuildRecord {
  const reason = "Superseded by newer build";
  record.status = {
    ...record.status,
    state: "failed",
    finishedAt,
    error: reason,
    cancelled: true,
    events: [
      ...record.status.events,
      {
        type: "status",
        status: {
          version: record.status.version,
          id: record.status.id,
          state: "failed",
          createdAt: record.status.createdAt,
          ...(record.status.startedAt && {
            startedAt: record.status.startedAt,
          }),
          finishedAt,
          error: reason,
        },
      },
    ],
  };
  record.metadata.resourceVersion = String(
    Number(record.metadata.resourceVersion) + 1,
  );
  return record;
}

export class MemoryBuildStore implements BuildStore {
  private readonly builds = new Map<string, BuildRecord>();
  private readonly uploads = new Map<Sha256Digest, UploadRecord>();
  private readonly active = new Map<string, string>();
  private readonly byImage = new Map<string, Set<string>>();

  async createBuild(record: BuildRecord): Promise<CreateBuildResult> {
    const existing = this.builds.get(record.metadata.name);
    if (existing) {
      if (!sameSpec(existing, record)) {
        throw new BuildStoreConflictError(
          "Build ID was already used for a different request",
        );
      }
      if (terminal(existing))
        return { record: clone(existing), created: false };
      const candidates = [...(this.byImage.get(existing.spec.imageKey) ?? [])]
        .map((id) => this.builds.get(id))
        .filter(
          (value): value is BuildRecord =>
            !!value &&
            !terminal(value) &&
            value.metadata.name !== existing.metadata.name,
        );
      const newerCandidate = candidates.find((candidate) =>
        newer(candidate, existing),
      );
      if (newerCandidate) {
        const superseded = supersede(existing, buildTimestamp(existing));
        this.builds.set(existing.metadata.name, clone(superseded));
        if (this.active.get(existing.spec.imageKey) === existing.metadata.name)
          this.active.delete(existing.spec.imageKey);
        return {
          record: clone(superseded),
          created: false,
          superseded: [clone(superseded)],
        };
      }
      const superseded: BuildRecord[] = [];
      for (const old of candidates) {
        const next = supersede(old, buildTimestamp(existing));
        this.builds.set(old.metadata.name, clone(next));
        superseded.push(clone(next));
      }
      this.active.set(existing.spec.imageKey, existing.metadata.name);
      return {
        record: clone(existing),
        created: false,
        ...(superseded.length && { superseded }),
      };
    }
    const candidates = [...(this.byImage.get(record.spec.imageKey) ?? [])]
      .map((id) => this.builds.get(id))
      .filter((value): value is BuildRecord => !!value && !terminal(value));
    const active = candidates.sort(compareBuildRecords).at(-1);
    if (active && newer(active, record)) {
      const superseded = supersede(record, record.status.createdAt);
      this.builds.set(record.metadata.name, clone(superseded));
      this.byImage.get(record.spec.imageKey)?.add(record.metadata.name);
      return {
        record: clone(superseded),
        created: false,
        superseded: [clone(superseded)],
      };
    }
    if (active) {
      const old = supersede(active, record.status.createdAt);
      this.builds.set(active.metadata.name, clone(old));
      if (this.active.get(record.spec.imageKey) === active.metadata.name)
        this.active.delete(record.spec.imageKey);
    }
    this.builds.set(record.metadata.name, clone(record));
    this.byImage.set(
      record.spec.imageKey,
      new Set([
        ...(this.byImage.get(record.spec.imageKey) ?? []),
        record.metadata.name,
      ]),
    );
    this.active.set(record.spec.imageKey, record.metadata.name);
    return {
      record: clone(record),
      created: true,
      superseded: active
        ? [clone(this.builds.get(active.metadata.name)!)]
        : undefined,
    };
  }

  async getBuild(id: string) {
    const value = this.builds.get(id);
    return value && clone(value);
  }

  async listBuilds() {
    return [...this.builds.values()].sort(compareBuildRecords).map(clone);
  }

  async replaceBuild(
    record: BuildRecord,
    expectedResourceVersion: string,
    reconcileLeaseToken?: string,
  ) {
    const current = this.builds.get(record.metadata.name);
    if (
      !current ||
      current.metadata.resourceVersion !== expectedResourceVersion
    ) {
      throw new BuildStoreConflictError(
        "Build record was concurrently modified",
      );
    }
    if (!sameSpec(current, record)) {
      throw new BuildStoreConflictError("Build specification is immutable");
    }
    if (
      reconcileLeaseToken &&
      current.status.reconcileLease?.token !== reconcileLeaseToken
    ) {
      throw new BuildStoreConflictError(
        "Build reconciliation lease ownership was lost",
      );
    }
    if (
      current.status.state === "succeeded" &&
      (record.status.state !== "succeeded" ||
        record.status.digest !== current.status.digest)
    ) {
      throw new BuildStoreConflictError(
        "A successful image digest is immutable",
      );
    }
    this.builds.set(record.metadata.name, clone(record));
    if (
      terminal(record) &&
      this.active.get(record.spec.imageKey) === record.metadata.name
    ) {
      this.active.delete(record.spec.imageKey);
    }
  }

  async acquireBuildReconciliationLease(
    id: string,
    holder: string,
    ttlMs: number,
  ): Promise<BuildReconciliationLease | undefined> {
    if (!holder || !Number.isFinite(ttlMs) || ttlMs <= 0)
      throw new BuildStoreConflictError("Invalid build reconciliation lease");
    const current = this.builds.get(id);
    if (!current) return;
    if (
      current.status.reconcileLease &&
      Date.parse(current.status.reconcileLease.expiresAt) > Date.now()
    ) {
      return;
    }
    const lease: BuildReconciliationLease = {
      holder,
      token: randomUUID(),
      expiresAt: new Date(Date.now() + ttlMs).toISOString(),
    };
    const next = clone(current);
    next.metadata.resourceVersion = String(
      Number(current.metadata.resourceVersion) + 1,
    );
    next.status.reconcileLease = lease;
    await this.replaceBuild(next, current.metadata.resourceVersion);
    return clone(lease);
  }

  async renewBuildReconciliationLease(
    id: string,
    token: string,
    ttlMs: number,
  ): Promise<boolean> {
    if (!token || !Number.isFinite(ttlMs) || ttlMs <= 0) return false;
    const current = this.builds.get(id);
    if (
      !current ||
      current.status.reconcileLease?.token !== token ||
      Date.parse(current.status.reconcileLease.expiresAt) <= Date.now()
    ) {
      return false;
    }
    const next = clone(current);
    next.metadata.resourceVersion = String(
      Number(current.metadata.resourceVersion) + 1,
    );
    next.status.reconcileLease!.expiresAt = new Date(
      Date.now() + ttlMs,
    ).toISOString();
    try {
      await this.replaceBuild(next, current.metadata.resourceVersion, token);
      return true;
    } catch (error) {
      if (error instanceof BuildStoreConflictError) return false;
      throw error;
    }
  }

  async releaseBuildReconciliationLease(
    id: string,
    token: string,
  ): Promise<void> {
    const current = this.builds.get(id);
    if (current?.status.reconcileLease?.token !== token) return;
    const next = clone(current);
    next.metadata.resourceVersion = String(
      Number(current.metadata.resourceVersion) + 1,
    );
    delete next.status.reconcileLease;
    await this.replaceBuild(next, current.metadata.resourceVersion, token);
  }

  async ownsBuild(imageKey: string, buildId: string): Promise<boolean> {
    return this.active.get(imageKey) === buildId;
  }

  async getUpload(digest: Sha256Digest) {
    const value = this.uploads.get(digest);
    return value && clone(value);
  }

  async createUpload(record: UploadRecord) {
    const current = this.uploads.get(record.spec.digest);
    if (current) return clone(current);
    this.uploads.set(record.spec.digest, clone(record));
    return clone(record);
  }

  async replaceUpload(record: UploadRecord, expectedResourceVersion: string) {
    const current = this.uploads.get(record.spec.digest);
    if (
      !current ||
      current.metadata.resourceVersion !== expectedResourceVersion
    ) {
      throw new BuildStoreConflictError(
        "Upload record was concurrently modified",
      );
    }
    this.uploads.set(record.spec.digest, clone(record));
  }

  async deleteUpload(digest: Sha256Digest) {
    this.uploads.delete(digest);
  }
}