fix(relay): let the staging admission workflow register and promote c5-c7 (#27491)

The staging admission config hard-wired staging-gce-c4, so new staging cells
had no reviewed path into the admission selector. Add c5-c7 as a second
staging wave at their US shape; production cell shapes are unchanged.

cloud/dev/scripts/operate-relay-asia-admission.mjs

@@ -12,9 +12,9 @@ const SHAPES = {
   staging: {
     directorOrigin: 'https://relay-staging.onorca.dev',
     domain: 'relay-staging.onorca.dev',
-    allCells: ['staging-gce-c4'],
-    registrationWaves: [['staging-gce-c4']],
-    promotionWaves: [['staging-gce-c4']]
+    allCells: ['staging-gce-c4', 'staging-gce-c5', 'staging-gce-c6', 'staging-gce-c7'],
+    registrationWaves: [['staging-gce-c4'], ['staging-gce-c5', 'staging-gce-c6', 'staging-gce-c7']],
+    promotionWaves: [['staging-gce-c4'], ['staging-gce-c5', 'staging-gce-c6', 'staging-gce-c7']]
   },
   production: {
     directorOrigin: 'https://relay.onorca.dev',
@@ -45,12 +45,26 @@ const SHAPES = {
 }
 
 const PRODUCTION_CANARY_CELL = 'production-gce-c27'
+const STAGING_PROOF_CELL = 'staging-gce-c4'
 const ASIA_REGION = 'asia-east2'
-// The US cells at the Asia shape live in the root region; every other cell here is Asia.
-const CELL_REGIONS = { 'production-gce-c32': 'us-central1', 'production-gce-c33': 'us-central1' }
+const ASIA_CELL_SHAPE = { region: ASIA_REGION, capacityRequests: 6_000, connectionHardCap: 3_000 }
+// The US cells at the Asia shape live in the root region; every other cell here is Asia, except
+// the staging dead-cell guard cells, which keep c2's ordinary US shape (staging.tfvars).
+const STAGING_US_CELL_SHAPE = { region: 'us-central1', capacityRequests: 4_000, connectionHardCap: 600 }
+const CELL_SHAPES = {
+  'production-gce-c32': { ...ASIA_CELL_SHAPE, region: 'us-central1' },
+  'production-gce-c33': { ...ASIA_CELL_SHAPE, region: 'us-central1' },
+  'staging-gce-c5': STAGING_US_CELL_SHAPE,
+  'staging-gce-c6': STAGING_US_CELL_SHAPE,
+  'staging-gce-c7': STAGING_US_CELL_SHAPE
+}
+
+function cellShape(cellId) {
+  return CELL_SHAPES[cellId] ?? ASIA_CELL_SHAPE
+}
 
 function cellRegion(cellId) {
-  return CELL_REGIONS[cellId] ?? ASIA_REGION
+  return cellShape(cellId).region
 }
 
 export function parseRelayAsiaAdmissionArguments(argv) {
@@ -156,9 +170,9 @@ async function verifyRuntime(fetchImpl, post, shape, cellId, imageDigest, requir
     runtime.region !== cellRegion(cellId) ||
     runtime.imageDigest !== imageDigest ||
     runtime.draining !== false ||
-    runtime.connectionCapacity?.hardCap !== 3_000 ||
+    runtime.connectionCapacity?.hardCap !== cellShape(cellId).connectionHardCap ||
     runtime.connectionCapacity?.unobservedBound !== 60
-  ) throw new Error(`${cellId} runtime does not match the reviewed Asia shape`)
+  ) throw new Error(`${cellId} runtime does not match its reviewed shape`)
   if (requireDirector) {
     const result = await post(`${shape.directorOrigin}/v1/admin/cell-status`, { v: 1, cellId })
     if (
@@ -272,11 +286,15 @@ export async function operateRelayAsiaAdmission(config, dependencies = {}) {
   const shape = SHAPES[config.environment]
   const fetchImpl = dependencies.fetch ?? fetch
   const post = dependencies.post ?? defaultPost(fetchImpl, config.token)
+  // Only the C4 proof keeps the server-enforced C4-only route (the Asia proof identity can call
+  // nothing else); other staging cells go through the general route, as production does.
+  const stagingProof = config.environment === 'staging' &&
+    config.cells.length === 1 && config.cells[0] === STAGING_PROOF_CELL
   const selectorPost = (path, body) => {
-    if (config.environment !== 'staging' || path !== '/v1/admin/admission-selector/apply') {
+    if (!stagingProof || path !== '/v1/admin/admission-selector/apply') {
       return post(`${shape.directorOrigin}${path}`, body)
     }
-    const state = selectorCellState({ membership: body.membership }, 'staging-gce-c4')
+    const state = selectorCellState({ membership: body.membership }, STAGING_PROOF_CELL)
     if (!['general', 'migration-only'].includes(state)) {
       throw new Error('staging proof can only transition C4 between reviewed states')
     }
@@ -418,8 +436,8 @@ export async function operateRelayAsiaAdmission(config, dependencies = {}) {
           cellId,
           cellUrl: cellOrigin(shape, cellId),
           region: cellRegion(cellId),
-          capacityRequests: 6_000,
-          connectionHardCap: 3_000,
+          capacityRequests: cellShape(cellId).capacityRequests,
+          connectionHardCap: cellShape(cellId).connectionHardCap,
           connectionUnobservedBound: 60
         }))
       },

cloud/dev/scripts/operate-relay-asia-admission.test.mjs

@@ -15,7 +15,7 @@ const digest = `sha256:${'a'.repeat(64)}`
 const membershipDigest = (membership) =>
   createHash('sha256').update(JSON.stringify(membership)).digest('hex')
 
-function harness(initialSelector, runtimeDigests = {}, runtimeRegions = {}) {
+function harness(initialSelector, runtimeDigests = {}, runtimeRegions = {}, runtimeHardCaps = {}) {
   const initialMembership = structuredClone(initialSelector.membership)
   let selector = structuredClone(initialSelector)
   const intents = new Map()
@@ -26,20 +26,21 @@ function harness(initialSelector, runtimeDigests = {}, runtimeRegions = {}) {
     const parsed = new URL(url)
     requests.push({ path: parsed.pathname, body })
     if (parsed.pathname === '/v1/admin/runtime-status') {
-      const cell = parsed.hostname.split('.')[0]
+      const [cell, domain] = parsed.hostname.split(/\.(.*)/)
+      const cellId = `${domain.startsWith('relay-staging.') ? 'staging' : 'production'}-gce-${cell}`
       return {
-        cellId: `production-gce-${cell}`,
+        cellId,
         cellUrl: parsed.origin,
-        region: runtimeRegions[`production-gce-${cell}`] ?? 'asia-east2',
-        imageDigest: runtimeDigests[`production-gce-${cell}`] ?? digest,
+        region: runtimeRegions[cellId] ?? 'asia-east2',
+        imageDigest: runtimeDigests[cellId] ?? digest,
         draining: false,
-        connectionCapacity: { hardCap: 3_000, unobservedBound: 60 }
+        connectionCapacity: { hardCap: runtimeHardCaps[cellId] ?? 3_000, unobservedBound: 60 }
       }
     }
     if (parsed.pathname === '/v1/admin/cell-status') {
       return {
         status: {
-          cellUrl: `https://${body.cellId.split('-').at(-1)}.relay.onorca.dev`,
+          cellUrl: `https://${body.cellId.split('-').at(-1)}.${parsed.hostname}`,
           runtime: { heartbeatFresh: true, ready: true }
         }
       }
@@ -857,3 +858,59 @@ test('runs each later cell\'s own canary with load aimed at that cell\'s region'
   assert.match(admissionWorkflow, /--preferred-region "\$\{CANARY_REGION\}"/)
   assert.match(admissionWorkflow, /CANARY_REGION: \$\{\{ steps\.inputs\.outputs\.canary_region \}\}/)
 })
+
+const stagingDeadCellWave = ['staging-gce-c5', 'staging-gce-c6', 'staging-gce-c7']
+const stagingUsShape = Object.fromEntries(stagingDeadCellWave.map((cellId) => [cellId, 'us-central1']))
+const stagingUsCaps = Object.fromEntries(stagingDeadCellWave.map((cellId) => [cellId, 600]))
+
+test('registers the staging dead-cell wave at its own US shape', async () => {
+  const subject = harness({
+    generation: 43,
+    membership: { existingOnly: ['staging-gce-c3'], migrationOnly: [], general: ['staging-gce-c1', 'staging-gce-c2', 'staging-gce-c4'] }
+  }, {}, stagingUsShape, stagingUsCaps)
+  const result = await operateRelayAsiaAdmission({
+    environment: 'staging', mode: 'register', cells: stagingDeadCellWave,
+    expectedGeneration: 43, imageDigest: digest, attemptId: 'staging_c5_c7_register', token: 'not-logged'
+  }, subject)
+  const request = subject.requests.find(({ path }) => path.endsWith('/add-migration-cells'))
+  assert.deepEqual(request.body.cells.map((cell) => cell.cellUrl), [
+    'https://c5.relay-staging.onorca.dev', 'https://c6.relay-staging.onorca.dev', 'https://c7.relay-staging.onorca.dev'
+  ])
+  assert.ok(request.body.cells.every((cell) =>
+    cell.region === 'us-central1' && cell.capacityRequests === 4_000 &&
+    cell.connectionHardCap === 600 && cell.connectionUnobservedBound === 60
+  ))
+  assert.deepEqual(new Set(Object.values(result.states)), new Set(['migration-only']))
+})
+
+test('promotes the staging dead-cell wave through the general route, not the C4 proof route', async () => {
+  const subject = harness({
+    generation: 44,
+    membership: {
+      existingOnly: ['staging-gce-c3'],
+      migrationOnly: [...stagingDeadCellWave],
+      general: ['staging-gce-c1', 'staging-gce-c2', 'staging-gce-c4']
+    }
+  }, {}, stagingUsShape, stagingUsCaps)
+  const result = await operateRelayAsiaAdmission({
+    environment: 'staging', mode: 'promote', cells: stagingDeadCellWave,
+    expectedGeneration: 44, imageDigest: digest, attemptId: 'staging_c5_c7_promote', token: 'not-logged'
+  }, subject)
+  assert.ok(subject.requests.some(({ path }) => path === '/v1/admin/admission-selector/apply'))
+  assert.ok(!subject.requests.some(({ path }) => path.endsWith('/apply-staging-asia-proof')))
+  assert.deepEqual(new Set(Object.values(result.states)), new Set(['general']))
+  assert.deepEqual(subject.selector().membership.general, [
+    'staging-gce-c1', 'staging-gce-c2', 'staging-gce-c4', ...stagingDeadCellWave
+  ])
+})
+
+test('refuses a staging wave that was not reviewed', () => {
+  assert.throws(() => parseRelayAsiaAdmissionArguments([
+    '--environment', 'staging', '--mode', 'promote', '--cell-ids', 'staging-gce-c4,staging-gce-c5',
+    '--expected-generation', '44', '--attempt-id', 'staging_mixed_wave', '--image-digest', digest
+  ]), /reviewed admission wave/)
+  assert.throws(() => parseRelayAsiaAdmissionArguments([
+    '--environment', 'staging', '--mode', 'register', '--cell-ids', 'staging-gce-c8',
+    '--expected-generation', '44', '--attempt-id', 'staging_unknown', '--image-digest', digest
+  ]), /cell-ids are invalid/)
+})