@zgeoff/atc 2.24.0 → 2.25.0

package/src/daemon/imp-provider.ts

@@ -1,6 +1,7 @@
 import { DaemonError } from '../protocol/daemon-error';
 import { LineDecoder } from '../protocol/line-decoder';
 import { isCompiledBinary } from '../shared/is-compiled-binary';
+import type { BrokerAuthHost } from './broker-auth-host';
 import { buildImpName } from './build-imp-name';
 import { buildTarArchive } from './build-tar-archive';
 import type {
@@ -96,6 +97,8 @@ export class ImpProvider implements ExecutionProvider {
   // namespace an impd token's imp patterns are checked against.
   readonly impPrefix: string;
 
+  readonly brokerAuth: BrokerAuthHost;
+
   private readonly port: ImpPort;
 
   private readonly target: ImpTargetOptions;
@@ -106,6 +109,14 @@ export class ImpProvider implements ExecutionProvider {
 
   private readonly hosts = new Map<string, ImpHost>();
 
+  // Each imp's readyings, sleeps, and lease returns run one after another,
+  // in the order they were asked for, so none of them gives back the lease
+  // a readying just took.
+  private readonly turns = new Map<string, Promise<void>>();
+
+  // The readyings of each imp still waiting for their turn.
+  private readonly waitingPrepares = new Map<string, number>();
+
   private readonly atcBinary: string | null;
 
   // Whether impd carries output offsets, once a prepare has read its
@@ -132,6 +143,19 @@ export class ImpProvider implements ExecutionProvider {
       dir,
       atc: target.guestATC ?? (this.atcBinary === null ? null : `${dir}/bin/atc`),
     };
+
+    this.brokerAuth = {
+      impPrefix: this.impPrefix,
+      port,
+      getImpName: (hostKey) => this.getImpName(hostKey),
+      createImp: (hostKey) =>
+        port.createImp({
+          name: this.getImpName(hostKey),
+          ...(target.image === undefined ? {} : { image: target.image }),
+          ...(target.memoryMib === undefined ? {} : { memoryMib: target.memoryMib }),
+        }),
+      destroyImp: (hostKey) => this.destroyHost(hostKey),
+    };
   }
 
   // The imp a host key runs on, under this provider's prefix.
@@ -141,7 +165,25 @@ export class ImpProvider implements ExecutionProvider {
 
   // Creates the host's imp when impd holds none, then takes the daemon's
   // lease, which boots or wakes the imp.
-  readonly prepareHost = async (request: HostRequest): Promise<void> => {
+  readonly prepareHost = (request: HostRequest): Promise<void> => {
+    const name = this.getImpName(request.host);
+
+    this.waitingPrepares.set(name, (this.waitingPrepares.get(name) ?? 0) + 1);
+
+    return this.withHostTurn(name, () => {
+      const left = (this.waitingPrepares.get(name) ?? 1) - 1;
+
+      if (left === 0) {
+        this.waitingPrepares.delete(name);
+      } else {
+        this.waitingPrepares.set(name, left);
+      }
+
+      return this.prepareHostInTurn(request);
+    });
+  };
+
+  private readonly prepareHostInTurn = async (request: HostRequest): Promise<void> => {
     const name = this.getImpName(request.host);
     const label = `atc-${request.daemonID}`;
 
@@ -187,9 +229,11 @@ export class ImpProvider implements ExecutionProvider {
       harnesses: 0,
       renewTimer: null,
       suspending: false,
+      isIdle: null,
     };
 
     host.suspending = false;
+    host.isIdle = request.isIdle ?? null;
 
     this.hosts.set(request.host, host);
     this.startRenewal(host, label);
@@ -201,6 +245,7 @@ export class ImpProvider implements ExecutionProvider {
       harnesses: 0,
       renewTimer: null,
       suspending: false,
+      isIdle: null,
     };
 
     this.hosts.set(spec.host, host);
@@ -225,11 +270,13 @@ export class ImpProvider implements ExecutionProvider {
         cwd: spec.cwd,
         cols: spec.cols,
         rows: spec.rows,
+        ...(spec.requireBroker === true ? { require: ['broker'] } : {}),
       },
       {
         offsets: this.offsets,
         reconnectDelaysMs: this.reconnectDelaysMs,
         ...(reporter === null ? {} : { ready: reporter.forward.listening }),
+        ...(spec.admit === undefined ? {} : { admit: spec.admit }),
         isSuspending: () => host.suspending,
         onDone: () => {
           reporter?.forward.stop();
@@ -271,37 +318,57 @@ export class ImpProvider implements ExecutionProvider {
   // then asks impd to sleep the imp without force. Another owner's lease
   // refuses the sleep: the daemon takes its lease back and the refusal
   // carries what impd showed of the other leases.
-  readonly suspendHost = async (hostKey: string): Promise<void> => {
-    const host = this.hosts.get(hostKey) ?? {
-      name: this.getImpName(hostKey),
-      harnesses: 0,
-      renewTimer: null,
-      suspending: false,
-    };
+  readonly suspendHost = (hostKey: string, isIdle?: () => boolean): Promise<void> =>
+    this.withHostTurn(this.getImpName(hostKey), async () => {
+      const host = this.hosts.get(hostKey) ?? {
+        name: this.getImpName(hostKey),
+        harnesses: 0,
+        renewTimer: null,
+        suspending: false,
+        isIdle: null,
+      };
+
+      this.hosts.set(hostKey, host);
+
+      // A host a harness runs on, or that a readying waits for, is busy.
+      const isBusy = () =>
+        isIdle !== undefined &&
+        (!isIdle() || host.harnesses > 0 || (this.waitingPrepares.get(host.name) ?? 0) > 0);
+
+      if (isBusy()) {
+        throw buildBusyRefusal(host.name);
+      }
 
-    this.hosts.set(hostKey, host);
-    this.stopRenewal(host);
+      this.stopRenewal(host);
 
-    host.suspending = true;
+      host.suspending = true;
 
-    try {
-      if (this.label !== null) {
-        await this.port.releaseLease(host.name, this.label);
-      }
+      try {
+        if (this.label !== null) {
+          await this.port.releaseLease(host.name, this.label);
+        }
 
-      await this.port.suspendImp(host.name);
-    } catch (error) {
-      host.suspending = false;
+        if (isBusy()) {
+          host.suspending = false;
 
-      if (error instanceof ImpPortError && error.code === 'LEASED') {
-        await this.restoreLease(host);
+          await this.restoreLease(host);
 
-        throw buildLeasedRefusal(host.name, error.data);
-      }
+          throw buildBusyRefusal(host.name);
+        }
 
-      throw toHostRefusal(error, host.name);
-    }
-  };
+        await this.port.suspendImp(host.name);
+      } catch (error) {
+        host.suspending = false;
+
+        if (error instanceof ImpPortError && error.code === 'LEASED') {
+          await this.restoreLease(host);
+
+          throw buildLeasedRefusal(host.name, error.data);
+        }
+
+        throw toHostRefusal(error, host.name);
+      }
+    });
 
   // A destroy ends every lease on the imp with it.
   readonly destroyHost = async (hostKey: string): Promise<void> => {
@@ -437,13 +504,7 @@ export class ImpProvider implements ExecutionProvider {
 
     host.renewTimer = setInterval(
       () => {
-        if (host.harnesses === 0) {
-          void this.tryReleaseLease(host);
-
-          return;
-        }
-
-        void this.tryRenewLease(host, label);
+        void this.refreshLease(host, label);
       },
       (this.leaseSeconds * 1000) / 3,
     );
@@ -473,6 +534,28 @@ export class ImpProvider implements ExecutionProvider {
     }
   }
 
+  // Runs a readying, a sleep, or a lease return of an imp after the one
+  // before it there. A turn never starts in its caller's tick, so a harness
+  // end the daemon still applies counts before the turn checks the host.
+  private async withHostTurn<T>(name: string, run: () => Promise<T>): Promise<T> {
+    const before = this.turns.get(name) ?? Promise.resolve();
+    const turn = Promise.withResolvers<void>();
+
+    this.turns.set(name, turn.promise);
+
+    try {
+      await before;
+
+      return await run();
+    } finally {
+      turn.resolve();
+
+      if (this.turns.get(name) === turn.promise) {
+        this.turns.delete(name);
+      }
+    }
+  }
+
   private async restoreLease(host: ImpHost): Promise<void> {
     if (this.label === null) {
       return;
@@ -487,19 +570,45 @@ export class ImpProvider implements ExecutionProvider {
     this.startRenewal(host, this.label);
   }
 
-  private async tryReleaseLease(host: ImpHost): Promise<boolean> {
-    this.stopRenewal(host);
+  // A host with no harness gives its lease back, unless a launch readies it,
+  // which keeps the lease renewed.
+  private async refreshLease(host: ImpHost, label: string): Promise<void> {
+    if (host.harnesses === 0) {
+      const released = await this.tryReleaseLease(host);
 
-    if (this.label === null) {
-      return false;
+      if (released) {
+        return;
+      }
     }
 
-    try {
-      return await this.port.releaseLease(host.name, this.label);
-    } catch {
-      return false;
+    if (host.renewTimer !== null) {
+      await this.tryRenewLease(host, label);
     }
   }
+
+  // Gives an imp's lease back once nothing runs or starts there, in the
+  // imp's turn, so it never drops a lease a readying took for a launch.
+  private tryReleaseLease(host: ImpHost): Promise<boolean> {
+    return this.withHostTurn(host.name, async () => {
+      const isBusy =
+        host.harnesses > 0 ||
+        host.suspending ||
+        (this.waitingPrepares.get(host.name) ?? 0) > 0 ||
+        host.isIdle?.() === false;
+
+      if (isBusy || this.label === null) {
+        return false;
+      }
+
+      this.stopRenewal(host);
+
+      try {
+        return await this.port.releaseLease(host.name, this.label);
+      } catch {
+        return false;
+      }
+    });
+  }
 }
 
 // What the provider follows of one imp: its name, how many harnesses run
@@ -510,6 +619,10 @@ interface ImpHost {
   harnesses: number;
   renewTimer: ReturnType<typeof setInterval> | null;
   suspending: boolean;
+
+  // The daemon's check that nothing runs or starts on the host, from its
+  // latest readying; null until one gives it.
+  isIdle: (() => boolean) | null;
 }
 
 function buildImpSessionName(sessionID: string): string {
@@ -571,6 +684,14 @@ function buildGuestEnv(env: Readonly<Record<string, string>>): Record<string, st
   return { TERM: 'xterm-256color', LANG: 'C.UTF-8', PATH: GUEST_PATH, ...own };
 }
 
+function buildBusyRefusal(name: string): DaemonError {
+  return new DaemonError(
+    'host_unavailable',
+    `imp ${name} stays awake: a session runs there or is starting there`,
+    { provider: 'imp', problem: 'busy' },
+  );
+}
+
 function buildLeasedRefusal(name: string, data: unknown): DaemonError {
   const record = typeof data === 'object' && data !== null ? data : {};
   const leases = 'leases' in record && Array.isArray(record.leases) ? record.leases : [];

package/src/daemon/local-pty-provider.ts

@@ -1,5 +1,6 @@
 import { mkdir } from 'node:fs/promises';
 import { spawn } from 'bun-pty';
+import { DaemonError } from '../protocol/daemon-error';
 import { collectCleanEnv } from '../shared/collect-clean-env';
 import type {
   CommandResult,
@@ -39,7 +40,17 @@ export class LocalPTYProvider implements ExecutionProvider {
 
   // The daemon's machine keeps no process the daemon lets go of, so a
   // detach ends the harness as a kill does, after dropping every listener.
+  // It has no credential broker, so a harness that requires one never
+  // starts here.
   readonly spawnHarness = (spec: HarnessSpec): HarnessHandle => {
+    if (spec.requireBroker === true) {
+      throw new DaemonError(
+        'auth_target_unsupported',
+        "the daemon's own machine has no credential broker to start the harness behind",
+        { provider: 'local-pty' },
+      );
+    }
+
     const pty = spawn(spec.bin, [...spec.args], {
       name: 'xterm-256color',
       cols: spec.cols,

package/src/daemon/restore-fleet.ts

@@ -185,7 +185,8 @@ function isSameSession(s: Session, entry: FleetEntry): boolean {
 
 // The refusals that leave one session without a terminal: its target is
 // misconfigured, gone from the config, changed, has no provider here, or
-// cannot start a terminal, or its agent refuses every start.
+// cannot start a terminal, or its agent refuses every start, or its
+// host's runtime auth binding refuses the launch, its broker included.
 const REFUSED_ADOPT_CODES: ReadonlySet<string> = new Set([
   'unsupported_operation',
   'unknown_target',
@@ -195,6 +196,19 @@ const REFUSED_ADOPT_CODES: ReadonlySet<string> = new Set([
   'host_unavailable',
   'auth_not_configured',
   'auth_target_unsupported',
+  'auth_impd_too_old',
+  'auth_token_scope',
+  'auth_token_too_broad',
+  'auth_imp_out_of_scope',
+  'auth_secret_not_grantable',
+  'auth_secret_mismatch',
+  'auth_runtime_mismatch',
+  'auth_grant_missing',
+  'auth_grants_mismatch',
+  'auth_rebind_required',
+  'auth_blocked',
+  'auth_binding_invalid',
+  'broker_not_ready',
 ]);
 
 async function tryAdoptTerminal(