Skip to content

Commit ac161ec

Browse files
committed
fix(isolated-vm): address review on host drain
- Keep counting a refused connection's bytes until its socket closes - Skip a redispatch the caller cancelled in between - Look up the owner's state after the admission waits - Parse IVM_REMOTE_WORKER_MAX_AGE_MS strictly, falling back below 1s - Hold the inspected execution open until the worker check finishes
1 parent d686032 commit ac161ec

4 files changed

Lines changed: 99 additions & 53 deletions

File tree

‎apps/sim/lib/core/config/env.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -459,7 +459,7 @@ export const env = createEnv({
459459
IVM_MAX_BROKERS_PER_EXECUTION: z.string().optional().default('1000'), // Max broker calls per sandbox task execution
460460
IVM_REMOTE_WORKER_URL: z.string().optional(), // ws(s)://<host>/v1/worker of a dedicated worker host; workers run there unless the isolated-vm-local-workers kill switch is on
461461
IVM_REMOTE_WORKER_TOKEN: z.string().optional(), // Token the worker host expects (its IVM_HOST_TOKEN, >= 32 chars); required with IVM_REMOTE_WORKER_URL
462-
IVM_REMOTE_WORKER_MAX_AGE_MS: z.string().optional().default('120000'), // Age at which a worker host connection stops taking executions, so host deploys drain; keep below the host's deregistration delay
462+
IVM_REMOTE_WORKER_MAX_AGE_MS: z.coerce.number().int().positive().optional(), // (default 120000) Age at which a worker host connection stops taking executions, so host deploys drain; keep below the host's deregistration delay
463463

464464
// Knowledge Base Processing Configuration - Shared across all processing methods
465465
KB_CONFIG_MAX_DURATION: z.number().optional().default(600), // Max processing duration in seconds (10 minutes)

‎apps/sim/lib/execution/isolated-vm-host.cjs‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -195,24 +195,26 @@ server.on('upgrade', (req, socket, head) => {
195195
function authenticate(ws, socket) {
196196
pendingConnections++
197197
let settled = false
198+
let authenticated = false
198199
let preAuthBytes = 0
199200
// Counts what arrives on the TCP socket, so an oversized frame is cut off as it streams in
200201
// rather than after ws has buffered all of it.
201202
const onPreAuthData = (chunk) => {
202203
// ws parses each chunk before this listener runs, so the chunk that authenticated still arrives.
203-
if (settled) return
204+
// A refused connection stays counted: its peer may ignore the close and keep sending.
205+
if (authenticated) return
204206
preAuthBytes += chunk.length
205207
if (preAuthBytes <= MAX_PRE_AUTH_BYTES) return
206-
settle()
207-
log('warn', 'Rejected worker connection', { reason: 'too much data before authentication' })
208+
if (settle()) {
209+
log('warn', 'Rejected worker connection', { reason: 'too much data before authentication' })
210+
}
208211
socket.destroy()
209212
}
210213
const settle = () => {
211214
if (settled) return false
212215
settled = true
213216
pendingConnections--
214217
clearTimeout(authTimeout)
215-
socket.off('data', onPreAuthData)
216218
return true
217219
}
218220
const refuse = (reason) => {
@@ -237,6 +239,8 @@ function authenticate(ws, socket) {
237239
return
238240
}
239241
if (!settle()) return
242+
authenticated = true
243+
socket.off('data', onPreAuthData)
240244
if (draining) {
241245
ws.close(1012, 'draining')
242246
return

‎apps/sim/lib/execution/isolated-vm.ts‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import { createLogger } from '@sim/logger'
66
import { getErrorMessage, toError } from '@sim/utils/errors'
77
import { filterUndefined } from '@sim/utils/object'
88
import { randomFloat } from '@sim/utils/random'
9-
import { env } from '@/lib/core/config/env'
9+
import { env, envNumber } from '@/lib/core/config/env'
1010
import { getConfiguredCacheProvider } from '@/lib/core/config/env-capabilities.server'
1111
import { isFeatureEnabled } from '@/lib/core/config/feature-flags'
1212
import { getRedisClient } from '@/lib/core/config/redis'
@@ -154,7 +154,10 @@ const MAX_EXECUTIONS_PER_WORKER = Number.parseInt(env.IVM_MAX_EXECUTIONS_PER_WOR
154154
* delay and never tells the host, so this bound is what drains a host deploy or scale-in: infra
155155
* sets that delay to at least this age plus the longest execution a release keeps alive.
156156
*/
157-
const REMOTE_WORKER_MAX_AGE_MS = Number.parseInt(env.IVM_REMOTE_WORKER_MAX_AGE_MS) || 120_000
157+
const REMOTE_WORKER_MAX_AGE_MS = envNumber(env.IVM_REMOTE_WORKER_MAX_AGE_MS, 120_000, {
158+
min: 1000,
159+
integer: true,
160+
})
158161
/** Spreads retirements so a pool's connections don't all reopen at once; only ever shortens the age. */
159162
const REMOTE_WORKER_MAX_AGE_JITTER = 0.2
160163
/** Times one execution moves to another worker because its connection closed before it was sent. */
@@ -1380,15 +1383,17 @@ function dispatchToWorker(
13801383
workerInfo.retiring = true
13811384
if (workerInfo.activeExecutions === 0) cleanupWorker(workerInfo.id)
13821385
// The owner's state may have been released above; look it up again rather than queue on it.
1383-
queueMicrotask(() =>
1386+
queueMicrotask(() => {
1387+
// An abort in between already settled the caller and has no queue entry to remove.
1388+
if (state.cancelled) return
13841389
enqueueExecution(
13851390
getOrCreateOwnerState(ownerState.ownerKey, ownerState.weight),
13861391
req,
13871392
resolve,
13881393
state,
13891394
brokers
13901395
)
1391-
)
1396+
})
13921397
return
13931398
}
13941399
resolve({
@@ -1549,7 +1554,6 @@ export async function executeInIsolatedVM(
15491554
): Promise<IsolatedVMExecutionResult> {
15501555
const ownerKey = normalizeOwnerKey(req.ownerKey)
15511556
const ownerWeight = normalizeOwnerWeight(req.ownerWeight)
1552-
const ownerState = getOrCreateOwnerState(ownerKey, ownerWeight)
15531557
const brokers = options?.brokers
15541558
const signal = options?.signal
15551559

@@ -1642,6 +1646,9 @@ export async function executeInIsolatedVM(
16421646
// An undetermined lease cannot reject the execution: the per-process pool and
16431647
// the per-owner active/queued limits above still bound this work.
16441648

1649+
// Looked up only after the waits above: another execution finishing meanwhile can release
1650+
// this owner's state, and work queued on a released state is never dispatched.
1651+
const ownerState = getOrCreateOwnerState(ownerKey, ownerWeight)
16451652
const state: ExecutionState = {
16461653
cancelled: false,
16471654
runInOutboundScope: captureOutboundScope(),

‎apps/sim/scripts/test-isolated-vm-remote-e2e.ts‎

Lines changed: 78 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,60 @@ function upgradeStatus(host: Host, options: { path?: string } = {}): Promise<num
199199
})
200200
}
201201

202+
/** A client-to-server text frame header, masked with a zero key so the payload goes as is. */
203+
function maskedTextFrameHeader(payloadBytes: number): Buffer {
204+
if (payloadBytes < 126) return Buffer.from([0x81, 0x80 | payloadBytes, 0, 0, 0, 0])
205+
const header = Buffer.alloc(14)
206+
header[0] = 0x81
207+
header[1] = 0x80 | 127
208+
header.writeBigUInt64BE(BigInt(payloadBytes), 2)
209+
return header
210+
}
211+
212+
/**
213+
* Upgrades a raw socket, optionally sends `firstFrame`, then declares a 32 MiB frame (under the
214+
* host's message limit) and streams 1 MiB of it without ever authenticating. Resolves with
215+
* whether the host closed the connection within a second.
216+
*/
217+
async function streamUnauthenticated(host: Host, firstFrame?: string): Promise<'closed' | 'open'> {
218+
const socket = connect(host.port, '127.0.0.1')
219+
socket.on('error', () => {})
220+
const closed = new Promise<void>((resolve) => socket.once('close', () => resolve()))
221+
await new Promise<void>((resolve) => socket.once('connect', () => resolve()))
222+
const handshake = new Promise<string>((resolve) =>
223+
socket.once('data', (data: Buffer) => resolve(data.toString()))
224+
)
225+
socket.write(
226+
[
227+
'GET /v1/worker HTTP/1.1',
228+
`Host: 127.0.0.1:${host.port}`,
229+
'Upgrade: websocket',
230+
'Connection: Upgrade',
231+
'Sec-WebSocket-Version: 13',
232+
`Sec-WebSocket-Key: ${Buffer.from(generateRandomBytes(16)).toString('base64')}`,
233+
'',
234+
'',
235+
].join('\r\n')
236+
)
237+
const response = await handshake
238+
assert.match(response, /^HTTP\/1\.1 101/, `Expected the upgrade to be accepted: ${response}`)
239+
if (firstFrame) {
240+
const payload = Buffer.from(firstFrame)
241+
socket.write(Buffer.concat([maskedTextFrameHeader(payload.length), payload]))
242+
}
243+
socket.write(maskedTextFrameHeader(32 * 1024 * 1024))
244+
const chunk = Buffer.alloc(64 * 1024)
245+
for (let sent = 0; sent < 1024 * 1024 && !socket.destroyed; sent += chunk.length) {
246+
socket.write(chunk)
247+
}
248+
const outcome = await Promise.race([
249+
closed.then(() => 'closed' as const),
250+
sleep(1000).then(() => 'open' as const),
251+
])
252+
socket.destroy()
253+
return outcome
254+
}
255+
202256
/** Opens a connection and sends `firstFrame`; resolves with the socket once it opens. */
203257
function openSocket(host: Host, firstFrame: string): Promise<WebSocket> {
204258
return new Promise((resolve, reject) => {
@@ -312,9 +366,15 @@ async function run() {
312366
const { executeInIsolatedVM } = await import('@/lib/execution/isolated-vm')
313367

314368
// The isolate has no timers; a slow upstream holds an execution open the way real ones run long.
369+
// `/hold` answers only when the check releases it, so an inspection never races the execution.
370+
let onHeld: (release: () => void) => void = () => {}
315371
const slowServer = createHttpServer((req, res) => {
316-
const delay = Number(new URL(req.url ?? '/', 'http://slow').searchParams.get('ms'))
317-
setTimeout(() => res.end('done'), delay)
372+
const url = new URL(req.url ?? '/', 'http://slow')
373+
if (url.pathname === '/hold') {
374+
onHeld(() => res.end('done'))
375+
return
376+
}
377+
setTimeout(() => res.end('done'), Number(url.searchParams.get('ms')))
318378
})
319379
await new Promise<void>((resolve) => slowServer.listen(0, '127.0.0.1', resolve))
320380
const slowAddress = slowServer.address()
@@ -345,12 +405,19 @@ async function run() {
345405

346406
/** Inspects workers while an execution holds one open, since an idle worker can retire at any time. */
347407
const whileAWorkerRuns = async (inspect: () => void) => {
348-
const running = execute(slowFetch(500))
349-
await sleep(200)
408+
const held = new Promise<() => void>((resolve) => {
409+
onHeld = resolve
410+
})
411+
const running = execute(
412+
`const r = await fetch('http://127.0.0.1:${slowAddress.port}/hold'); return await r.text()`
413+
)
414+
const release = await held
350415
try {
351416
inspect()
352417
} finally {
353-
await running
418+
release()
419+
const result = await running
420+
assert.equal(result.result, 'done', result.error?.message)
354421
}
355422
}
356423

@@ -437,48 +504,16 @@ async function run() {
437504
})
438505

439506
await check('drops a connection that streams a large frame before authenticating', async () => {
440-
const socket = connect(probeHost.port, '127.0.0.1')
441-
socket.on('error', () => {})
442-
const closed = new Promise<void>((resolve) => socket.once('close', () => resolve()))
443-
await new Promise<void>((resolve) => socket.once('connect', () => resolve()))
444-
const handshake = new Promise<string>((resolve) =>
445-
socket.once('data', (data: Buffer) => resolve(data.toString()))
446-
)
447-
socket.write(
448-
[
449-
'GET /v1/worker HTTP/1.1',
450-
`Host: 127.0.0.1:${probeHost.port}`,
451-
'Upgrade: websocket',
452-
'Connection: Upgrade',
453-
'Sec-WebSocket-Version: 13',
454-
`Sec-WebSocket-Key: ${Buffer.from(generateRandomBytes(16)).toString('base64')}`,
455-
'',
456-
'',
457-
].join('\r\n')
458-
)
459-
const response = await handshake
460-
assert.match(response, /^HTTP\/1\.1 101/, `Expected the upgrade to be accepted: ${response}`)
461-
462-
// A masked text frame that declares 32 MiB, under the host's message limit, and never authenticates.
463-
const header = Buffer.alloc(14)
464-
header[0] = 0x81
465-
header[1] = 0x80 | 127
466-
header.writeBigUInt64BE(BigInt(32 * 1024 * 1024), 2)
467-
socket.write(header)
468-
const chunk = Buffer.alloc(64 * 1024)
469-
for (let sent = 0; sent < 1024 * 1024 && !socket.destroyed; sent += chunk.length) {
470-
socket.write(chunk)
471-
}
472-
const outcome = await Promise.race([
473-
closed.then(() => 'closed' as const),
474-
sleep(1000).then(() => 'open' as const),
475-
])
476-
socket.destroy()
477-
assert.equal(outcome, 'closed', 'The host kept buffering an unauthenticated frame')
507+
assert.equal(await streamUnauthenticated(probeHost), 'closed', 'The host kept buffering')
478508
const health = await fetch(`http://127.0.0.1:${probeHost.port}/health`)
479509
assert.equal(health.status, 200)
480510
})
481511

512+
await check('drops a refused connection that ignores the close and keeps sending', async () => {
513+
const refusedAuth = authFrame(generateRandomHex(32))
514+
assert.equal(await streamUnauthenticated(probeHost, refusedAuth), 'closed')
515+
})
516+
482517
await check('rejects an unparseable request target and keeps serving', async () => {
483518
assert.equal(await upgradeStatus(probeHost, { path: '//[' }), 400)
484519
const health = await fetch(`http://127.0.0.1:${probeHost.port}/health`)

0 commit comments

Comments
 (0)