starred/orca
1
0
Fork 0
mirror of https://github.com/stablyai/orca.git synced 2026-09-21 16:02:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(relay): batch control lease renewals per cell instead of one write transaction per host (#21303)

Browse Source 
* perf(relay): batch control lease renewals per cell instead of one write transaction per host

Every connected desktop renewed its own control lease with its own single-row
write transaction every 30s. At ~14,000 hosts that is ~470 write transactions
per second fleet-wide, each with its own transaction id, all updating the same
few heap pages of relay_assignments and relay_assignment_activity_leases.
Sampling three onsets at 250ms showed no lock queue and no slow statement:
60-144 backends piled into LWLock:BufferContent and Timeout/SpinDelay inside
that one statement, and Query Insights attributed 152 of 157 seconds of
lightweight-lock wait in the onset minute to it.

The heartbeat now enqueues a due renewal instead of issuing it. A cell flushes
its queue once per second, or as soon as 500 rows are waiting, through one
statement that unnests the parameter arrays and applies the same CTE row-wise.
Concurrent writers drop from the host count to the cell count, and transaction
ids with them. Measured against a 20,000-row table: 1 row 4.1ms, 12 rows
3.3ms, 100 rows 6.6ms, 500 rows 22.7ms.

Per-session semantics are unchanged. Each enqueue still resolves on a renewal
and rejects with the outcome as its message, so the completed-attempt counter,
the staleness guard, and every close path route exactly as before, and one
outcome per row is recorded against the flush latency.

Lock order is (user_id, relay_host_id), the primary key of relay_assignments,
applied in JavaScript and repeated as the statement's ORDER BY. EXPLAIN
confirms LockRows sits above that Sort, so a batch acquires its assignment rows
in one global order. Every writer in the store locks a host's assignment row
before its migration or lease rows and only ever touches one host, so a batch
can only wait on a row a single-host writer holds, never the reverse.

One statement also means one contended row could fail the whole batch, so a
failed batch degrades to the per-host statements it replaced rather than
costing every other host on the cell its renewal.

* perf(relay): batch control lease renewals per cell instead of one write transaction per host

Every connected desktop renewed its own control lease with its own single-row
write transaction every 30s. At ~14,000 hosts that is ~470 write transactions
per second fleet-wide, each with its own transaction id, all updating the same
few heap pages of relay_assignments and relay_assignment_activity_leases.
Sampling three onsets at 250ms showed no lock queue and no slow statement:
60-144 backends piled into LWLock:BufferContent and Timeout/SpinDelay inside
that one statement, and Query Insights attributed 152 of 157 seconds of
lightweight-lock wait in the onset minute to it.

The heartbeat now enqueues a due renewal instead of issuing it. A cell flushes
its queue every second, or as soon as 200 rows are waiting, through one
statement that unnests the parameter arrays and applies the same CTE row-wise.
Concurrent writers drop from the host count to the cell count, and transaction
ids with them. Per-host buffer traffic is unchanged: 30 hits for one row, 25.5
per host at 12 rows, 30.1 per host at 200, against the 28.7 the single-row
statement reports in production.

Per-session semantics are unchanged. Each enqueue still resolves on a renewal
and rejects with the outcome as its message, so the completed-attempt counter,
the staleness guard, and every close path route as before, and one outcome per
row is recorded against the flush latency.

Row locks live until the statement commits, so a batch that waited on a
contended row would hold every other row's lock for that whole wait. The
assignment pass therefore takes its locks with SKIP LOCKED and reports a
contended host as assignment_lock_unavailable, which the registry retries on the
next tick instead of closing the control. That keeps the hold to the statement's
own execution: 11.5ms for 200 rows against a 20,000-row table, and 9.4ms with a
host wedged in a per-host transaction, where a blocking FOR UPDATE spends the
pool's whole 1s lock_timeout and then fails every row in the flush.

An unlocked present_assignment probe separates a host with no assignment row
from one the skip passed over, so a skipped row can never be mistaken for a
missing assignment and close a live desktop.

With no wait on the assignment pass the lock order is only needed for the two
later passes, and it holds: every writer takes a host's assignment row before
that host's lease rows, and a host whose assignment row is held was skipped, so
the batch never reaches its lease. markMigrationTargetRegistered is the one
writer that locks a migration row first, and it takes no further locks.

* fix(relay): answer every row of a control-renewal batch from its own lease update

Review findings on the batched renewal.

Two control leases on one host in one batch made the second report
control_activity_not_found although both were renewed: the assignment UPDATE is
offered the same target row twice, applies one source row and returns one, so
the other row_index never came back. The verdict now reads renewed_lease, which
has a row per input row, and the assignment UPDATE groups per host so it also
stops taking an arbitrary one of the two expiries instead of the later one.

The queue now partitions per (userId, relayHostId) rather than per activity, so
a second control activity for one host opens the next flush instead of sharing
this one. Belt to the statement fix, not a substitute: the store API has to be
right for the rows it accepts.

renewControlActivities recorded no outcome for a one-row flush that threw, and
none at all when every row failed validation, where a mixed batch recorded its
invalid_* rows. Both now record in a finally, the way the single-row path's
finally always did, and the error-to-outcome mapping both paths share is one
function.

The four flush fields the runtime metrics event emits had no log-based metric,
so add them next to the existing controlRenewalLatencyMs* entries. Applying the
Terraform is a separate manual step.

