feat: harden self-managed reconciliation

2026-09-05 10:17:55 +00:00
parent 8e9d207915
commit 321f4e807a
40 changed files with 4977 additions and 404 deletions
@@ -12,7 +12,7 @@ import {
  type Sha256Digest,
 } from "../shared/build-protocol";
 import { resolveComposeArch } from "./arch";
-import { apiRequest, type ApiRequestInit } from "./api";
+import { apiRequest, type ApiRequestInit, type ApiRequestOptions } from "./api";
 import { DEFAULT_REGISTRY } from "./config";
 import {
  enumerateWorkspace,
@@ -23,10 +23,132 @@ import {
 const execFileAsync = promisify(execFile);
 const UPLOAD_CHUNK_BYTES = 8 * 1024 * 1024;
 const DEFAULT_POLL_INTERVAL_MS = 1_000;
+const BUILD_POLL_REQUEST_TIMEOUT_MS = 300_000;
+const MAX_BUILD_POLL_ATTEMPTS = 3;
+
+export const MAX_CONCURRENT_REQUESTS = 20;
+export const MAX_REQUESTS_PER_SECOND = 40;
+
+export type SchedulerClock = {
+  now(): number;
+};
+
+export type SchedulerSleep = (ms: number) => Promise<void>;
+
+export class TaskScheduler {
+  private inflight = 0;
+  private maxInflight: number;
+  private maxPerSecond: number;
+  private requestStarts: number[] = [];
+  private waiting: Array<() => void> = [];
+  private rateGate: Promise<void> = Promise.resolve();
+  private clock: SchedulerClock;
+  private sleep: SchedulerSleep;
+
+  constructor(options?: {
+    maxInflight?: number;
+    maxPerSecond?: number;
+    clock?: SchedulerClock;
+    sleep?: SchedulerSleep;
+  }) {
+    this.maxInflight = options?.maxInflight ?? MAX_CONCURRENT_REQUESTS;
+    this.maxPerSecond = options?.maxPerSecond ?? MAX_REQUESTS_PER_SECOND;
+    if (!Number.isSafeInteger(this.maxInflight) || this.maxInflight < 1)
+      throw new RangeError("maxInflight must be a positive integer");
+    if (!Number.isSafeInteger(this.maxPerSecond) || this.maxPerSecond < 1)
+      throw new RangeError("maxPerSecond must be a positive integer");
+    this.clock = options?.clock ?? { now: () => Date.now() };
+    this.sleep = options?.sleep ?? ((ms) => Bun.sleep(ms));
+  }
+
+  get currentInflight(): number {
+    return this.inflight;
+  }
+
+  get maxConcurrent(): number {
+    return this.maxInflight;
+  }
+
+  get currentRequestStarts(): number {
+    this.pruneOldStarts();
+    return this.requestStarts.length;
+  }
+
+  private pruneOldStarts(): void {
+    const cutoff = this.clock.now() - 1000;
+    while (this.requestStarts.length > 0 && this.requestStarts[0]! <= cutoff) {
+      this.requestStarts.shift();
+    }
+  }
+
+  private async acquire(): Promise<void> {
+    if (this.inflight < this.maxInflight) {
+      this.inflight++;
+      return;
+    }
+    await new Promise<void>((resolve) => {
+      this.waiting.push(resolve);
+    });
+  }
+
+  private release(): void {
+    if (this.waiting.length > 0) {
+      const next = this.waiting.shift()!;
+      next();
+    } else {
+      this.inflight--;
+    }
+  }
+
+  private waitForRateLimit(): Promise<void> {
+    const reservation = this.rateGate.then(async () => {
+      for (;;) {
+        this.pruneOldStarts();
+        if (this.requestStarts.length < this.maxPerSecond) {
+          this.requestStarts.push(this.clock.now());
+          return;
+        }
+        const oldest = this.requestStarts[0]!;
+        await this.sleep(Math.max(1, oldest + 1000 - this.clock.now()));
+      }
+    });
+    this.rateGate = reservation.catch(() => {});
+    return reservation;
+  }
+
+  async run<T>(fn: () => Promise<T>): Promise<T> {
+    await this.acquire();
+    try {
+      await this.waitForRateLimit();
+      return await fn();
+    } finally {
+      this.release();
+    }
+  }
+}
+
+async function runConcurrent<T>(
+  values: T[],
+  concurrency: number,
+  run: (value: T) => Promise<void>,
+): Promise<void> {
+  let index = 0;
+  const worker = async () => {
+    for (;;) {
+      const current = index++;
+      if (current >= values.length) return;
+      await run(values[current]!);
+    }
+  };
+  await Promise.all(
+    Array.from({ length: Math.min(concurrency, values.length) }, worker),
+  );
+}

 export type ApiRequester = <T>(
  path: string,
  init?: ApiRequestInit,
+  options?: ApiRequestOptions,
 ) => Promise<T>;

 export type BuildOptions = {
@@ -35,6 +157,7 @@ export type BuildOptions = {
  pollIntervalMs?: number;
  sleep?: (milliseconds: number) => Promise<void>;
  snapshot?: WorkspaceSnapshot;
+  scheduler?: TaskScheduler;
 };

 type BuildPlan = {
@@ -191,29 +314,42 @@ async function uploadBlob(
  digest: Sha256Digest,
  data: Uint8Array,
  request: ApiRequester,
+  scheduler: TaskScheduler,
+  project?: string,
 ): Promise<void> {
-  const path = `/blobs/${encodeURIComponent(digest)}/uploads`;
-  const progress = await request<{ offset: number; complete: boolean }>(path, {
-    method: "POST",
-    json: { size: data.byteLength },
-  });
+  const uploadPath = `/blobs/${encodeURIComponent(digest)}/uploads`;
+  const projectQuery = project ? `?project=${encodeURIComponent(project)}` : "";
+  const path = `${uploadPath}${projectQuery}`;
+  const progress = await scheduler.run(() =>
+    request<{ offset: number; complete: boolean }>(path, {
+      method: "POST",
+      json: { size: data.byteLength },
+    }),
+  );
  let offset = progress.offset;
  while (!progress.complete && offset < data.byteLength) {
    const chunk = data.subarray(offset, offset + UPLOAD_CHUNK_BYTES);
-    const uploaded = await request<{ offset: number }>(path, {
-      method: "PATCH",
-      headers: {
-        "content-type": "application/octet-stream",
-        "upload-offset": String(offset),
-      },
-      body: chunk,
-    });
+    const uploaded = await scheduler.run(() =>
+      request<{ offset: number }>(path, {
+        method: "PATCH",
+        headers: {
+          "content-type": "application/octet-stream",
+          "upload-offset": String(offset),
+        },
+        body: chunk,
+      }),
+    );
    if (uploaded.offset <= offset)
      throw new Error(`Blob upload for ${digest} made no progress`);
    offset = uploaded.offset;
  }
  if (!progress.complete) {
-    await request(`${path}/complete`, { method: "POST", json: {} });
+    await scheduler.run(() =>
+      request(`${uploadPath}/complete${projectQuery}`, {
+        method: "POST",
+        json: {},
+      }),
+    );
  }
 }

@@ -221,30 +357,36 @@ export async function uploadWorkspaceSnapshot(
  snapshot: WorkspaceSnapshot,
  request: ApiRequester = apiRequest,
  reporter?: BuildReporter,
+  scheduler?: TaskScheduler,
+  project?: string,
 ): Promise<void> {
  const blobs = new Map(snapshot.blobs.map((blob) => [blob.digest, blob.data]));
  blobs.set(snapshot.digest, serializeWorkspaceManifest(snapshot.manifest));
+  const requestScheduler = scheduler ?? new TaskScheduler();

  for (;;) {
-    const negotiation = await request<SnapshotNegotiation>(
-      "/snapshots/negotiate",
-      {
+    const negotiation = await requestScheduler.run(() =>
+      request<SnapshotNegotiation>("/snapshots/negotiate", {
        method: "POST",
-        json: { workspace: snapshot.digest },
-      },
+        json: { workspace: snapshot.digest, ...(project && { project }) },
+      }),
    );
    if (negotiation.ready) return;
    if (negotiation.missing.length === 0)
      throw new Error(
        "Snapshot negotiation is incomplete but reported no missing blobs",
      );
-    for (const digest of negotiation.missing) {
-      const data = blobs.get(digest);
-      if (!data)
-        throw new Error(`Server requested unknown workspace blob ${digest}`);
-      await reporter?.progress?.(`Uploading ${digest}`);
-      await uploadBlob(digest, data, request);
-    }
+    await runConcurrent(
+      negotiation.missing,
+      requestScheduler.maxConcurrent,
+      async (digest) => {
+        const data = blobs.get(digest);
+        if (!data)
+          throw new Error(`Server requested unknown workspace blob ${digest}`);
+        await reporter?.progress?.(`Uploading ${digest}`);
+        await uploadBlob(digest, data, request, requestScheduler, project);
+      },
+    );
  }
 }

@@ -265,6 +407,52 @@ async function reportBuildEvent(
  return event.sequence;
 }

+function isTransientBuildPollError(error: unknown): boolean {
+  if (!(error instanceof Error) || error.name === "AbortError") return false;
+  if (error.name === "TimeoutError" || error instanceof TypeError) return true;
+  const code =
+    "code" in error && typeof error.code === "string"
+      ? error.code
+      : error.cause &&
+          typeof error.cause === "object" &&
+          "code" in error.cause &&
+          typeof error.cause.code === "string"
+        ? error.cause.code
+        : undefined;
+  return (
+    code === "ECONNABORTED" ||
+    code === "ECONNRESET" ||
+    code === "ECONNREFUSED" ||
+    code === "EAI_AGAIN" ||
+    code === "ETIMEDOUT"
+  );
+}
+
+async function requestBuildPoll<T>(
+  request: ApiRequester,
+  path: string,
+  init: ApiRequestInit | undefined,
+  pollIntervalMs: number,
+  sleep: (milliseconds: number) => Promise<void>,
+): Promise<T> {
+  for (let attempt = 1; attempt <= MAX_BUILD_POLL_ATTEMPTS; attempt++) {
+    try {
+      return await request<T>(path, init, {
+        timeoutMs: BUILD_POLL_REQUEST_TIMEOUT_MS,
+      });
+    } catch (error) {
+      if (
+        !isTransientBuildPollError(error) ||
+        attempt === MAX_BUILD_POLL_ATTEMPTS
+      ) {
+        throw error;
+      }
+      await sleep(pollIntervalMs);
+    }
+  }
+  throw new Error("Build poll retries exhausted");
+}
+
 async function waitForBuild(
  id: string,
  request: ApiRequester,
@@ -277,8 +465,12 @@ async function waitForBuild(
  let sequence = 0;
  const reportedStates = new Set<BuildStatus["state"]>();
  for (;;) {
-    const events = await request<BuildEvent[]>(
+    const events = await requestBuildPoll<BuildEvent[]>(
+      request,
      `/builds/${encodeURIComponent(id)}/events?after=${sequence}`,
+      undefined,
+      pollIntervalMs,
+      sleep,
    );
    for (const event of events)
      sequence = Math.max(
@@ -287,12 +479,15 @@ async function waitForBuild(
      );
    if (status.state === "succeeded" || status.state === "failed")
      return status;
-    status = await request<BuildStatus>(
+    status = await requestBuildPoll<BuildStatus>(
+      request,
      `/builds/${encodeURIComponent(id)}/reconcile`,
      {
        method: "POST",
        json: {},
      },
+      pollIntervalMs,
+      sleep,
    );
    if (status.state !== "succeeded" && status.state !== "failed")
      await sleep(pollIntervalMs);
@@ -351,7 +546,13 @@ export async function buildServices(
      return plan ? [plan] : [];
    },
  );
-  await uploadWorkspaceSnapshot(snapshot, request, reporter);
+  await uploadWorkspaceSnapshot(
+    snapshot,
+    request,
+    reporter,
+    options.scheduler,
+    project,
+  );

  const images: Record<string, string> = {};
  for (const plan of plans) {
@@ -372,10 +573,14 @@ export async function buildServices(
        workspace: snapshot.digest,
      },
    };
-    const initial = await request<BuildStatus>("/builds", {
-      method: "POST",
-      json: buildRequest,
-    });
+    const initial = await request<BuildStatus>(
+      "/builds",
+      {
+        method: "POST",
+        json: buildRequest,
+      },
+      { timeoutMs: 300_000 },
+    );
    const status = await waitForBuild(
      id,
      request,