Throttle agent usage writes and cache inventory per request

This commit is contained in:
2026-10-05 15:47:25 -04:00
parent a86499a0f9
commit 847399a4ff
4 changed files with 187 additions and 5 deletions
+3 -2
View File
@@ -6,6 +6,7 @@ import { authenticationHandle } from '$lib/server/auth/guard.js';
import { invalidateFleet } from '$lib/server/pipeline/bus.js'; import { invalidateFleet } from '$lib/server/pipeline/bus.js';
import { startPerformanceMonitor } from '$lib/server/performance.js'; import { startPerformanceMonitor } from '$lib/server/performance.js';
import { withOperationSource } from '$lib/server/hosts/operation-context.js'; import { withOperationSource } from '$lib/server/hosts/operation-context.js';
import { withHostRequestContext } from '$lib/server/hosts/request-context.js';
let booted = false; let booted = false;
function boot(): void { function boot(): void {
@@ -23,8 +24,8 @@ function boot(): void {
export const handle: Handle = async (input) => { export const handle: Handle = async (input) => {
boot(); boot();
const started = performance.now(); const started = performance.now();
const response = await withOperationSource(`${input.event.request.method} ${input.event.route.id ?? 'Web request'}`, const response = await withHostRequestContext(() => withOperationSource(`${input.event.request.method} ${input.event.route.id ?? 'Web request'}`,
() => authenticationHandle(input)); () => authenticationHandle(input)));
const elapsed = performance.now() - started; const elapsed = performance.now() - started;
response.headers.append('Server-Timing', `app;dur=${elapsed.toFixed(1)}`); response.headers.append('Server-Timing', `app;dur=${elapsed.toFixed(1)}`);
if (process.env.EC_PERF_LOG === 'true' && elapsed > 250 && !input.event.route.id?.startsWith('/api/agent/')) { if (process.env.EC_PERF_LOG === 'true' && elapsed > 250 && !input.event.route.id?.startsWith('/api/agent/')) {
+10 -3
View File
@@ -6,6 +6,7 @@ import { db } from '../db/index.js';
import { ensureHostSchema } from './schema.js'; import { ensureHostSchema } from './schema.js';
import { successfulObservation, routineObservationSql, importantOperationSql, operationTargetSql, safeOperationTarget } from './observations.js'; import { successfulObservation, routineObservationSql, importantOperationSql, operationTargetSql, safeOperationTarget } from './observations.js';
import { currentOperationSource } from './operation-context.js'; import { currentOperationSource } from './operation-context.js';
import { parseHostInventory } from './request-context.js';
import { fleetSettings } from '../settings.js'; import { fleetSettings } from '../settings.js';
import { recordInstanceResources, validateInstanceResources } from './calibration.js'; import { recordInstanceResources, validateInstanceResources } from './calibration.js';
import { refreshHostAvailability, scheduleAvailabilityRefresh } from './availability.js'; import { refreshHostAvailability, scheduleAvailabilityRefresh } from './availability.js';
@@ -38,7 +39,7 @@ function presentDecommission(row: DecommissionRow | undefined): HostDecommission
} }
function hostFromRow(row: HostRow): Host { function hostFromRow(row: HostRow): Host {
const enabling = database().prepare('SELECT requested_at FROM ec_host_enable_requests WHERE host_id=?').get(row.id) as { requested_at: string } | undefined; const enabling = database().prepare('SELECT requested_at FROM ec_host_enable_requests WHERE host_id=?').get(row.id) as { requested_at: string } | undefined;
const inventory: HostInventory = JSON.parse(row.inventory); const inventory = parseHostInventory(row.id, row.inventory);
const connected = !!row.credential_active && !!row.last_heartbeat && Date.now() - Date.parse(row.last_heartbeat) <= HEARTBEAT_MAX_AGE_MS; const connected = !!row.credential_active && !!row.last_heartbeat && Date.now() - Date.parse(row.last_heartbeat) <= HEARTBEAT_MAX_AGE_MS;
const reasons: string[] = []; const reasons: string[] = [];
if (!row.enrolled_at) reasons.push('Agent has not enrolled.'); if (!row.enrolled_at) reasons.push('Agent has not enrolled.');
@@ -304,9 +305,15 @@ export function enrollHost(token: unknown, rawInventory: unknown): { hostId: str
} }
export function authenticateCredential(token: string): Host | null { export function authenticateCredential(token: string): Host | null {
if (!token || token.length > 200) return null; if (!token || token.length > 200) return null;
const row = database().prepare('SELECT id,host_id FROM ec_host_credentials WHERE token_hash=? AND revoked_at IS NULL').get(digest(token)) as { id: string; host_id: string } | undefined; const row = database().prepare('SELECT id,host_id,last_used_at FROM ec_host_credentials WHERE token_hash=? AND revoked_at IS NULL')
.get(digest(token)) as { id: string; host_id: string; last_used_at: string | null } | undefined;
if (!row) return null; if (!row) return null;
database().prepare('UPDATE ec_host_credentials SET last_used_at=? WHERE id=?').run(now(), row.id); const timestamp = Date.now();
if (row.last_used_at === null || timestamp - Date.parse(row.last_used_at) >= 60_000) {
database().prepare(`UPDATE ec_host_credentials SET last_used_at=? WHERE id=? AND revoked_at IS NULL
AND (last_used_at IS NULL OR last_used_at<=?)`)
.run(new Date(timestamp).toISOString(), row.id, new Date(timestamp - 60_000).toISOString());
}
return getHost(row.host_id); return getHost(row.host_id);
} }
export function revokeCredentials(hostId: string): void { export function revokeCredentials(hostId: string): void {
+28
View File
@@ -0,0 +1,28 @@
import { AsyncLocalStorage } from 'node:async_hooks';
import type { HostInventory } from '../../hosts.js';
type InventoryCache = Map<string, { json: string; inventory: HostInventory }>;
const requests = new AsyncLocalStorage<{ inventories: InventoryCache | null }>();
export function withHostRequestContext<T>(work: () => T | Promise<T>): Promise<T> {
const context = { inventories: new Map() as InventoryCache | null };
return requests.run(context, async () => {
try { return await work(); }
finally {
// Detached work can inherit this context. Release request data and
// stop caching when the response is ready, even if that work continues.
context.inventories = null;
}
});
}
export function parseHostInventory(hostId: string, json: string): HostInventory {
const cache = requests.getStore()?.inventories;
const saved = cache?.get(hostId);
if (saved?.json === json) return saved.inventory;
const inventory: HostInventory = JSON.parse(json);
// Compare the stored JSON so heartbeats and transaction rollbacks within
// a request remain visible. Other host state is always read afresh.
cache?.set(hostId, { json, inventory });
return inventory;
}
+146
View File
@@ -0,0 +1,146 @@
import { setImmediate } from 'node:timers/promises';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import type { HostInventory } from '../../src/lib/hosts.js';
import { newFixture, type Fixture } from '../helpers/fixture.js';
let fixture: Fixture;
const epoch = Date.parse('2026-10-05T12:00:00Z');
const inventory = { hostname: 'request-cache-host', protocolVersion: 1, dockerRunning: true, directoriesWritable: true,
storageAvailable: 100 * 1024 ** 3, runtimeImages: ['plutainer:test'], libraries: [{ id: 'test-library', state: 'ready' }] };
beforeEach(async () => {
vi.resetModules(); vi.useFakeTimers({ toFake: ['Date'] }); vi.setSystemTime(epoch);
fixture = await newFixture([]);
});
afterEach(async () => { vi.restoreAllMocks(); vi.useRealTimers(); await fixture.cleanup(); });
async function enrolled() {
const hosts = await import('../../src/lib/server/hosts/index.js');
const context = await import('../../src/lib/server/hosts/request-context.js');
const host = hosts.createHost({ name: 'Request cache host' });
const { credential } = hosts.enrollHost(hosts.createEnrollment(host.id).token, inventory);
// Finish enrollment's independent availability refresh before counting parses.
await setImmediate();
return { ...hosts, ...context, hostId: host.id, credential };
}
describe('agent credential usage writes', () => {
it('writes on first use and at most once per minute, including after reopening the database', async () => {
const api = await enrolled();
const { db, closeDb } = await import('../../src/lib/server/db/index.js');
db().exec(`CREATE TABLE credential_usage_writes (credential_id TEXT, at TEXT);
CREATE TRIGGER count_credential_usage AFTER UPDATE OF last_used_at ON ec_host_credentials
BEGIN INSERT INTO credential_usage_writes VALUES (NEW.id, NEW.last_used_at); END;`);
const writes = () => db().prepare('SELECT at FROM credential_usage_writes').all();
const authenticate = () => expect(api.authenticateCredential(api.credential)?.id).toBe(api.hostId);
authenticate(); authenticate(); authenticate();
expect(writes()).toEqual([{ at: new Date(epoch).toISOString() }]);
vi.setSystemTime(epoch + 59_999);
authenticate();
closeDb(); authenticate();
expect(writes()).toHaveLength(1);
vi.setSystemTime(epoch + 60_000);
authenticate(); authenticate();
expect(writes()).toEqual([{ at: new Date(epoch).toISOString() }, { at: new Date(epoch + 60_000).toISOString() }]);
expect(api.authenticateCredential('invalid')).toBeNull();
expect(writes()).toHaveLength(2);
});
it('checks revocation immediately and records first use of a replacement credential', async () => {
const api = await enrolled();
const { db } = await import('../../src/lib/server/db/index.js');
await api.withHostRequestContext(() => {
expect(api.authenticateCredential(api.credential)?.id).toBe(api.hostId);
vi.setSystemTime(epoch + 1000);
const replacement = api.rotateCredential(api.hostId);
expect(api.authenticateCredential(api.credential)).toBeNull();
expect(api.authenticateCredential(replacement)?.id).toBe(api.hostId);
expect(db().prepare('SELECT last_used_at FROM ec_host_credentials WHERE revoked_at IS NULL').get())
.toEqual({ last_used_at: new Date(epoch + 1000).toISOString() });
api.revokeCredentials(api.hostId);
expect(api.authenticateCredential(replacement)).toBeNull();
expect(api.getHost(api.hostId)?.credentialActive).toBe(false);
});
});
});
describe('request-scoped host inventory parsing', () => {
it('parses once across authentication, host reads, and operation handling within a request', async () => {
const api = await enrolled();
const serialized = JSON.stringify(inventory);
const parse = vi.spyOn(JSON, 'parse');
await api.withHostRequestContext(async () => {
const authenticated = api.authenticateCredential(api.credential)!;
const operation = api.queueHostOperation(api.hostId, 'instance.inspect', { name: 'server' });
api.pendingHostOperationRows(api.hostId, 'read', false);
api.reportHostOperationStatus(api.hostId, operation.id, { state: 'accepted' });
await Promise.resolve();
for (let i = 0; i < 10; i++) expect(api.getHost(api.hostId)?.inventory).toBe(authenticated.inventory);
expect(api.requireReadyHost(api.hostId).inventory).toBe(authenticated.inventory);
expect(api.listHosts()[0].inventory).toBe(authenticated.inventory);
});
expect(parse.mock.calls.filter(([value]) => value === serialized)).toHaveLength(1);
await api.withHostRequestContext(() => api.getHost(api.hostId));
expect(parse.mock.calls.filter(([value]) => value === serialized)).toHaveLength(2);
});
it('keeps host metadata fresh and reparses changed inventory, including after rollback', async () => {
const api = await enrolled();
const { db } = await import('../../src/lib/server/db/index.js');
await api.withHostRequestContext(() => {
const original = api.getHost(api.hostId)!;
const updated = api.updateHost(api.hostId, { name: 'Renamed host', mode: 'draining' });
expect(updated).toMatchObject({ name: 'Renamed host', mode: 'draining', readiness: 'draining' });
expect(updated.inventory).toBe(original.inventory);
const heartbeat = api.heartbeat(api.hostId, { ...inventory, storageAvailable: 123 });
expect(heartbeat.inventory).not.toBe(original.inventory);
expect(api.getHost(api.hostId)?.inventory).toBe(heartbeat.inventory);
expect(api.getHost(api.hostId)?.inventory.storageAvailable).toBe(123);
expect(() => db().transaction(() => {
db().prepare('UPDATE ec_hosts SET inventory=? WHERE id=?').run(JSON.stringify({ ...inventory, storageAvailable: 456 }), api.hostId);
expect(api.getHost(api.hostId)?.inventory.storageAvailable).toBe(456);
throw new Error('rollback');
})()).toThrow('rollback');
expect(api.getHost(api.hostId)?.inventory.storageAvailable).toBe(123);
});
});
it('isolates overlapping requests and does not cache reads outside a request', async () => {
const api = await enrolled();
let release!: () => void;
const gate = new Promise<void>(resolve => { release = resolve; });
let first!: HostInventory;
const request = api.withHostRequestContext(async () => {
first = api.getHost(api.hostId)!.inventory;
await gate;
expect(api.getHost(api.hostId)!.inventory).toBe(first);
});
try {
await api.withHostRequestContext(async () => {
const second = api.getHost(api.hostId)!.inventory;
expect(second).not.toBe(first);
await Promise.resolve();
expect(api.getHost(api.hostId)!.inventory).toBe(second);
});
expect(api.getHost(api.hostId)!.inventory).not.toBe(api.getHost(api.hostId)!.inventory);
} finally { release(); await request; }
});
it.each([false, true])('releases cached inventory from inherited work after request completion (failed: %s)', async failed => {
const api = await enrolled();
let release!: () => void;
const gate = new Promise<void>(resolve => { release = resolve; });
let detached!: Promise<void>;
let original!: HostInventory;
const request = api.withHostRequestContext(() => {
original = api.getHost(api.hostId)!.inventory;
detached = gate.then(() => {
const later = api.getHost(api.hostId)!.inventory;
expect(later).not.toBe(original);
expect(api.getHost(api.hostId)!.inventory).not.toBe(later);
});
if (failed) throw new Error('request failed');
});
if (failed) await expect(request).rejects.toThrow('request failed'); else await request;
release(); await detached;
});
});