controlRenewalLatencyMsP50/P95/Max now measure a batched row's flush duration
rather than its own statement latency. Left named as they are for history, with
a line at the emit site recording that the meaning changed here.
This commit is contained in:
Jinwoo Hong
2026-09-17 19:25:54 -04:00
committed by GitHub
parent ddbad2218b
commit 91ade4b82a
11 changed files with 1392 additions and 134 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cloud/apps/relay/src/assignment-store.ts
+159 -119
View File
@@ -45,6 +45,15 @@ import {
ASSIGNMENT_CONNECTION_HEADROOM_QUERY
} from './assignment-connection-headroom-query.js'
import { AssignmentIdentityQueue } from './assignment-identity-queue.js'
import {
CONTROL_RENEWAL_BATCH_SQL,
CONTROL_RENEWAL_STATEMENT_OUTCOMES,
controlRenewalBatchParams,
orderedControlRenewalRows,
readControlRenewalOutcomes,
type ControlRenewalOutcome,
type ControlRenewalRequest
} from './control-renewal-statement.js'
import {
REGIONAL_REHOME_DEFAULT_HOST_COOLDOWN_MS
} from './database.js'
@@ -101,21 +110,8 @@ type RelayAssignmentStoreOptions = {
recordControlRenewal?: (durationMs: number, outcome: ControlRenewalOutcome) => void
}
export type ControlRenewalOutcome =
| 'renewed'
| 'assignment_not_found'
| 'activity_cell_not_authoritative'
| 'control_activity_not_found'
| 'control_activity_moved'
| 'database_error'
export type { ControlRenewalOutcome, ControlRenewalRequest }
const CONTROL_RENEWAL_OUTCOMES = new Set<ControlRenewalOutcome>([
'renewed',
'assignment_not_found',
'activity_cell_not_authoritative',
'control_activity_not_found',
'control_activity_moved'
])
export type RelayAssignment = AssignmentIdentity & {
cellId: string
cellUrl: string
@@ -3478,132 +3474,149 @@ export class RelayAssignmentStore {
})
}
// Kept as the single-row contract for callers and tests: resolves on a
// renewal and throws the outcome (or the driver's own error) otherwise.
async renewControlActivity(
identity: AssignmentIdentity,
input: { activityId: string; cellId: string; expiresAt: number }
): Promise<void> {
validateActivityId(input.activityId)
const now = this.now()
const maximumExpiresAt =
now +
ASSIGNMENT_LIMITS.activityLeaseMs +
RELAY_PROTOCOL_LIMITS.controlPingIntervalMs * 2
if (
!Number.isSafeInteger(input.expiresAt) ||
input.expiresAt <= now ||
input.expiresAt > maximumExpiresAt
) {
if (!controlRenewalExpiryIsValid(input.expiresAt, now)) {
throw new Error('invalid_activity_expiry')
}
const startedAt = performance.now()
let outcome: ControlRenewalOutcome = 'database_error'
try {
outcome =
this.database.dialect === 'postgres'
? await this.renewPostgresControlActivity(identity, input, now)
: await this.renewTransactionalControlActivity(identity, input, now)
outcome = await this.renewOneControlActivity({ identity, ...input }, now)
if (outcome !== 'renewed') throw new Error(outcome)
} catch (error) {
const message = String((error as { message?: unknown }).message)
if (CONTROL_RENEWAL_OUTCOMES.has(message as ControlRenewalOutcome)) {
outcome = message as ControlRenewalOutcome
}
outcome = controlRenewalOutcomeOfError(error)
throw error
} finally {
this.recordControlRenewal?.(performance.now() - startedAt, outcome)
}
}
private async renewPostgresControlActivity(
identity: AssignmentIdentity,
input: { activityId: string; cellId: string; expiresAt: number },
// Renews every due control lease on a cell in one write transaction, returning
// one outcome per input row in input order. Never throws for a multi-row batch:
// a caller routes its own session on its own outcome.
async renewControlActivities(
rows: readonly ControlRenewalRequest[]
): Promise<ControlRenewalOutcome[]> {
const now = this.now()
const outcomes = new Array<ControlRenewalOutcome>(rows.length)
const accepted: Array<ControlRenewalRequest & { index: number }> = []
for (const [index, row] of rows.entries()) {
const rejection = controlRenewalRejection(row, now)
if (rejection) outcomes[index] = rejection
else accepted.push({ ...row, index })
}
const startedAt = performance.now()
try {
if (accepted.length === 0) return outcomes
let results: ControlRenewalOutcome[]
try {
results = await this.executeControlRenewals(accepted, now)
} catch (error) {
if (rows.length === 1) {
outcomes[accepted[0]!.index] = controlRenewalOutcomeOfError(error)
throw error
}
results = accepted.map(() => 'database_error')
}
for (const [position, row] of accepted.entries()) outcomes[row.index] = results[position]!
return outcomes
} finally {
// Every path, so a rethrown lone renewal and an all-invalid batch are
// counted the same as a batch that reached PostgreSQL.
const durationMs = performance.now() - startedAt
for (const outcome of outcomes) this.recordControlRenewal?.(durationMs, outcome)
}
}
private async executeControlRenewals(
rows: readonly ControlRenewalRequest[],
now: number
): Promise<ControlRenewalOutcome[]> {
if (rows.length === 1) return [await this.renewOneControlActivity(rows[0]!, now)]
if (this.database.dialect !== 'postgres') {
// Correctness over throughput: the SQLite writer is serialized anyway, and
// this is the dialect the unit suites run on.
return await this.renewControlActivitiesInSeries(rows, now)
}
const ordered = orderedControlRenewalRows(rows.map((row, index) => ({ ...row, index })))
const outcomes = new Array<ControlRenewalOutcome>(rows.length)
try {
const parsed = readControlRenewalOutcomes(
await this.database.query(
CONTROL_RENEWAL_BATCH_SQL,
controlRenewalBatchParams(ordered, now)
),
ordered.length
)
for (const [position, row] of ordered.entries()) outcomes[row.index] = parsed[position]!
return outcomes
} catch (error) {
// One statement means one contended assignment row can fail the whole
// batch, so a failure degrades to the per-host statements this replaced
// rather than costing every other host on the cell its renewal.
console.warn(
JSON.stringify({
event: 'orca_relay_control_renewal_batch_failed',
rows: ordered.length,
message: String((error as { message?: unknown }).message)
})
)
await Promise.all(
ordered.map(async (row) => {
try {
outcomes[row.index] = await this.renewOneControlActivity(row, now)
} catch {
outcomes[row.index] = 'database_error'
}
})
)
return outcomes
}
}
private async renewControlActivitiesInSeries(
rows: readonly ControlRenewalRequest[],
now: number
): Promise<ControlRenewalOutcome[]> {
const outcomes: ControlRenewalOutcome[] = []
for (const row of rows) {
try {
outcomes.push(await this.renewOneControlActivity(row, now))
} catch {
outcomes.push('database_error')
}
}
return outcomes
}
// Returns the outcome; a driver or pool failure reaches the caller unchanged.
private async renewOneControlActivity(
row: ControlRenewalRequest,
now: number
): Promise<ControlRenewalOutcome> {
const row = (
await this.database.query(
`WITH assignment_state AS MATERIALIZED (
SELECT cell_id, assignment_epoch
FROM relay_assignments
WHERE user_id = ? AND relay_host_id = ?
FOR UPDATE
), migration_state AS MATERIALIZED (
SELECT migration.assignment_epoch
FROM relay_assignment_migrations migration
JOIN assignment_state assignment
ON migration.target_cell_id = assignment.cell_id
AND migration.assignment_epoch = assignment.assignment_epoch
WHERE migration.user_id = ? AND migration.relay_host_id = ?
AND migration.source_cell_id = ?
AND migration.completed_at IS NULL AND migration.aborted_at IS NULL
FOR UPDATE OF migration
), authorization_state AS MATERIALIZED (
SELECT 1 AS authorized
FROM assignment_state assignment
WHERE assignment.cell_id = ? OR EXISTS (SELECT 1 FROM migration_state)
), lease_state AS MATERIALIZED (
SELECT lease.activity_kind, lease.cell_id
FROM relay_assignment_activity_leases lease
CROSS JOIN authorization_state
WHERE lease.user_id = ? AND lease.relay_host_id = ? AND lease.activity_id = ?
FOR UPDATE OF lease
), renewed_lease AS (
UPDATE relay_assignment_activity_leases lease
SET expires_at = GREATEST(lease.expires_at, ?),
updated_at = GREATEST(lease.updated_at, ?)
FROM lease_state state
WHERE lease.user_id = ? AND lease.relay_host_id = ? AND lease.activity_id = ?
AND state.activity_kind = 'control' AND state.cell_id = ?
RETURNING 1
), renewed_assignment AS (
UPDATE relay_assignments assignment
SET lease_expires_at = GREATEST(assignment.lease_expires_at, ?),
last_activity_at = GREATEST(assignment.last_activity_at, ?)
WHERE assignment.user_id = ? AND assignment.relay_host_id = ?
AND EXISTS (SELECT 1 FROM renewed_lease)
RETURNING 1
)
SELECT CASE
WHEN NOT EXISTS (SELECT 1 FROM assignment_state)
THEN 'assignment_not_found'
WHEN NOT EXISTS (SELECT 1 FROM authorization_state)
THEN 'activity_cell_not_authoritative'
WHEN NOT EXISTS (SELECT 1 FROM lease_state)
THEN 'control_activity_not_found'
WHEN EXISTS (
SELECT 1 FROM lease_state
WHERE activity_kind <> 'control' OR cell_id <> ?
) THEN 'control_activity_moved'
WHEN EXISTS (SELECT 1 FROM renewed_assignment) THEN 'renewed'
ELSE 'control_activity_not_found'
END AS outcome`,
[
identity.userId,
identity.relayHostId,
identity.userId,
identity.relayHostId,
input.cellId,
input.cellId,
identity.userId,
identity.relayHostId,
input.activityId,
input.expiresAt,
now,
identity.userId,
identity.relayHostId,
input.activityId,
input.cellId,
input.expiresAt,
now,
identity.userId,
identity.relayHostId,
input.cellId
]
)
)[0]
if (!row) throw new Error('missing_control_renewal_outcome')
const outcome = text(row, 'outcome') as ControlRenewalOutcome
if (!CONTROL_RENEWAL_OUTCOMES.has(outcome)) throw new Error('invalid_control_renewal_outcome')
return outcome
if (this.database.dialect === 'postgres') {
return readControlRenewalOutcomes(
await this.database.query(
CONTROL_RENEWAL_BATCH_SQL,
controlRenewalBatchParams([row], now)
),
1
)[0]!
}
try {
return await this.renewTransactionalControlActivity(row.identity, row, now)
} catch (error) {
const message = String((error as { message?: unknown }).message)
if (!CONTROL_RENEWAL_STATEMENT_OUTCOMES.has(message as ControlRenewalOutcome)) throw error
return message as ControlRenewalOutcome
}
}
private async renewTransactionalControlActivity(
@@ -7900,6 +7913,33 @@ function validateActivityId(activityId: string): void {
if (!activityId || activityId.length > 256) throw new Error('invalid_activity_id')
}
function controlRenewalExpiryIsValid(expiresAt: number, now: number): boolean {
const maximumExpiresAt =
now + ASSIGNMENT_LIMITS.activityLeaseMs + RELAY_PROTOCOL_LIMITS.controlPingIntervalMs * 2
return Number.isSafeInteger(expiresAt) && expiresAt > now && expiresAt <= maximumExpiresAt
}
// A renewal that threw still owes the metric an outcome: the message carries one
// when the statement decided it, and anything else is the driver failing.
function controlRenewalOutcomeOfError(error: unknown): ControlRenewalOutcome {
const message = String((error as { message?: unknown }).message)
return CONTROL_RENEWAL_STATEMENT_OUTCOMES.has(message as ControlRenewalOutcome)
? (message as ControlRenewalOutcome)
: 'database_error'
}
function controlRenewalRejection(
row: ControlRenewalRequest,
now: number
): ControlRenewalOutcome | null {
try {
validateActivityId(row.activityId)
} catch {
return 'invalid_activity_id'
}
return controlRenewalExpiryIsValid(row.expiresAt, now) ? null : 'invalid_activity_expiry'
}
function activityKind(row: SqlRow): AssignmentActivityKind {
const value = text(row, 'activity_kind')
if (!(value in ACTIVITY_REQUEST_UNITS)) throw new Error('invalid_activity_kind')
cloud/apps/relay/src/control-lease-recovery-postgres.test.ts
+8 -2
View File
@@ -3,6 +3,7 @@ import { ASSIGNMENT_LIMITS, RELAY_CLOSE_CODE } from '@orca-cloud/relay-contract'
import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
import type WebSocket from 'ws'
import { RelayAssignmentStore } from './assignment-store.js'
import { CONTROL_RENEWAL_BATCH_INTERVAL_MS } from './control-renewal-batch.js'
import type { RelayConfig } from './config.js'
import type { RelayCredentialStore } from './credential-store.js'
import { openRelayDatabase, type RelayDatabase } from './database.js'
@@ -133,6 +134,10 @@ describePostgres('expired control lease after a database outage', () => {
internals.heartbeat(session)
}
// A due renewal leaves the heartbeat as a batch enqueue, so a poll has to
// outlast the batch window before it can call the renewal missing.
const renewalPoll = { timeout: CONTROL_RENEWAL_BATCH_INTERVAL_MS + 4_000 }
const leaseRows = async (relayHostId: string) =>
await database.query(
`SELECT activity_id, cell_id FROM relay_assignment_activity_leases
@@ -149,7 +154,8 @@ describePostgres('expired control lease after a database outage', () => {
.poll(
async () =>
socket.close.mock.calls.length > 0 ||
(await leaseRows(relayHostId)).length === expectedRows
(await leaseRows(relayHostId)).length === expectedRows,
renewalPoll
)
.toBe(true)
}
@@ -212,7 +218,7 @@ describePostgres('expired control lease after a database outage', () => {
).rejects.toThrow('control_activity_moved')
heartbeat(registry, session)
await expect.poll(() => socket.close.mock.calls.length).toBe(1)
await expect.poll(() => socket.close.mock.calls.length, renewalPoll).toBe(1)
expect(socket.close).toHaveBeenCalledWith(RELAY_CLOSE_CODE.DRAINING, 'control activity moved')
expect(await leaseRows(identity.relayHostId)).toEqual([
cloud/apps/relay/src/control-renewal-batch-store.test.ts
+291
View File
@@ -0,0 +1,291 @@
import { ASSIGNMENT_LIMITS } from '@orca-cloud/relay-contract'
import { describe, expect, it, vi } from 'vitest'
import { RelayAssignmentStore } from './assignment-store.js'
import { CONTROL_RENEWAL_BATCH_SQL } from './control-renewal-statement.js'
import type { ControlRenewalOutcome } from './control-renewal-statement.js'
import {
openInMemoryRelayDatabase,
type RelayDatabase,
type SqlRow
} from './database.js'
const now = 1_900_000_000_000
const expiresAt = now + 105_000
function renewal(userId: string, relayHostId: string, expiry = expiresAt) {
return {
identity: { userId, relayHostId },
activityId: 'control:cell-a:1',
cellId: 'cell-a',
expiresAt: expiry
}
}
// A PostgreSQL-dialect database that answers the renewal statement without a
// server, so the statement count and its parameter arrays are observable.
class RenewalStatementProbe implements RelayDatabase {
readonly dialect = 'postgres' as const
readonly statements: Array<{ sql: string; params: unknown[] }> = []
failuresRemaining = 0
constructor(private readonly outcomeFor: (userId: string) => ControlRenewalOutcome) {}
async query(sql: string, params: unknown[] = []): Promise<SqlRow[]> {
this.statements.push({ sql, params })
if (this.failuresRemaining > 0) {
this.failuresRemaining -= 1
throw new Error('canceling statement due to statement timeout')
}
const userIds = params[0] as string[]
return userIds.map((userId, index) => ({
row_index: String(index + 1),
outcome: this.outcomeFor(userId)
}))
}
async queryLocked(): Promise<SqlRow[]> {
throw new Error('unexpected_locked_query')
}
async transaction<T>(): Promise<T> {
// Renewals must never open one: that is the write transaction per host this
// batch exists to remove.
throw new Error('unexpected_transaction')
}
async close(): Promise<void> {}
}
describe('batched control renewals on PostgreSQL', () => {
it('spends one statement on every host that came due', async () => {
const probe = new RenewalStatementProbe(() => 'renewed')
const store = new RelayAssignmentStore(probe, () => now)
const outcomes = await store.renewControlActivities([
renewal('user-a', 'host000000000001'),
renewal('user-a', 'host000000000002'),
renewal('user-b', 'host000000000003')
])
expect(outcomes).toEqual(['renewed', 'renewed', 'renewed'])
expect(probe.statements).toHaveLength(1)
expect(probe.statements[0]!.sql).toBe(CONTROL_RENEWAL_BATCH_SQL)
expect(probe.statements[0]!.params[0]).toEqual(['user-a', 'user-a', 'user-b'])
expect(probe.statements[0]!.params[4]).toEqual([expiresAt, expiresAt, expiresAt])
})
it('locks assignment rows in primary-key order and still answers in input order', async () => {
const probe = new RenewalStatementProbe((userId) =>
userId === 'user-b' ? 'control_activity_moved' : 'renewed'
)
const store = new RelayAssignmentStore(probe, () => now)
const outcomes = await store.renewControlActivities([
renewal('user-c', 'host000000000003'),
renewal('user-a', 'host000000000002'),
renewal('user-b', 'host000000000001'),
renewal('user-a', 'host000000000001')
])
// (user_id, relay_host_id) is the primary key of relay_assignments, and the
// statement's ORDER BY repeats it: no batch can queue against another in a
// different sequence.
expect(probe.statements[0]!.params[0]).toEqual(['user-a', 'user-a', 'user-b', 'user-c'])
expect(probe.statements[0]!.params[1]).toEqual([
'host000000000001',
'host000000000002',
'host000000000001',
'host000000000003'
])
expect(outcomes).toEqual([
'renewed',
'renewed',
'control_activity_moved',
'renewed'
])
})
it('keeps a malformed request out of the statement and fails only that row', async () => {
const probe = new RenewalStatementProbe(() => 'renewed')
const store = new RelayAssignmentStore(probe, () => now)
const outcomes = await store.renewControlActivities([
renewal('user-a', 'host000000000001'),
renewal('user-a', 'host000000000002', now + ASSIGNMENT_LIMITS.activityLeaseMs * 10),
{ ...renewal('user-a', 'host000000000003'), activityId: '' },
renewal('user-a', 'host000000000004')
])
expect(outcomes).toEqual([
'renewed',
'invalid_activity_expiry',
'invalid_activity_id',
'renewed'
])
expect(probe.statements[0]!.params[1]).toEqual(['host000000000001', 'host000000000004'])
})
it('degrades to one statement per host when the batch statement fails', async () => {
const probe = new RenewalStatementProbe(() => 'renewed')
probe.failuresRemaining = 1
const store = new RelayAssignmentStore(probe, () => now)
const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined)
try {
const outcomes = await store.renewControlActivities([
renewal('user-a', 'host000000000001'),
renewal('user-a', 'host000000000002')
])
expect(outcomes).toEqual(['renewed', 'renewed'])
expect(probe.statements).toHaveLength(3)
expect(probe.statements[1]!.params[1]).toEqual(['host000000000001'])
expect(probe.statements[2]!.params[1]).toEqual(['host000000000002'])
expect(JSON.parse(String(warn.mock.calls[0]![0]))).toMatchObject({
event: 'orca_relay_control_renewal_batch_failed',
rows: 2
})
} finally {
warn.mockRestore()
}
})
it('reports a host that fails its own fallback statement without touching the rest', async () => {
const probe = new RenewalStatementProbe(() => 'renewed')
probe.failuresRemaining = 2
const store = new RelayAssignmentStore(probe, () => now)
const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined)
try {
const outcomes = await store.renewControlActivities([
renewal('user-a', 'host000000000001'),
renewal('user-a', 'host000000000002')
])
expect(outcomes.filter((outcome) => outcome === 'renewed')).toHaveLength(1)
expect(outcomes.filter((outcome) => outcome === 'database_error')).toHaveLength(1)
} finally {
warn.mockRestore()
}
})
it('reports a contended assignment row apart from a missing one', async () => {
const probe = new RenewalStatementProbe((userId) =>
userId === 'user-b' ? 'assignment_lock_unavailable' : 'renewed'
)
const store = new RelayAssignmentStore(probe, () => now)
const outcomes = await store.renewControlActivities([
renewal('user-a', 'host000000000001'),
renewal('user-b', 'host000000000002')
])
// Retryable: SKIP LOCKED passed over the row rather than queueing the whole
// flush behind whoever held it.
expect(outcomes).toEqual(['renewed', 'assignment_lock_unavailable'])
})
it('counts a lone renewal that threw before rethrowing it', async () => {
const probe = new RenewalStatementProbe(() => 'renewed')
probe.failuresRemaining = 1
const recordControlRenewal = vi.fn()
const store = new RelayAssignmentStore(probe, () => now, { recordControlRenewal })
// A one-row flush keeps the pre-batch contract and rethrows, but the metric
// still owes an outcome for the attempt.
await expect(
store.renewControlActivities([renewal('user-a', 'host000000000001')])
).rejects.toThrow('statement timeout')
expect(recordControlRenewal).toHaveBeenCalledTimes(1)
expect(recordControlRenewal.mock.calls[0]![1]).toBe('database_error')
})
it('counts a batch in which every row was rejected before the statement', async () => {
const probe = new RenewalStatementProbe(() => 'renewed')
const recordControlRenewal = vi.fn()
const store = new RelayAssignmentStore(probe, () => now, { recordControlRenewal })
const outcomes = await store.renewControlActivities([
{ ...renewal('user-a', 'host000000000001'), activityId: '' },
renewal('user-a', 'host000000000002', now - 1)
])
expect(outcomes).toEqual(['invalid_activity_id', 'invalid_activity_expiry'])
expect(probe.statements).toHaveLength(0)
expect(recordControlRenewal.mock.calls.map((call) => call[1])).toEqual([
'invalid_activity_id',
'invalid_activity_expiry'
])
})
it('counts one renewal metric per row against the flush latency', async () => {
const probe = new RenewalStatementProbe((userId) =>
userId === 'user-b' ? 'assignment_not_found' : 'renewed'
)
const recordControlRenewal = vi.fn()
const store = new RelayAssignmentStore(probe, () => now, { recordControlRenewal })
await store.renewControlActivities([
renewal('user-a', 'host000000000001'),
renewal('user-b', 'host000000000002')
])
expect(recordControlRenewal).toHaveBeenCalledTimes(2)
expect(recordControlRenewal.mock.calls.map((call) => call[1])).toEqual([
'renewed',
'assignment_not_found'
])
})
})
describe('batched control renewals on SQLite', () => {
it('renews every host through the transactional path', async () => {
let clock = now
const database = await openInMemoryRelayDatabase()
try {
const store = new RelayAssignmentStore(database, () => clock)
await store.reconcileCells([
{ id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }
])
const hosts = ['host000000000001', 'host000000000002']
const requests = []
for (const relayHostId of hosts) {
const identity = { userId: 'user-a', relayHostId }
const assignment = await store.assign(identity)
await store.activateControl(identity, {
cellId: assignment.cellId,
assignmentEpoch: assignment.assignmentEpoch,
generation: 1
})
requests.push({
identity,
activityId: `control:${assignment.cellId}:1`,
cellId: assignment.cellId,
expiresAt: clock + 105_000
})
}
// A host with no assignment at all must not cost the others their renewal.
requests.push({
identity: { userId: 'user-a', relayHostId: 'host000000000009' },
activityId: 'control:cell-a:1',
cellId: 'cell-a',
expiresAt: clock + 105_000
})
clock += 1_000
const outcomes = await store.renewControlActivities(requests)
expect(outcomes).toEqual(['renewed', 'renewed', 'assignment_not_found'])
const leases = await database.query(
`SELECT relay_host_id, expires_at FROM relay_assignment_activity_leases
WHERE user_id = ? ORDER BY relay_host_id ASC`,
['user-a']
)
expect(leases.map((lease) => Number(lease.expires_at))).toEqual([
now + 105_000,
now + 105_000
])
} finally {
await database.close()
}
})
})
cloud/apps/relay/src/control-renewal-batch.test.ts
+231
View File
@@ -0,0 +1,231 @@
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import {
CONTROL_RENEWAL_BATCH_INTERVAL_MS,
CONTROL_RENEWAL_BATCH_MAX_ROWS,
ControlRenewalBatch,
type ControlRenewalFlush
} from './control-renewal-batch.js'
import type {
ControlRenewalOutcome,
ControlRenewalRequest
} from './control-renewal-statement.js'
// Settles into the outcome the caller saw, attached at enqueue so a rejection is
// never momentarily unhandled.
function outcomeOf(renewal: Promise<void>): Promise<string> {
return renewal.then(
() => 'renewed',
(error: unknown) => String((error as { message?: unknown }).message)
)
}
function request(
host: string,
expiresAt = 1_000,
activityId = 'control:cell-a:1'
): ControlRenewalRequest {
return {
identity: { userId: 'user-a', relayHostId: host },
activityId,
cellId: 'cell-a',
expiresAt
}
}
describe('control renewal batch', () => {
beforeEach(() => vi.useFakeTimers())
afterEach(() => {
vi.clearAllTimers()
vi.useRealTimers()
})
it('spends one call on every renewal that came due in the window', async () => {
const renew = vi.fn(async (rows: readonly ControlRenewalRequest[]) =>
rows.map((): ControlRenewalOutcome => 'renewed')
)
const batch = new ControlRenewalBatch(renew)
const settled = [
batch.enqueue(request('host0000000000a1')),
batch.enqueue(request('host0000000000a2')),
batch.enqueue(request('host0000000000a3'))
]
expect(renew).not.toHaveBeenCalled()
await vi.advanceTimersByTimeAsync(CONTROL_RENEWAL_BATCH_INTERVAL_MS)
await expect(Promise.all(settled)).resolves.toEqual([undefined, undefined, undefined])
expect(renew).toHaveBeenCalledOnce()
expect(renew.mock.calls[0]![0].map((row) => row.identity.relayHostId)).toEqual([
'host0000000000a1',
'host0000000000a2',
'host0000000000a3'
])
})
it('routes each outcome back to the caller that asked for it', async () => {
const outcomes: ControlRenewalOutcome[] = [
'renewed',
'assignment_not_found',
'control_activity_moved'
]
const batch = new ControlRenewalBatch(async () => outcomes)
const first = outcomeOf(batch.enqueue(request('host0000000000b1')))
const second = outcomeOf(batch.enqueue(request('host0000000000b2')))
const third = outcomeOf(batch.enqueue(request('host0000000000b3')))
await vi.advanceTimersByTimeAsync(CONTROL_RENEWAL_BATCH_INTERVAL_MS)
await expect(Promise.all([first, second, third])).resolves.toEqual([
'renewed',
'assignment_not_found',
'control_activity_moved'
])
})
it('flushes on reaching the row ceiling instead of waiting out the window', async () => {
const renew = vi.fn(async (rows: readonly ControlRenewalRequest[]) =>
rows.map((): ControlRenewalOutcome => 'renewed')
)
const batch = new ControlRenewalBatch(renew)
for (let row = 0; row < CONTROL_RENEWAL_BATCH_MAX_ROWS - 1; row++) {
void batch.enqueue(request(`host${String(row).padStart(12, '0')}`))
}
expect(renew).not.toHaveBeenCalled()
void batch.enqueue(request('host0000000000zz'))
await vi.advanceTimersByTimeAsync(0)
expect(renew).toHaveBeenCalledOnce()
expect(renew.mock.calls[0]![0]).toHaveLength(CONTROL_RENEWAL_BATCH_MAX_ROWS)
// The window timer must not fire a second, empty statement.
await vi.advanceTimersByTimeAsync(CONTROL_RENEWAL_BATCH_INTERVAL_MS)
expect(renew).toHaveBeenCalledOnce()
})
it('does not hold a new window behind a statement still in PostgreSQL', async () => {
let release!: (outcomes: ControlRenewalOutcome[]) => void
const renew = vi
.fn<(rows: readonly ControlRenewalRequest[]) => Promise<ControlRenewalOutcome[]>>()
.mockImplementationOnce(
async () => await new Promise<ControlRenewalOutcome[]>((resolve) => (release = resolve))
)
.mockResolvedValue(['renewed'])
const batch = new ControlRenewalBatch(renew)
const stalled = batch.enqueue(request('host0000000000c1'))
await vi.advanceTimersByTimeAsync(CONTROL_RENEWAL_BATCH_INTERVAL_MS)
const next = batch.enqueue(request('host0000000000c2'))
await vi.advanceTimersByTimeAsync(CONTROL_RENEWAL_BATCH_INTERVAL_MS)
expect(renew).toHaveBeenCalledTimes(2)
await expect(next).resolves.toBeUndefined()
release(['renewed'])
await expect(stalled).resolves.toBeUndefined()
})
it('reports the driver failure to every caller in the flush', async () => {
const batch = new ControlRenewalBatch(async () => {
throw new Error('pool timeout')
})
const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined)
try {
const first = outcomeOf(batch.enqueue(request('host0000000000d1')))
const second = outcomeOf(batch.enqueue(request('host0000000000d2')))
await vi.advanceTimersByTimeAsync(CONTROL_RENEWAL_BATCH_INTERVAL_MS)
await expect(Promise.all([first, second])).resolves.toEqual([
'pool timeout',
'pool timeout'
])
} finally {
warn.mockRestore()
}
})
it('supersedes a second attempt for one lease and answers both callers', async () => {
const renew = vi.fn(async (rows: readonly ControlRenewalRequest[]) =>
rows.map((): ControlRenewalOutcome => 'renewed')
)
const batch = new ControlRenewalBatch(renew)
const earlier = batch.enqueue(request('host0000000000e1', 1_000))
const later = batch.enqueue(request('host0000000000e1', 2_000))
await vi.advanceTimersByTimeAsync(CONTROL_RENEWAL_BATCH_INTERVAL_MS)
expect(renew.mock.calls[0]![0]).toEqual([
expect.objectContaining({ expiresAt: 2_000 })
])
await expect(earlier).resolves.toBeUndefined()
await expect(later).resolves.toBeUndefined()
})
it('holds a second activity for one host back to the next flush', async () => {
const renew = vi.fn(async (rows: readonly ControlRenewalRequest[]) =>
rows.map((): ControlRenewalOutcome => 'renewed')
)
const batch = new ControlRenewalBatch(renew)
const first = batch.enqueue(request('host0000000000h1', 1_000, 'control:cell-a:1'))
const second = batch.enqueue(request('host0000000000h1', 1_000, 'control:cell-a:2'))
const other = batch.enqueue(request('host0000000000h2'))
await vi.advanceTimersByTimeAsync(CONTROL_RENEWAL_BATCH_INTERVAL_MS)
// One statement updates a host's assignment row once, so the host appears in
// one flush only; the newer generation leads the next one.
expect(renew.mock.calls[0]![0].map((row) => row.activityId)).toEqual([
'control:cell-a:1',
'control:cell-a:1'
])
await expect(Promise.all([first, other])).resolves.toEqual([undefined, undefined])
await vi.advanceTimersByTimeAsync(CONTROL_RENEWAL_BATCH_INTERVAL_MS)
expect(renew).toHaveBeenCalledTimes(2)
expect(renew.mock.calls[1]![0].map((row) => row.activityId)).toEqual(['control:cell-a:2'])
await expect(second).resolves.toBeUndefined()
})
it('stays quiet for a fast flush that renewed everything', async () => {
const flushes: ControlRenewalFlush[] = []
const batch = new ControlRenewalBatch(
async () => ['renewed'],
() => ({ cellId: 'cell-a' }),
(flush) => flushes.push(flush)
)
const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined)
try {
void batch.enqueue(request('host0000000000f1'))
await vi.advanceTimersByTimeAsync(CONTROL_RENEWAL_BATCH_INTERVAL_MS)
expect(warn).not.toHaveBeenCalled()
expect(flushes).toEqual([
{ rows: 1, durationMs: expect.any(Number), outcomes: { renewed: 1 } }
])
} finally {
warn.mockRestore()
}
})
it('logs one line with the outcome counts when a flush did not renew everything', async () => {
const batch = new ControlRenewalBatch(
async () => ['renewed', 'control_activity_not_found'],
() => ({ cellId: 'cell-a' })
)
const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined)
try {
void outcomeOf(batch.enqueue(request('host0000000000g1')))
const missing = outcomeOf(batch.enqueue(request('host0000000000g2')))
await vi.advanceTimersByTimeAsync(CONTROL_RENEWAL_BATCH_INTERVAL_MS)
await expect(missing).resolves.toBe('control_activity_not_found')
expect(warn).toHaveBeenCalledOnce()
expect(JSON.parse(String(warn.mock.calls[0]![0]))).toMatchObject({
event: 'orca_relay_control_renewal_flush',
cellId: 'cell-a',
rows: 2,
outcomes: { renewed: 1, control_activity_not_found: 1 }
})
} finally {
warn.mockRestore()
}
})
})
cloud/apps/relay/src/control-renewal-batch.ts
+148
View File
@@ -0,0 +1,148 @@
import { performance } from 'node:perf_hooks'
import type {
ControlRenewalOutcome,
ControlRenewalRequest
} from './control-renewal-statement.js'
// One flush per second turns the fleet's control-lease write rate into a
// function of the cell count rather than the host count: a cell's ~10 due
// renewals per second become one write transaction instead of ten. Well inside
// the 105s lease runway, so a host that misses a window is never at risk.
export const CONTROL_RENEWAL_BATCH_INTERVAL_MS = 1_000
// Ceiling on the parameter arrays. Row locks live until the statement commits,
// so this is what bounds how long one flush holds them: measured at 11.5ms for
// 200 rows against a 20,000-row table, and 9.4ms with a host wedged in a
// per-host transaction.
export const CONTROL_RENEWAL_BATCH_MAX_ROWS = 200
// A flush slower than this is the only latency worth a line; the metrics event
// carries the distribution.
const CONTROL_RENEWAL_SLOW_FLUSH_MS = 250
export type ControlRenewalFlush = {
rows: number
durationMs: number
outcomes: Record<string, number>
}
type PendingWaiter = { resolve: () => void; reject: (error: unknown) => void }
type PendingRenewal = { request: ControlRenewalRequest; waiters: PendingWaiter[] }
type QueuedRenewal = { request: ControlRenewalRequest; waiter: PendingWaiter }
// Per host, not per activity: one statement updates a host's assignment row
// once, so two activities for the same host must not share a flush.
function pendingKey(request: ControlRenewalRequest): string {
return [request.identity.userId, request.identity.relayHostId].join('\u0000')
}
// Collects the control-lease renewals a cell owes and spends one statement on
// them. Each caller still gets the single-renewal contract: the promise resolves
// on `renewed` and rejects with the outcome as its message otherwise, so callers
// keep their per-session error routing unchanged.
export class ControlRenewalBatch {
private pending = new Map<string, PendingRenewal>()
// Renewals a host cannot contribute to the flush being built; they open the
// next one.
private deferred: QueuedRenewal[] = []
private timer: ReturnType<typeof setTimeout> | null = null
constructor(
private readonly renew: (
rows: readonly ControlRenewalRequest[]
) => Promise<ControlRenewalOutcome[]>,
private readonly logFields: () => Record<string, unknown> = () => ({}),
private readonly observe?: (flush: ControlRenewalFlush) => void
) {}
enqueue(request: ControlRenewalRequest): Promise<void> {
return new Promise<void>((resolve, reject) => {
this.admit({ request, waiter: { resolve, reject } })
})
}
private admit(queued: QueuedRenewal): void {
const key = pendingKey(queued.request)
const existing = this.pending.get(key)
if (existing && existing.request.activityId !== queued.request.activityId) {
this.deferred.push(queued)
this.scheduleFlush()
return
}
if (existing) {
// A second attempt at the same lease inside one window supersedes the
// first expiry; both callers still hear the outcome they waited for.
existing.request = {
...queued.request,
expiresAt: Math.max(existing.request.expiresAt, queued.request.expiresAt)
}
existing.waiters.push(queued.waiter)
return
}
this.pending.set(key, { request: queued.request, waiters: [queued.waiter] })
if (this.pending.size >= CONTROL_RENEWAL_BATCH_MAX_ROWS) {
void this.flush()
return
}
this.scheduleFlush()
}
private scheduleFlush(): void {
this.timer ??= setTimeout(() => {
this.timer = null
void this.flush()
}, CONTROL_RENEWAL_BATCH_INTERVAL_MS)
this.timer.unref?.()
}
// Flushes run concurrently on purpose: a statement stalled in PostgreSQL must
// not hold back the renewals that came due while it was waiting.
async flush(): Promise<void> {
if (this.timer) {
clearTimeout(this.timer)
this.timer = null
}
const batch = [...this.pending.values()]
this.pending = new Map()
// Re-admitted against the empty map, so a host deferred out of this flush
// leads the next one.
const deferred = this.deferred
this.deferred = []
for (const queued of deferred) this.admit(queued)
if (batch.length === 0) return
const startedAt = performance.now()
let outcomes: ControlRenewalOutcome[]
try {
outcomes = await this.renew(batch.map((entry) => entry.request))
} catch (error) {
for (const entry of batch) for (const waiter of entry.waiters) waiter.reject(error)
this.report(batch.length, performance.now() - startedAt, { flush_failed: batch.length })
return
}
const counts: Record<string, number> = {}
for (const [index, entry] of batch.entries()) {
const outcome = outcomes[index] ?? 'database_error'
counts[outcome] = (counts[outcome] ?? 0) + 1
for (const waiter of entry.waiters) {
if (outcome === 'renewed') waiter.resolve()
else waiter.reject(new Error(outcome))
}
}
this.report(batch.length, performance.now() - startedAt, counts)
}
private report(rows: number, durationMs: number, outcomes: Record<string, number>): void {
this.observe?.({ rows, durationMs, outcomes })
const renewed = outcomes.renewed ?? 0
if (durationMs <= CONTROL_RENEWAL_SLOW_FLUSH_MS && renewed === rows) return
console.warn(
JSON.stringify({
event: 'orca_relay_control_renewal_flush',
...this.logFields(),
rows,
durationMs: Math.round(durationMs),
outcomes
})
)
}
}
cloud/apps/relay/src/control-renewal-postgres.test.ts
+166 -3
View File
@@ -1,8 +1,11 @@
import { performance } from 'node:perf_hooks'
import { ASSIGNMENT_LIMITS, RELAY_PROTOCOL_LIMITS } from '@orca-cloud/relay-contract'
import { afterAll, beforeAll, describe, expect, it } from 'vitest'
import { RelayAssignmentStore } from './assignment-store.js'
import { CONTROL_RENEWAL_BATCH_SQL } from './control-renewal-statement.js'
import {
openRelayDatabase,
POSTGRES_LOCK_TIMEOUT_MS,
type RelayDatabase,
type RelayLockOptions,
type SqlRow
@@ -22,7 +25,9 @@ const targetCell = {
capacityRequests: 100
}
const userId = 'control-renewal-postgres-user'
const identities = Array.from({ length: 6 }, (_, index) => ({
// Indexes 0-5 belong to the single-renewal cases below, which mutate their
// host's migration and lease state; the batch cases own 6-13.
const identities = Array.from({ length: 14 }, (_, index) => ({
userId,
relayHostId: `controlrenewal${index + 1}`
}))
@@ -72,7 +77,7 @@ class StallFirstRenewalQueryDatabase implements RelayDatabase {
constructor(private readonly database: RelayDatabase) {}
async query(sql: string, params?: unknown[]): Promise<SqlRow[]> {
if (this.stallNext && sql.includes('WITH assignment_state AS MATERIALIZED')) {
if (this.stallNext && sql === CONTROL_RENEWAL_BATCH_SQL) {
this.stallNext = false
this.stalled.resolve()
await this.continue.promise
@@ -103,7 +108,7 @@ class RenewalQueryProbeDatabase implements RelayDatabase {
constructor(private readonly database: RelayDatabase) {}
async query(sql: string, params?: unknown[]): Promise<SqlRow[]> {
if (sql.includes('WITH assignment_state AS MATERIALIZED')) this.renewalQueries++
if (sql === CONTROL_RENEWAL_BATCH_SQL) this.renewalQueries++
return await this.database.query(sql, params)
}
@@ -332,6 +337,164 @@ describePostgres('PostgreSQL control renewal', () => {
).rejects.toThrow('invalid_activity_expiry')
})
it('renews every due host in one autocommitted statement', async () => {
const probe = new RenewalQueryProbeDatabase(database)
const store = new RelayAssignmentStore(probe, () => now)
const batch = identities.slice(6, 10)
now += 30_000
const expiresAt = now + 105_000
const outcomes = await store.renewControlActivities(
batch.map((identity) => ({
identity,
activityId: controlId(sourceCell.id),
cellId: sourceCell.id,
expiresAt
}))
)
expect(outcomes).toEqual(['renewed', 'renewed', 'renewed', 'renewed'])
expect(probe.renewalQueries).toBe(1)
expect(probe.transactions).toBe(0)
const leases = await database.query(
`SELECT relay_host_id, expires_at FROM relay_assignment_activity_leases
WHERE user_id = ? AND activity_id = ? ORDER BY relay_host_id ASC`,
[userId, controlId(sourceCell.id)]
)
expect(
leases
.filter((lease) =>
batch.some((identity) => identity.relayHostId === lease.relay_host_id)
)
.map((lease) => Number(lease.expires_at))
).toEqual([expiresAt, expiresAt, expiresAt, expiresAt])
})
it('reports each host its own verdict inside one batch', async () => {
const store = new RelayAssignmentStore(database, () => now)
now += 30_000
const expiresAt = now + 105_000
const live = identities[10]!
const absent = { userId, relayHostId: 'controlrenewalgone' }
const outcomes = await store.renewControlActivities([
{ identity: absent, activityId: controlId(sourceCell.id), cellId: sourceCell.id, expiresAt },
{ identity: live, activityId: controlId(sourceCell.id), cellId: sourceCell.id, expiresAt },
{
identity: live,
activityId: controlId(targetCell.id),
cellId: targetCell.id,
expiresAt
}
])
expect(outcomes).toEqual([
'assignment_not_found',
'renewed',
'activity_cell_not_authoritative'
])
})
it('passes over a host whose assignment row is held and renews the rest', async () => {
const store = new RelayAssignmentStore(database, () => now)
now += 30_000
const expiresAt = now + 105_000
const held = identities[11]!
const free = identities[12]!
const locked = signal()
const release = signal()
// Holds the row the way every per-host transactional path does.
const holder = database.transaction(async (transaction) => {
await transaction.queryLocked(
`SELECT * FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`,
[held.userId, held.relayHostId]
)
locked.resolve()
await release.promise
})
await locked.promise
const startedAt = performance.now()
const outcomes = await store.renewControlActivities(
[held, free].map((identity) => ({
identity,
activityId: controlId(sourceCell.id),
cellId: sourceCell.id,
expiresAt
}))
)
const elapsedMs = performance.now() - startedAt
release.resolve()
await holder
expect(outcomes).toEqual(['assignment_lock_unavailable', 'renewed'])
// It skipped rather than queued: a blocking FOR UPDATE would have spent the
// pool's whole lock_timeout here and failed the free host too.
expect(elapsedMs).toBeLessThan(POSTGRES_LOCK_TIMEOUT_MS)
const lease = (
await database.query(
`SELECT expires_at FROM relay_assignment_activity_leases
WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`,
[free.userId, free.relayHostId, controlId(sourceCell.id)]
)
)[0]
expect(Number(lease!.expires_at)).toBe(expiresAt)
})
it('renews both of one host\u2019s control leases in a single batch', async () => {
const store = new RelayAssignmentStore(database, () => now)
const identity = identities[13]!
// Two live control leases on one host. Written directly because
// activateControl retires the prior generation, and what is under test is the
// statement's row-wise behaviour, not how the second lease came to exist.
await database.query(
`INSERT INTO relay_assignment_activity_leases
(user_id, relay_host_id, activity_id, activity_kind, cell_id,
request_units, expires_at, updated_at)
VALUES (?, ?, ?, 'control', ?, 1, ?, ?)`,
[
identity.userId,
identity.relayHostId,
`control:${sourceCell.id}:2`,
sourceCell.id,
now,
now
]
)
now += 30_000
const expiresAt = now + 105_000
const outcomes = await store.renewControlActivities(
[1, 2].map((generation) => ({
identity,
activityId: `control:${sourceCell.id}:${generation}`,
cellId: sourceCell.id,
expiresAt: expiresAt - generation
}))
)
expect(outcomes).toEqual(['renewed', 'renewed'])
const leases = await database.query(
`SELECT activity_id, expires_at FROM relay_assignment_activity_leases
WHERE user_id = ? AND relay_host_id = ? ORDER BY activity_id ASC`,
[identity.userId, identity.relayHostId]
)
expect(leases.map((lease) => Number(lease.expires_at))).toEqual([
expiresAt - 1,
expiresAt - 2
])
// The assignment row is written once, carrying the later of the two.
const row = (
await database.query(
`SELECT lease_expires_at, last_activity_at FROM relay_assignments
WHERE user_id = ? AND relay_host_id = ?`,
[identity.userId, identity.relayHostId]
)
)[0]
expect(Number(row!.lease_expires_at)).toBe(expiresAt - 1)
expect(Number(row!.last_activity_at)).toBe(now)
})
it('uses one autocommitted PostgreSQL statement for a steady renewal', async () => {
const probe = new RenewalQueryProbeDatabase(database)
const store = new RelayAssignmentStore(probe, () => now)
cloud/apps/relay/src/control-renewal-statement.ts
+230
View File
@@ -0,0 +1,230 @@
import type { AssignmentIdentity } from './assignment-identity-queue.js'
import type { SqlRow } from './database.js'
export type ControlRenewalOutcome =
| 'renewed'
| 'assignment_not_found'
| 'activity_cell_not_authoritative'
| 'control_activity_not_found'
| 'control_activity_moved'
// The host's assignment row was already locked by one of the per-host
// transactional paths. Retryable, and never a reason to close a control: the
// next tick is 15s away and the lease has 105s on it.
| 'assignment_lock_unavailable'
// Decided per row before the statement runs, so one malformed request cannot
// cost the rest of the batch its renewal.
| 'invalid_activity_id'
| 'invalid_activity_expiry'
| 'database_error'
// Outcomes the statement itself can report. `database_error` is raised by the
// driver, and `invalid_activity_expiry` is decided per row before the statement
// is built, so neither can come back as a row.
export const CONTROL_RENEWAL_STATEMENT_OUTCOMES = new Set<ControlRenewalOutcome>([
'renewed',
'assignment_not_found',
'activity_cell_not_authoritative',
'control_activity_not_found',
'control_activity_moved',
'assignment_lock_unavailable'
])
export type ControlRenewalRequest = {
identity: AssignmentIdentity
activityId: string
cellId: string
expiresAt: number
}
// LOCK ORDER - (user_id, relay_host_id), the primary key of relay_assignments,
// applied here and repeated as the statement's ORDER BY so it holds whether the
// planner walks the primary-key index or sorts under the LockRows node.
//
// The batch never waits for an assignment row: SKIP LOCKED reports a contended
// host separately instead. That is what bounds how long a flush holds its locks
// to its own execution time, because row locks live until the statement commits,
// and it is why one host wedged in a per-host transaction cannot stall the
// renewals of every other host sharing the flush.
//
// With no wait on the assignment pass, the deadlock question reduces to the two
// later passes. Every writer in this store locks a host's assignment row before
// that host's lease rows (`assignmentRow` then `lockAssignmentActivities`), and
// a host whose assignment row is held was skipped, so the batch never reaches
// that host's lease: the lease pass cannot wait either.
// `markMigrationTargetRegistered` is the one writer that locks a migration row
// without the assignment row first. It takes no further locks, so it can delay a
// mid-migration row by up to the pool's lock_timeout but cannot close a cycle.
export function orderedControlRenewalRows<Row extends { identity: AssignmentIdentity }>(
rows: readonly Row[]
): Row[] {
return [...rows].sort(
(left, right) =>
left.identity.userId.localeCompare(right.identity.userId) ||
left.identity.relayHostId.localeCompare(right.identity.relayHostId)
)
}
// One statement renewing every due control lease on this cell, row-wise over the
// unnested parameter arrays. Logic per row is what the single-row predecessor
// did: lock the assignment, admit the caller's cell either as the current cell or
// as the source of an active forward migration, lock that host's control lease,
// push both expiries forward, and report one outcome. The one addition is
// `present_assignment`, an unlocked probe that separates a host with no
// assignment row at all from one whose row SKIP LOCKED passed over - the first
// closes the control, the second retries.
export const CONTROL_RENEWAL_BATCH_SQL = `WITH renewal_input AS MATERIALIZED (
SELECT
renewal.ordinality AS row_index,
renewal.user_id,
renewal.relay_host_id,
renewal.activity_id,
renewal.cell_id,
renewal.expires_at
FROM unnest(?::text[], ?::text[], ?::text[], ?::text[], ?::bigint[])
WITH ORDINALITY AS renewal(
user_id, relay_host_id, activity_id, cell_id, expires_at, ordinality
)
), present_assignment AS MATERIALIZED (
SELECT input.row_index
FROM renewal_input input
JOIN relay_assignments assignment
ON assignment.user_id = input.user_id
AND assignment.relay_host_id = input.relay_host_id
), assignment_state AS MATERIALIZED (
SELECT input.row_index, assignment.cell_id, assignment.assignment_epoch
FROM renewal_input input
JOIN relay_assignments assignment
ON assignment.user_id = input.user_id
AND assignment.relay_host_id = input.relay_host_id
ORDER BY assignment.user_id, assignment.relay_host_id
FOR UPDATE OF assignment SKIP LOCKED
), migration_state AS MATERIALIZED (
SELECT locked.row_index
FROM assignment_state locked
JOIN renewal_input input ON input.row_index = locked.row_index
JOIN relay_assignment_migrations migration
ON migration.user_id = input.user_id
AND migration.relay_host_id = input.relay_host_id
AND migration.source_cell_id = input.cell_id
AND migration.target_cell_id = locked.cell_id
AND migration.assignment_epoch = locked.assignment_epoch
AND migration.completed_at IS NULL AND migration.aborted_at IS NULL
ORDER BY migration.user_id, migration.relay_host_id
FOR UPDATE OF migration
), authorization_state AS MATERIALIZED (
SELECT locked.row_index
FROM assignment_state locked
JOIN renewal_input input ON input.row_index = locked.row_index
WHERE locked.cell_id = input.cell_id
OR EXISTS (
SELECT 1 FROM migration_state moving
WHERE moving.row_index = locked.row_index
)
), lease_state AS MATERIALIZED (
SELECT authorized.row_index, lease.activity_kind, lease.cell_id
FROM authorization_state authorized
JOIN renewal_input input ON input.row_index = authorized.row_index
JOIN relay_assignment_activity_leases lease
ON lease.user_id = input.user_id
AND lease.relay_host_id = input.relay_host_id
AND lease.activity_id = input.activity_id
ORDER BY lease.user_id, lease.relay_host_id, lease.activity_id
FOR UPDATE OF lease
), renewed_lease AS (
UPDATE relay_assignment_activity_leases lease
SET expires_at = GREATEST(lease.expires_at, input.expires_at),
updated_at = GREATEST(lease.updated_at, ?)
FROM lease_state state
JOIN renewal_input input ON input.row_index = state.row_index
WHERE lease.user_id = input.user_id
AND lease.relay_host_id = input.relay_host_id
AND lease.activity_id = input.activity_id
AND state.activity_kind = 'control' AND state.cell_id = input.cell_id
RETURNING state.row_index
), renewed_assignment AS (
-- Grouped per host: an UPDATE whose FROM offers a target row more than
-- once applies one source row and returns one, so two leases on one
-- host would leave the assignment carrying the wrong expiry. The
-- aggregate hands it exactly one row, carrying the later expiry.
UPDATE relay_assignments assignment
SET lease_expires_at = GREATEST(assignment.lease_expires_at, renewed.expires_at),
last_activity_at = GREATEST(assignment.last_activity_at, ?)
FROM (
SELECT input.user_id, input.relay_host_id, MAX(input.expires_at) AS expires_at
FROM renewed_lease renewed
JOIN renewal_input input ON input.row_index = renewed.row_index
GROUP BY input.user_id, input.relay_host_id
) renewed
WHERE assignment.user_id = renewed.user_id
AND assignment.relay_host_id = renewed.relay_host_id
RETURNING renewed.user_id
)
SELECT input.row_index, CASE
WHEN NOT EXISTS (
SELECT 1 FROM present_assignment present
WHERE present.row_index = input.row_index
) THEN 'assignment_not_found'
WHEN NOT EXISTS (
SELECT 1 FROM assignment_state locked WHERE locked.row_index = input.row_index
) THEN 'assignment_lock_unavailable'
WHEN NOT EXISTS (
SELECT 1 FROM authorization_state authorized
WHERE authorized.row_index = input.row_index
) THEN 'activity_cell_not_authoritative'
WHEN NOT EXISTS (
SELECT 1 FROM lease_state state WHERE state.row_index = input.row_index
) THEN 'control_activity_not_found'
WHEN EXISTS (
SELECT 1 FROM lease_state state
WHERE state.row_index = input.row_index
AND (state.activity_kind <> 'control' OR state.cell_id <> input.cell_id)
) THEN 'control_activity_moved'
-- Read from renewed_lease, which has one row per input row. The
-- assignment update collapses to one row per host, so it cannot answer
-- for a host that brought two leases to the same batch.
WHEN EXISTS (
SELECT 1 FROM renewed_lease renewed
WHERE renewed.row_index = input.row_index
) THEN 'renewed'
ELSE 'control_activity_not_found'
END AS outcome
FROM renewal_input input
ORDER BY input.row_index`
export function controlRenewalBatchParams(
rows: readonly ControlRenewalRequest[],
now: number
): unknown[] {
return [
rows.map((row) => row.identity.userId),
rows.map((row) => row.identity.relayHostId),
rows.map((row) => row.activityId),
rows.map((row) => row.cellId),
rows.map((row) => row.expiresAt),
now,
now
]
}
// Rows come back ordered by row_index, which is the 1-based position in the
// statement's parameter arrays.
export function readControlRenewalOutcomes(
rows: SqlRow[],
expected: number
): ControlRenewalOutcome[] {
if (rows.length !== expected) throw new Error('missing_control_renewal_outcome')
return rows.map((row, position) => {
if (Number(row.row_index) !== position + 1) {
throw new Error('misordered_control_renewal_outcome')
}
const outcome = row.outcome
if (
typeof outcome !== 'string' ||
!CONTROL_RENEWAL_STATEMENT_OUTCOMES.has(outcome as ControlRenewalOutcome)
) {
throw new Error('invalid_control_renewal_outcome')
}
// SAFETY: the membership check above is what narrows this string.
return outcome as ControlRenewalOutcome
})
}
cloud/apps/relay/src/host-session-registry.test.ts
+109 -1
View File
@@ -12,6 +12,12 @@ import type WebSocket from 'ws'
import type { RelayAssignmentStore } from './assignment-store.js'
import type { RelayConfig } from './config.js'
import type { RelayCredentialStore } from './credential-store.js'
import { CONTROL_RENEWAL_BATCH_INTERVAL_MS } from './control-renewal-batch.js'
import {
CONTROL_RENEWAL_STATEMENT_OUTCOMES,
type ControlRenewalOutcome,
type ControlRenewalRequest
} from './control-renewal-statement.js'
import { HostSessionRegistry, type HostSession } from './host-session-registry.js'
import { relayHostLogDigest } from './relay-host-log-digest.js'
import type { RelayRuntimeObserver } from './relay-observability.js'
@@ -23,6 +29,12 @@ import {
import type { RelayTokenClaims } from './relay-token-verifier.js'
import { ProcessQueuedByteBudget } from './splice-forwarder.js'
// A due renewal leaves the heartbeat as a batch enqueue, so the store only sees
// the tick once the batch window closes.
async function closeRenewalWindow(): Promise<void> {
await vi.advanceTimersByTimeAsync(CONTROL_RENEWAL_BATCH_INTERVAL_MS)
}
class FakeSocket extends EventEmitter {
readonly OPEN = 1
readonly CLOSING = 2
@@ -109,6 +121,7 @@ function createRegistry(
activate: ActivateSession
acquireActivity: ReturnType<typeof vi.fn>
renewControlActivity: ReturnType<typeof vi.fn>
renewControlActivities: ReturnType<typeof vi.fn>
releaseActivity: ReturnType<typeof vi.fn>
observer: {
recordAuth: ReturnType<typeof vi.fn>
@@ -119,12 +132,38 @@ function createRegistry(
const acquireActivity = vi.fn().mockResolvedValue(undefined)
const renewControlActivity = vi.fn().mockResolvedValue(undefined)
const releaseActivity = vi.fn().mockResolvedValue(true)
// Mirrors the store's own batch semantics over the single-renewal mock: a known
// outcome becomes that row's verdict, and any other failure reaches the caller
// as the driver's error. Keeps every per-call expectation below aimed at the
// renewal a session actually asked for.
const renewControlActivities = vi.fn(
async (rows: readonly ControlRenewalRequest[]): Promise<ControlRenewalOutcome[]> =>
await Promise.all(
rows.map(async (row): Promise<ControlRenewalOutcome> => {
try {
await renewControlActivity(row.identity, {
activityId: row.activityId,
cellId: row.cellId,
expiresAt: row.expiresAt
})
return 'renewed'
} catch (error) {
const message = String((error as { message?: unknown }).message)
if (!CONTROL_RENEWAL_STATEMENT_OUTCOMES.has(message as ControlRenewalOutcome)) {
throw error
}
return message as ControlRenewalOutcome
}
})
)
)
const assignments = {
activateControl,
markMigrationTargetRegistered: vi.fn().mockResolvedValue(undefined),
resolve: vi.fn().mockResolvedValue({ cellId: config.cellId }),
acquireActivity,
renewControlActivity,
renewControlActivities,
releaseActivity
} as unknown as RelayAssignmentStore
const observer = {
@@ -176,6 +215,7 @@ function createRegistry(
activate,
acquireActivity,
renewControlActivity,
renewControlActivities,
releaseActivity,
observer
}
@@ -689,7 +729,8 @@ describe('host session cleanup races', () => {
expect(original).not.toBeNull()
await activate(new FakeSocket() as unknown as WebSocket, identity, original, 2, false, 1)
vi.advanceTimersByTime(15_000)
await vi.advanceTimersByTimeAsync(15_000)
await closeRenewalWindow()
expect(renewControlActivity).toHaveBeenCalledOnce()
expect(renewControlActivity).toHaveBeenCalledWith(
@@ -715,6 +756,7 @@ describe('host session cleanup races', () => {
})
)
await vi.advanceTimersByTimeAsync(15_000)
await closeRenewalWindow()
const replacement = new FakeSocket()
await h.activate(replacement as unknown as WebSocket, identity, session, 1, true, 1)
reject(new Error('activity_cell_not_authoritative'))
@@ -737,6 +779,7 @@ describe('host session cleanup races', () => {
})
)
await vi.advanceTimersByTimeAsync(15_000)
await closeRenewalWindow()
h.registry.drainHost({
attemptId: 'attempt',
userId: identity.sub,
@@ -762,6 +805,7 @@ describe('host session cleanup races', () => {
for (let interval = 0; interval < 4; interval++) {
await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs)
socket.emit('message', Buffer.from(JSON.stringify({ type: 'pong' })), false)
await closeRenewalWindow()
}
const pings = socket.send.mock.calls.filter((call) => String(call[0]).includes('"ping"'))
@@ -791,8 +835,10 @@ describe('host session cleanup races', () => {
try {
await activate(socket as unknown as WebSocket, identity, null, 1, false, 1)
await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs)
await closeRenewalWindow()
expect(renewControlActivity).toHaveBeenCalledOnce()
await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs)
await closeRenewalWindow()
expect(renewControlActivity).toHaveBeenCalledTimes(2)
} finally {
warn.mockRestore()
@@ -812,6 +858,7 @@ describe('host session cleanup races', () => {
await activate(socket as unknown as WebSocket, identity, null, 1, false, 1)
await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs * 2)
await closeRenewalWindow()
expect(renewControlActivity).toHaveBeenCalledTimes(2)
stalled.resolve(undefined)
@@ -831,13 +878,16 @@ describe('host session cleanup races', () => {
await activate(socket as unknown as WebSocket, identity, null, 1, false, 1)
await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs * 2)
await closeRenewalWindow()
expect(renewControlActivity).toHaveBeenCalledTimes(2)
stalled.resolve(undefined)
await vi.advanceTimersByTimeAsync(0)
await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs)
await closeRenewalWindow()
expect(renewControlActivity).toHaveBeenCalledTimes(2)
await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs)
await closeRenewalWindow()
expect(renewControlActivity).toHaveBeenCalledTimes(3)
registry.drain(0)
@@ -855,6 +905,7 @@ describe('host session cleanup races', () => {
await activate(socket as unknown as WebSocket, identity, null, 1, false, 1)
await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs)
await closeRenewalWindow()
expect(acquireActivity).toHaveBeenCalledWith(
{ userId: identity.sub, relayHostId: identity.relayHostId },
@@ -880,6 +931,7 @@ describe('host session cleanup races', () => {
await activate(socket as unknown as WebSocket, identity, null, 1, false, 1)
await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs)
await closeRenewalWindow()
expect(acquireActivity).not.toHaveBeenCalled()
expect(socket.close).toHaveBeenCalledWith(RELAY_CLOSE_CODE.DRAINING, 'control activity moved')
@@ -899,6 +951,7 @@ describe('host session cleanup races', () => {
await activate(socket as unknown as WebSocket, identity, null, 1, false, 1)
await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs)
await closeRenewalWindow()
expect(socket.close).toHaveBeenCalledWith(
RELAY_CLOSE_CODE.DRAINING,
@@ -918,6 +971,7 @@ describe('host session cleanup races', () => {
await activate(socket as unknown as WebSocket, identity, null, 1, false, 1)
await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs)
await closeRenewalWindow()
expect(socket.close).toHaveBeenCalledWith(
RELAY_CLOSE_CODE.DRAINING,
@@ -939,6 +993,7 @@ describe('host session cleanup races', () => {
for (let interval = 0; interval < 3; interval++) {
await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs)
socket.emit('message', Buffer.from(JSON.stringify({ type: 'pong' })), false)
await closeRenewalWindow()
}
expect(renewControlActivity).toHaveBeenCalledTimes(2)
@@ -966,6 +1021,7 @@ describe('control renewal cadence across a rebind', () => {
const beat = async (target: FakeSocket): Promise<void> => {
await vi.advanceTimersByTimeAsync(ping)
target.emit('message', Buffer.from(JSON.stringify({ type: 'pong' })), false)
await closeRenewalWindow()
}
// Age the session so its attempt counter is well above zero.
@@ -993,6 +1049,57 @@ describe('control renewal cadence across a rebind', () => {
})
})
describe('control renewals shared by one batch', () => {
beforeEach(() => vi.useFakeTimers())
afterEach(() => {
vi.clearAllTimers()
vi.useRealTimers()
})
it('renews two due hosts in one call and leaves a stale one alone', async () => {
const activateControl = vi
.fn<RelayAssignmentStore['activateControl']>()
.mockResolvedValueOnce('control:production-gce-c3:1')
.mockResolvedValueOnce('control:production-gce-c3:1')
const { registry, activate, renewControlActivities } = createRegistry(activateControl)
const other = { ...identity, sub: 'user-2', relayHostId: 'ponmlkjihgfedcba' }
const staleSocket = new FakeSocket()
const liveSocket = new FakeSocket()
await activate(staleSocket as unknown as WebSocket, identity, null, 1, false, 1)
await activate(liveSocket as unknown as WebSocket, other, null, 1, false, 1)
const stale = registry.get({ userId: identity.sub, relayHostId: identity.relayHostId })!
const live = registry.get({ userId: other.sub, relayHostId: other.relayHostId })!
// Both come due inside the same window, and one socket goes away while the
// statement is still in PostgreSQL.
let release!: () => void
renewControlActivities.mockImplementationOnce(
async (rows: readonly ControlRenewalRequest[]) => {
staleSocket.close()
await new Promise<void>((resolve) => (release = resolve))
return rows.map((): ControlRenewalOutcome => 'renewed')
}
)
await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs)
const staleDueAt = stale.activityRenewalDueAt
await closeRenewalWindow()
release()
await vi.advanceTimersByTimeAsync(0)
expect(renewControlActivities).toHaveBeenCalledOnce()
expect(
renewControlActivities.mock.calls[0]![0].map(
(row: ControlRenewalRequest) => row.identity.relayHostId
)
).toEqual([identity.relayHostId, other.relayHostId])
expect(live.activityRenewalCompletedAttempt).toBe(1)
expect(stale.activityRenewalCompletedAttempt).toBe(0)
expect(stale.activityRenewalDueAt).toBe(staleDueAt)
registry.drain(0)
vi.advanceTimersByTime(0)
})
})
describe('control lease recovery after the session is gone', () => {
beforeEach(() => vi.useFakeTimers())
afterEach(() => {
@@ -1019,6 +1126,7 @@ describe('control lease recovery after the session is gone', () => {
new Promise<void>((_resolve, reject) => (failRenewal = reject))
)
await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs)
await closeRenewalWindow()
expect(renewControlActivity).toHaveBeenCalledOnce()
const newer = new FakeSocket()
cloud/apps/relay/src/host-session-registry.ts
+24 -9
View File
@@ -26,6 +26,7 @@ import type WebSocket from 'ws'
import type { RawData } from 'ws'
import type { RelayConfig } from './config.js'
import type { RelayAssignmentStore } from './assignment-store.js'
import { ControlRenewalBatch } from './control-renewal-batch.js'
import { RelayCredentialStore, type CredentialReservation } from './credential-store.js'
import { HostCloseReasonMemory } from './host-close-reason-memory.js'
import { relayHostLogDigest } from './relay-host-log-digest.js'
@@ -303,6 +304,15 @@ export class HostSessionRegistry {
private readonly cellIncarnation?: string
) {}
// Renewals leave the heartbeat as an enqueue: one statement per cell per
// window replaces one write transaction per host, which is what keeps the
// shared PostgreSQL instance out of buffer-header contention.
private readonly controlRenewals = new ControlRenewalBatch(
async (rows) => await this.assignments.renewControlActivities(rows),
() => this.logIdentity(),
(flush) => this.observer.recordControlRenewalFlush?.(flush)
)
// Uniform over [CONTROL_LEASE_MS - jitter, CONTROL_LEASE_MS + jitter).
private controlLeaseExpiresAt(): number {
const offset = Math.floor((this.random() * 2 - 1) * CONTROL_LEASE_JITTER_MS)
@@ -1384,15 +1394,13 @@ export class HostSessionRegistry {
session.controlActivityId === controlActivityId &&
session.authorityRevision === authorityRevision &&
attempt > session.activityRenewalCompletedAttempt
void this.assignments
.renewControlActivity(
{ userId: session.identity.sub, relayHostId: session.relayHostId },
{
activityId: controlActivityId,
cellId: this.config.cellId,
expiresAt: startedAt + CONTROL_ACTIVITY_LEASE_MS
}
)
void this.controlRenewals
.enqueue({
identity: { userId: session.identity.sub, relayHostId: session.relayHostId },
activityId: controlActivityId,
cellId: this.config.cellId,
expiresAt: startedAt + CONTROL_ACTIVITY_LEASE_MS
})
.then(() => {
if (!current()) return
session.activityRenewalCompletedAttempt = attempt
@@ -1458,6 +1466,13 @@ export class HostSessionRegistry {
session.socket?.close(RELAY_CLOSE_CODE.DRAINING, 'control activity moved')
return
}
if (error instanceof Error && error.message === 'assignment_lock_unavailable') {
// A per-host transaction held the row, so the batch passed over it
// rather than making every other host in the flush wait. The next
// tick is 15s away against a 105s lease, and the flush line already
// reports the count, so this needs no line of its own.
return
}
console.warn('[orca-relay] control activity renewal failed')
})
// Terminal handler: a throw inside the async catch above (e.g. a
cloud/apps/relay/src/relay-observability.ts
+22
View File
@@ -1,6 +1,7 @@
import { monitorEventLoopDelay, performance } from 'node:perf_hooks'
import { RELAY_REGION_METRIC_SEGMENTS, type RelayRegion } from '@orca-cloud/relay-contract'
import type { ControlRenewalOutcome } from './assignment-store.js'
import type { ControlRenewalFlush } from './control-renewal-batch.js'
import type { CellInventoryHoldCounts } from './cell-inventory-hold-samples.js'
import type { PostgresPoolPressureCounts } from './postgres-pool-pressure.js'
import type { RelayReadinessGraceEvent, RelayReadinessObservation } from './relay-readiness.js'
@@ -53,6 +54,7 @@ export interface RelayRuntimeObserver {
recordReconnect(): void
recordSql(durationMs: number, success: boolean): void
recordControlRenewal?(durationMs: number, outcome: ControlRenewalOutcome): void
recordControlRenewalFlush?(flush: ControlRenewalFlush): void
recordControlActivityRecovery?(success: boolean): void
recordAssignmentAdmission?(outcome: AssignmentAdmissionOutcome): void
recordAssignmentRejectionReason?(lane: AssignmentAdmissionLane, reason: string): void
@@ -119,6 +121,8 @@ type RelayMetricDeltas = {
controlRttObserved: number
controlRenewalLatenciesMs: number[]
controlRenewalsByOutcome: Record<string, number>
controlRenewalFlushLatenciesMs: number[]
controlRenewalFlushRowsMax: number
controlActivityRecoveries: number
controlActivityRecoveryFailures: number
}
@@ -164,6 +168,8 @@ const emptyDeltas = (): RelayMetricDeltas => ({
controlRttObserved: 0,
controlRenewalLatenciesMs: [],
controlRenewalsByOutcome: {},
controlRenewalFlushLatenciesMs: [],
controlRenewalFlushRowsMax: 0,
controlActivityRecoveries: 0,
controlActivityRecoveryFailures: 0
})
@@ -274,6 +280,14 @@ export class RelayObservability implements RelayRuntimeObserver {
(this.deltas.controlRenewalsByOutcome[outcome] ?? 0) + 1
}
recordControlRenewalFlush(flush: ControlRenewalFlush): void {
this.deltas.controlRenewalFlushLatenciesMs.push(flush.durationMs)
this.deltas.controlRenewalFlushRowsMax = Math.max(
this.deltas.controlRenewalFlushRowsMax,
flush.rows
)
}
recordControlActivityRecovery(success: boolean): void {
if (success) this.deltas.controlActivityRecoveries++
else this.deltas.controlActivityRecoveryFailures++
@@ -379,6 +393,7 @@ export class RelayObservability implements RelayRuntimeObserver {
roundMs(percentile(deltas.clientAcceptStageSamplesMs[stage], 0.95))
const controlRtt = latencySummary(deltas.controlRttSamplesMs)
const controlRenewal = latencySummary(deltas.controlRenewalLatenciesMs)
const controlRenewalFlush = latencySummary(deltas.controlRenewalFlushLatenciesMs)
const memory = process.memoryUsage()
const p99 = this.eventLoop.count === 0 ? 0 : this.eventLoop.percentile(99) / 1_000_000
this.eventLoop.reset()
@@ -447,9 +462,16 @@ export class RelayObservability implements RelayRuntimeObserver {
deltas.controlRenewalsByOutcome.control_activity_not_found ?? 0,
controlActivityRecoveriesDelta: deltas.controlActivityRecoveries,
controlActivityRecoveryFailuresDelta: deltas.controlActivityRecoveryFailures,
// Meaning changed when renewals began batching: for a batched row this is
// the flush's duration, not that row's own statement latency. The
// per-flush fields below are the ones to read for statement cost.
controlRenewalLatencyMsP50: controlRenewal.p50,
controlRenewalLatencyMsP95: controlRenewal.p95,
controlRenewalLatencyMsMax: controlRenewal.max,
controlRenewalFlushesDelta: deltas.controlRenewalFlushLatenciesMs.length,
controlRenewalFlushRowsMax: deltas.controlRenewalFlushRowsMax,
controlRenewalFlushLatencyMsP95: controlRenewalFlush.p95,
controlRenewalFlushLatencyMsMax: controlRenewalFlush.max,
httpLatencyMsMax: roundMs(deltas.httpLatencyMsMax),
heapUsedBytes: memory.heapUsed,
heapTotalBytes: memory.heapTotal,
cloud/infra/terraform/relay-observability.tf
+4
View File
@@ -60,6 +60,10 @@ locals {
control_renewal_latency_ms_p50 = { field = "controlRenewalLatencyMsP50", description = "Control renewal latency p50 in the interval." }
control_renewal_latency_ms_p95 = { field = "controlRenewalLatencyMsP95", description = "Control renewal latency p95 in the interval." }
control_renewal_latency_ms_max = { field = "controlRenewalLatencyMsMax", description = "Maximum control renewal latency in the interval." }
control_renewal_flushes = { field = "controlRenewalFlushesDelta", description = "Batched control-renewal statements issued in the interval, one per cell per flush window." }
control_renewal_flush_rows_max = { field = "controlRenewalFlushRowsMax", description = "Largest number of hosts renewed by a single statement in the interval; the row ceiling is what bounds how long one flush holds its row locks." }
control_renewal_flush_ms_p95 = { field = "controlRenewalFlushLatencyMsP95", description = "Batched control-renewal statement duration p95 in the interval. Row locks live until the statement commits, so this is the lock hold." }
control_renewal_flush_ms_max = { field = "controlRenewalFlushLatencyMsMax", description = "Maximum batched control-renewal statement duration in the interval." }
control_renewals = { field = "controlRenewalsDelta", description = "Control renewal attempts in the interval." }
control_renewal_successes = { field = "controlRenewalSuccessesDelta", description = "Successful control renewals in the interval." }
control_renewal_lease_misses = { field = "controlRenewalLeaseMissesDelta", description = "Control renewals that found their activity lease missing." }