relay: skip the edge row read for hosts above the edge leave ceiling (#41150)

* relay: skip the edge row read for hosts above the edge leave ceiling

With the edge on, every assign of a host outside the cohort ran leaveEdge (a
relay_assignments read) and every resolve ran edgeAssignment, on the director's
3-connection pool, ahead of the in-memory lanes. At 1% this saturated the pool at
peak: acquire timeouts and 503s for GCE hosts, and some cell heartbeats.

New optional edge.json flag leaveCeilingPermille: the highest rollout ever applied.
A host above it, the current rollout and the env percent, on no allow or deny list,
cannot have an edge row, so both reads are skipped. Absent: today's behaviour.
The evacuation sweep reports orca_relay_edge_row_above_ceiling if it finds an edge
row above the ceiling (set too low).

* relay: edge leave ceiling: void a ceiling below the rollout; document operating rules; test reads under it

cloud/apps/relay/src/app.ts

@@ -46,6 +46,7 @@ import {
   EdgeMoveDeferredError,
   EdgeMoveBusyError,
   inEdgeCohort,
+  mayBeOnEdge,
   mapEpochFloor,
   type EdgeAssignments
 } from './edge-relay.js'
@@ -485,8 +486,9 @@ export function createRelayApp(
       edge?.ready === true &&
       body.data.regionCorrection?.action !== 'report' &&
       inEdgeCohort(edge.config, claims.relayHostId, requestedRegion, claims.sub)
-    if (edge?.pushes && !toEdge) {
+    if (edge?.pushes && !toEdge && mayBeOnEdge(edge.config, claims.relayHostId, claims.sub)) {
       // Any lane below may re-place a host whose row is on the edge; its devices come back first.
+      // Hosts that were never in a cohort skip this row read (most asks; the pool is small).
       try {
         leftEdge = await edge.leaveEdge(identity)
       } catch (error) {
@@ -820,7 +822,11 @@ export function createRelayApp(
         relayHostId: body.data.relayHostId
       }
       // A host on the relay edge: the liveness rule below would read its row as no answer.
-      const onEdge = await operations.edgeAssignments?.edgeAssignment(identity)
+      const edgeAssignments = operations.edgeAssignments
+      const onEdge =
+        edgeAssignments && mayBeOnEdge(edgeAssignments.config, identity.relayHostId, identity.userId)
+          ? await edgeAssignments.edgeAssignment(identity)
+          : null
       if (onEdge) {
         return context.json({
           v: 1,

cloud/apps/relay/src/edge-relay-app.test.ts

@@ -1,6 +1,8 @@
 import { afterEach, describe, expect, it, vi } from 'vitest'
 import type { RelayConfig } from './config.js'
 import { openInMemoryRelayDatabase, type RelayDatabase } from './database/database.js'
+import type { EdgeControlFlags } from './edge-relay-control.js'
+import { edgeCohortBucketPermille } from './edge-relay-cohort.js'
 import { EdgeAssignments, type EdgeRelayConfig } from './edge-relay.js'
 
 const fakes = vi.hoisted(() => ({
@@ -52,7 +54,7 @@ function config(): RelayConfig {
   }
 }
 
-async function harness(options: { seeded?: boolean } = {}) {
+async function harness(options: { seeded?: boolean; control?: () => EdgeControlFlags } = {}) {
   const database = await openInMemoryRelayDatabase()
   databases.push(database)
   const edgeConfig: EdgeRelayConfig = {
@@ -60,7 +62,8 @@ async function harness(options: { seeded?: boolean } = {}) {
     rolloutPercent: 0,
     allowHosts: new Set([COHORT_HOST]),
     denyHosts: new Set(),
-    regions: new Set(['us-central1'])
+    regions: new Set(['us-central1']),
+    ...(options.control ? { control: options.control } : {})
   }
   const edge = new EdgeAssignments(database, edgeConfig)
   if (options.seeded !== false) { await edge.ensureEdgeCell() }
@@ -127,4 +130,31 @@ describe('director routes with the relay edge', () => {
     expect(await response.json()).toMatchObject({ cellUrl: 'https://us-c1.relay.example.test' })
     expect(await h.edge.edgeAssignment({ userId: 'user-1', relayHostId: COHORT_HOST })).toBeNull()
   })
+
+  it('skips the edge row read for hosts above the leave ceiling, on assign and resolve', async () => {
+    const flags: EdgeControlFlags = { mode: 'on', rolloutPermille: 0, leaveCeilingPermille: 10 }
+    const h = await harness({ control: () => flags })
+    vi.spyOn(h.edge, 'pushes', 'get').mockReturnValue(true)
+    const leave = vi.spyOn(h.edge, 'leaveEdge').mockResolvedValue(false)
+    const read = vi.spyOn(h.edge, 'edgeAssignment')
+    const hosts = ['outsidehost00001', 'outsidehost00002', 'outsidehost00003', 'outsidehost00004']
+    const [above, other] = hosts.filter((host) => edgeCohortBucketPermille(host) >= 10)
+    expect((await h.assign(above!)).status).toBe(200)
+    await h.resolve(above!)
+    expect(leave).not.toHaveBeenCalled()
+    expect(read).not.toHaveBeenCalled()
+    // Without a ceiling every host reads, as before.
+    delete flags.leaveCeilingPermille
+    expect((await h.assign(other!)).status).toBe(200)
+    await h.resolve(other!)
+    expect(leave).toHaveBeenCalledTimes(1)
+    expect(read).toHaveBeenCalledTimes(1)
+    // With the ceiling, a host under it (here: everyone under 1000) still reads.
+    flags.leaveCeilingPermille = 1000
+    const third = hosts.find((host) => host !== above && host !== other)!
+    expect((await h.assign(third)).status).toBe(200)
+    await h.resolve(third)
+    expect(leave).toHaveBeenCalledTimes(2)
+    expect(read).toHaveBeenCalledTimes(2)
+  })
 })

cloud/apps/relay/src/edge-relay-cohort.test.ts

@@ -0,0 +1,56 @@
+import { describe, expect, it } from 'vitest'
+import { edgeCohortBucketPermille, mayBeOnEdge, type EdgeRelayConfig } from './edge-relay-cohort.js'
+import { parseEdgeControl, type EdgeControlFlags } from './edge-relay-control.js'
+
+function config(flags: EdgeControlFlags, overrides: Partial<EdgeRelayConfig> = {}): EdgeRelayConfig {
+  return {
+    cellUrl: 'https://edge.relay.example.test',
+    rolloutPercent: 0,
+    allowHosts: new Set(),
+    denyHosts: new Set(),
+    regions: new Set(['us-central1']),
+    control: () => flags,
+    ...overrides
+  }
+}
+
+// Hosts with known buckets: one under 10 permille, one at or above 20.
+const hosts = Array.from({ length: 400 }, (_, index) => `ceilinghost${String(index).padStart(5, '0')}`)
+const low = hosts.find((host) => edgeCohortBucketPermille(host) < 10)!
+const high = hosts.find((host) => edgeCohortBucketPermille(host) >= 20)!
+
+describe('mayBeOnEdge', () => {
+  it('is true for every host without a ceiling (today)', () => {
+    expect(mayBeOnEdge(config({ mode: 'on', rolloutPermille: 0 }), high)).toBe(true)
+  })
+
+  it('is false only above the ceiling, the current rollout and the env percent', () => {
+    expect(mayBeOnEdge(config({ mode: 'on', rolloutPermille: 10, leaveCeilingPermille: 10 }), low)).toBe(true)
+    expect(mayBeOnEdge(config({ mode: 'on', rolloutPermille: 10, leaveCeilingPermille: 10 }), high)).toBe(false)
+    // A rollout above a stale ceiling still counts.
+    expect(mayBeOnEdge(config({ mode: 'on', rolloutPermille: 1000, leaveCeilingPermille: 10 }), high)).toBe(true)
+    expect(mayBeOnEdge(config({ mode: 'on', leaveCeilingPermille: 0 }, { rolloutPercent: 100 }), high)).toBe(true)
+  })
+
+  it('keeps allow and deny listed hosts and accounts, and ignores the mode', () => {
+    const flags: EdgeControlFlags = { mode: 'off', leaveCeilingPermille: 0 }
+    expect(mayBeOnEdge(config(flags), low)).toBe(false)
+    expect(mayBeOnEdge(config(flags, { allowHosts: new Set([high]) }), high)).toBe(true)
+    expect(mayBeOnEdge(config(flags, { denyHosts: new Set([high]) }), high)).toBe(true)
+    expect(mayBeOnEdge(config(flags, { allowUsers: new Set(['user-1']) }), high, 'user-1')).toBe(true)
+    expect(mayBeOnEdge(config({ mode: 'down', leaveCeilingPermille: 10 }), low)).toBe(true)
+  })
+})
+
+describe('edge control leaveCeilingPermille', () => {
+  it('parses, and a bad value voids the object', () => {
+    expect(parseEdgeControl({ v: 1, flags: { mode: 'on', rolloutPermille: 10, leaveCeilingPermille: 10 } })).toEqual({
+      flags: { mode: 'on', rolloutPermille: 10, leaveCeilingPermille: 10 },
+      ignoredKeys: []
+    })
+    expect(parseEdgeControl({ v: 1, flags: { mode: 'on' } })).toEqual({ flags: { mode: 'on' }, ignoredKeys: [] })
+    expect(parseEdgeControl({ v: 1, flags: { mode: 'on', leaveCeilingPermille: 1001 } })).toBeNull()
+    // A ceiling below the rollout is a mistake: the last applied flags stay.
+    expect(parseEdgeControl({ v: 1, flags: { mode: 'on', rolloutPermille: 20, leaveCeilingPermille: 10 } })).toBeNull()
+  })
+})

cloud/apps/relay/src/edge-relay-cohort.ts

@@ -66,6 +66,24 @@ export function inEdgeCohort(
   return edgeCohortBucketPermille(relayHostId) < edgeRolloutPermille(config.rolloutPercent, flags)
 }
 
+// Whether the host's assignment row may be on the edge, so an ask that is not moving it there must
+// read the row (leave) first. A host reaches the edge only through the cohort (its bucket under the
+// rollout) or an allow list, and cohorts only grow with the rollout, so a host above the highest
+// rollout ever applied (leaveCeilingPermille) never had an edge row. Unset: every host may.
+// Mode is ignored: hosts leave under off and down too.
+export function mayBeOnEdge(config: EdgeRelayConfig, relayHostId: string, userId?: string): boolean {
+  const flags = config.control?.() ?? EDGE_CONTROL_DEFAULTS
+  if (flags.leaveCeilingPermille === undefined) { return true }
+  if (config.allowHosts.has(relayHostId) || config.denyHosts.has(relayHostId)) { return true }
+  if (userId !== undefined && config.allowUsers?.has(userId)) { return true }
+  const ceiling = Math.max(
+    flags.leaveCeilingPermille,
+    edgeRolloutPermille(config.rolloutPercent, flags),
+    config.rolloutPercent * 10
+  )
+  return edgeCohortBucketPermille(relayHostId) < ceiling
+}
+
 // Another director holds the host's row, or the host is on the edge and cannot leave yet (its
 // desktop is still connected there): retry shortly.
 export class EdgeMoveBusyError extends Error {

cloud/apps/relay/src/edge-relay-control.ts

@@ -17,6 +17,14 @@ export type EdgeControlFlags = {
   mode: EdgeControlMode
   // Replaces the env rollout percent when present (1‰ steps).
   rolloutPermille?: number
+  // Highest rollout (permille) ever applied. Present: a host above it (and above the current
+  // rollout, and on no allow or deny list) cannot have a row on the edge, so its asks skip the
+  // edge row read (mayBeOnEdge). Absent: every host's ask reads it. Operating rules:
+  // - never lower it below a rollout that was applied (a file with it below rolloutPermille is void);
+  // - raise it in its own write, at least a minute (many polls) before a higher rollout, so no
+  //   director still skips a host another director has just moved;
+  // - unset it before removing an allow list entry (the ceiling covers buckets, not the lists).
+  leaveCeilingPermille?: number
 }
 
 export const EDGE_CONTROL_DEFAULTS: EdgeControlFlags = { mode: 'on' }
@@ -25,10 +33,19 @@ export const EDGE_CONTROL_APPLIED_EVENT = 'orca_relay_edge_control_applied'
 
 const EdgeControlFlagsSchema = z.object({
   mode: z.enum(['on', 'off', 'down']).default(EDGE_CONTROL_DEFAULTS.mode),
-  rolloutPermille: z.number().int().min(0).max(1_000).optional()
+  rolloutPermille: z.number().int().min(0).max(1_000).optional(),
+  leaveCeilingPermille: z.number().int().min(0).max(1_000).optional()
 })
 const KNOWN_FLAG_KEYS = new Set(Object.keys(EdgeControlFlagsSchema.shape))
-const EdgeControlObjectSchema = z.object({ v: z.literal(1), flags: EdgeControlFlagsSchema })
+const EdgeControlObjectSchema = z.object({
+  v: z.literal(1),
+  flags: EdgeControlFlagsSchema.refine(
+    (flags) =>
+      flags.leaveCeilingPermille === undefined ||
+      flags.rolloutPermille === undefined ||
+      flags.leaveCeilingPermille >= flags.rolloutPermille
+  )
+})
 
 // Unknown keys (a newer writer's, or a typo) are dropped and reported; a known key with a bad value
 // voids the whole object, so the last applied flags stay.
@@ -42,8 +59,15 @@ export const parseEdgeControl: ControlFlagParse<EdgeControlFlags> = (body) => {
     named !== null && typeof named === 'object'
       ? Object.keys(named).filter((key) => !KNOWN_FLAG_KEYS.has(key)).sort()
       : []
-  const { mode, rolloutPermille } = parsed.data.flags
-  return { flags: rolloutPermille === undefined ? { mode } : { mode, rolloutPermille }, ignoredKeys }
+  const { mode, rolloutPermille, leaveCeilingPermille } = parsed.data.flags
+  return {
+    flags: {
+      mode,
+      ...(rolloutPermille === undefined ? {} : { rolloutPermille }),
+      ...(leaveCeilingPermille === undefined ? {} : { leaveCeilingPermille })
+    },
+    ignoredKeys
+  }
 }
 
 export type EdgeControl = {

cloud/apps/relay/src/edge-relay-evacuation.test.ts

@@ -124,6 +124,18 @@ describe('edge evacuation under runtime control', () => {
     expect(h.admin.drained.sort()).toEqual([hostId(1), hostId(2)])
   })
 
+  it('reports an edge row above the leave ceiling (the ceiling was set too low) and still drains it', async () => {
+    vi.spyOn(console, 'warn').mockImplementation(() => {})
+    const errors = vi.spyOn(console, 'error').mockImplementation(() => {})
+    const flags: EdgeControlFlags = { mode: 'on', rolloutPermille: 0, leaveCeilingPermille: 0 }
+    const h = await harness({ control: () => flags })
+    await h.insert(hostId(1))
+    expect(await h.evacuation.sweep()).toBe(1)
+    expect(errors.mock.calls.map((call) => String(call[0]))).toContain(
+      JSON.stringify({ event: 'orca_relay_edge_row_above_ceiling' })
+    )
+  })
+
   it(`sends at most ${EDGE_EVACUATION_CONCURRENCY} drains at once`, async () => {
     vi.spyOn(console, 'warn').mockImplementation(() => {})
     const h = await harness()

cloud/apps/relay/src/edge-relay-evacuation.ts

@@ -5,7 +5,7 @@
 // retries fences that did not land after a host left (edge-relay-fence.ts).
 import type { RelayDatabase } from './database/database.js'
 import { EdgeAdminUnavailableError, type EdgeAdmin } from './edge-admin-client.js'
-import { inEdgeCohort, text, type EdgeRelayConfig, type Identity } from './edge-relay-cohort.js'
+import { inEdgeCohort, mayBeOnEdge, text, type EdgeRelayConfig, type Identity } from './edge-relay-cohort.js'
 import { forEachBounded, retryPendingFences } from './edge-relay-fence.js'
 
 // A cohort change (director/edge.json, or a new revision) re-arms the sweep (controlChanged), so a
@@ -120,6 +120,11 @@ export class EdgeEvacuation {
         // hosts that the next ask put straight back (a drain every few minutes). A host that asks
         // for a region outside the scope leaves on that ask.
         if (inEdgeCohort(this.config, identity.relayHostId, 'any', identity.userId)) { continue }
+        if (!mayBeOnEdge(this.config, identity.relayHostId, identity.userId)) {
+          // leaveCeilingPermille is set below a rollout this host was moved under: its asks skip
+          // the leave read, so it would be re-placed without its devices. Raise the ceiling.
+          console.error(JSON.stringify({ event: 'orca_relay_edge_row_above_ceiling' }))
+        }
         if (this.drainedAt.has(drainKey(identity))) { continue }
         leaving.push(identity)
         // Resume after the last host taken, so the rest of the page is read next tick.

cloud/apps/relay/src/edge-relay.ts

@@ -37,6 +37,7 @@ export {
   edgeCohortBucket,
   inEdgeCohort,
   mapEpochFloor,
+  mayBeOnEdge,
   type EdgeRelayConfig
 } from './edge-relay-cohort.js'