From 847399a4ffe348e8a9eb8e2108b8ba79c5b21f0b Mon Sep 17 00:00:00 2001 From: Josh Wright Date: Mon, 5 Oct 2026 15:47:25 -0400 Subject: [PATCH] Throttle agent usage writes and cache inventory per request --- src/hooks.server.ts | 5 +- src/lib/server/hosts/index.ts | 13 ++- src/lib/server/hosts/request-context.ts | 28 +++++ tests/unit/host-request-cost.test.ts | 146 ++++++++++++++++++++++++ 4 files changed, 187 insertions(+), 5 deletions(-) create mode 100644 src/lib/server/hosts/request-context.ts create mode 100644 tests/unit/host-request-cost.test.ts diff --git a/src/hooks.server.ts b/src/hooks.server.ts index 4717552..b7b9e8c 100644 --- a/src/hooks.server.ts +++ b/src/hooks.server.ts @@ -6,6 +6,7 @@ import { authenticationHandle } from '$lib/server/auth/guard.js'; import { invalidateFleet } from '$lib/server/pipeline/bus.js'; import { startPerformanceMonitor } from '$lib/server/performance.js'; import { withOperationSource } from '$lib/server/hosts/operation-context.js'; +import { withHostRequestContext } from '$lib/server/hosts/request-context.js'; let booted = false; function boot(): void { @@ -23,8 +24,8 @@ function boot(): void { export const handle: Handle = async (input) => { boot(); const started = performance.now(); - const response = await withOperationSource(`${input.event.request.method} ${input.event.route.id ?? 'Web request'}`, - () => authenticationHandle(input)); + const response = await withHostRequestContext(() => withOperationSource(`${input.event.request.method} ${input.event.route.id ?? 'Web request'}`, + () => authenticationHandle(input))); const elapsed = performance.now() - started; 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/')) { diff --git a/src/lib/server/hosts/index.ts b/src/lib/server/hosts/index.ts index 453550e..714878e 100644 --- a/src/lib/server/hosts/index.ts +++ b/src/lib/server/hosts/index.ts @@ -6,6 +6,7 @@ import { db } from '../db/index.js'; import { ensureHostSchema } from './schema.js'; import { successfulObservation, routineObservationSql, importantOperationSql, operationTargetSql, safeOperationTarget } from './observations.js'; import { currentOperationSource } from './operation-context.js'; +import { parseHostInventory } from './request-context.js'; import { fleetSettings } from '../settings.js'; import { recordInstanceResources, validateInstanceResources } from './calibration.js'; import { refreshHostAvailability, scheduleAvailabilityRefresh } from './availability.js'; @@ -38,7 +39,7 @@ function presentDecommission(row: DecommissionRow | undefined): HostDecommission } 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 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 reasons: string[] = []; 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 { 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; - 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); } export function revokeCredentials(hostId: string): void { diff --git a/src/lib/server/hosts/request-context.ts b/src/lib/server/hosts/request-context.ts new file mode 100644 index 0000000..4f3aa63 --- /dev/null +++ b/src/lib/server/hosts/request-context.ts @@ -0,0 +1,28 @@ +import { AsyncLocalStorage } from 'node:async_hooks'; +import type { HostInventory } from '../../hosts.js'; + +type InventoryCache = Map; +const requests = new AsyncLocalStorage<{ inventories: InventoryCache | null }>(); + +export function withHostRequestContext(work: () => T | Promise): Promise { + 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; +} diff --git a/tests/unit/host-request-cost.test.ts b/tests/unit/host-request-cost.test.ts new file mode 100644 index 0000000..1d5e97f --- /dev/null +++ b/tests/unit/host-request-cost.test.ts @@ -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(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(resolve => { release = resolve; }); + let detached!: Promise; + 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; + }); +});