feat(relay): correct regional placement only when the source is idle (#20105)

* feat(relay): correct regional placement only at an idle source

* test(relay): lock source activity capacity semantics

cloud/dev/scripts/deploy-relay-blue-green.mjs

@@ -10,6 +10,7 @@ export const DIRECTOR_REGIONAL_PLACEMENT_SECRET =
   'orca-cloud-relay-regional-placement-enabled'
 export const DIRECTOR_REGIONAL_PLACEMENT_ENV =
   'ORCA_RELAY_REGIONAL_PLACEMENT_ENABLED'
+export const DIRECTOR_CORRECTION_COHORT_ENV = 'ORCA_RELAY_REGION_CORRECTION_COHORT_PERCENT'
 export const DIRECTOR_REHOME_IDENTITY_ENV =
   'ORCA_RELAY_REHOME_DIRECTOR_SERVICE_ACCOUNT'
 export const DIRECTOR_REHOME_AUDIENCE_ENV = 'ORCA_RELAY_REHOME_AUDIENCE'
@@ -228,6 +229,13 @@ export function directorCellSetAddition(currentValue, desiredValue) {
   return { changed: additions.length > 0, value: JSON.stringify(desired) }
 }
 
+export function correctionCohortPercent(value) {
+  if (!/^(?:[0-9]|[1-9][0-9]|100)$/.test(String(value))) {
+    throw new Error('region correction cohort must be an integer from 0 to 100')
+  }
+  return String(value)
+}
+
 export function directorDeploymentEnvironment(config) {
   const imageDigest = config.image?.match(/@(sha256:[a-f0-9]{64})$/)?.[1]
   if (config.image !== undefined && imageDigest === undefined) {
@@ -238,6 +246,10 @@ export function directorDeploymentEnvironment(config) {
     ORCA_RELAY_ADMISSION_SELECTOR_VERSION: SELECTOR_REVISION_MARKER,
     ...(imageDigest === undefined ? {} : { ORCA_RELAY_IMAGE_DIGEST: imageDigest })
   }
+  if (config['region-correction-cohort-percent'] !== undefined &&
+      config['region-correction-cohort-percent'] !== 'preserve') {
+    environment[DIRECTOR_CORRECTION_COHORT_ENV] = correctionCohortPercent(config['region-correction-cohort-percent'])
+  }
   const serviceAccount = projectServiceAccount(config, 'capacity-service-account')
   const asiaProofServiceAccount = projectServiceAccount(config, 'asia-proof-service-account')
   const rehomeDirectorServiceAccount = projectServiceAccount(
@@ -302,7 +314,8 @@ export function parseArguments(argv) {
       values['rehome-director-service-account'] !== undefined ||
       values['rehome-audience'] !== undefined ||
       values['expected-rehome-generation'] !== undefined ||
-      values['rehome-control-origin'] !== undefined
+      values['rehome-control-origin'] !== undefined ||
+      values['region-correction-cohort-percent'] !== undefined
     ) {
       throw new Error('director configuration arguments require --role director')
     }
@@ -785,6 +798,14 @@ export async function deployDirector(config, tag, overrides = {}) {
     config['prune-revisions'] === 'true' ? CONNECTION_CAPACITY_PROTOCOL : undefined
   const currentEnvironment = revisionEnvironment(servingRevision)
   const deploymentEnvironment = directorDeploymentEnvironment(config)
+  deploymentEnvironment[DIRECTOR_CORRECTION_COHORT_ENV] ??= correctionCohortPercent(
+    currentEnvironment[DIRECTOR_CORRECTION_COHORT_ENV] ?? '0'
+  )
+  if (config['region-correction-cohort-percent'] !== undefined &&
+      config['region-correction-cohort-percent'] !== 'preserve' &&
+      config['expected-rehome-generation'] === undefined) {
+    throw new Error('cohort changes require an exact disabled regional-rehome generation')
+  }
   const mutableEnvironment = {
     ...deploymentEnvironment,
     [DIRECTOR_REGIONAL_PLACEMENT_ENV]: ''

cloud/dev/scripts/deploy-relay-blue-green.test.mjs

@@ -4,6 +4,8 @@ import { test } from 'node:test'
 import { fileURLToPath } from 'node:url'
 import {
   activeRevision,
+  correctionCohortPercent,
+  DIRECTOR_CORRECTION_COHORT_ENV,
   cloudRunTrafficTag,
   DIRECTOR_ADMISSION_ENVIRONMENT,
   DIRECTOR_REGIONAL_PLACEMENT_ENV,
@@ -812,3 +814,43 @@ test('waits for authenticated target readiness without hiding other capacity err
     /forbidden/
   )
 })
+
+
+test('validates bounded correction cohorts and leaves unspecified values to serving inheritance', () => {
+  for (const value of ['0', '1', '100']) assert.equal(correctionCohortPercent(value), value)
+  for (const value of ['-1', '101', '1.5', '', '01', 'true', '1\n']) {
+    assert.throws(() => correctionCohortPercent(value), /integer from 0 to 100/)
+  }
+  assert.equal(directorDeploymentEnvironment({})[DIRECTOR_CORRECTION_COHORT_ENV], undefined)
+  assert.equal(directorDeploymentEnvironment({ 'region-correction-cohort-percent': 'preserve' })[DIRECTOR_CORRECTION_COHORT_ENV], undefined)
+  assert.equal(directorDeploymentEnvironment({ 'region-correction-cohort-percent': '1' })[DIRECTOR_CORRECTION_COHORT_ENV], '1')
+})
+
+test('inherits the cohort on candidate and rollback revisions without resetting an enabled cohort', async () => {
+  const harness = directorHarness()
+  harness.state.revisions.get('relay-00001-old').env[DIRECTOR_CORRECTION_COHORT_ENV] = '3'
+  await deployDirector({}, 'candidate-new', harness.operations)
+  for (const revision of ['relay-00002-new', 'relay-00003-new']) {
+    assert.equal(harness.state.revisions.get(revision).env[DIRECTOR_CORRECTION_COHORT_ENV], '3')
+  }
+})
+
+test('starts an unstamped cohort at zero and rejects a cohort change without disabled-control proof', async () => {
+  const harness = directorHarness()
+  await assert.rejects(deployDirector({ 'region-correction-cohort-percent': '1' },
+    'candidate-new', harness.operations), /exact disabled regional-rehome generation/)
+  assert.equal(harness.state.activeRevision, 'relay-00001-old')
+  assert.equal(harness.state.nextRevision, 2)
+  await deployDirector({}, 'candidate-new', harness.operations)
+  assert.equal(harness.state.revisions.get('relay-00003-new').env[DIRECTOR_CORRECTION_COHORT_ENV], '0')
+})
+
+test('sets a reviewed cohort only behind repeated disabled-control verification', async () => {
+  const harness = directorHarness()
+  let verified = 0
+  const config = { 'region-correction-cohort-percent': '1', 'expected-rehome-generation': '7' }
+  await deployDirector(config, 'candidate-new', { ...harness.operations,
+    assertRegionalRehomeDisabled: async () => { verified++ } })
+  assert.ok(verified >= 2)
+  assert.equal(harness.state.revisions.get('relay-00003-new').env[DIRECTOR_CORRECTION_COHORT_ENV], '1')
+})

cloud/dev/scripts/read-relay-serving-regional-placement-version.mjs

@@ -53,7 +53,7 @@ export function readRelayServingRegionalPlacementVersion(input, dependencies = {
   try {
     service = run(gcloudArguments('services', input))
   } catch (error) {
-    if (error?.code === 'NOT_FOUND') return { version: input.bootstrap_version }
+    if (error?.code === 'NOT_FOUND') return { version: input.bootstrap_version, cohort_percent: '0' }
     throw error
   }
   const serving = (service.status?.traffic ?? []).filter(
@@ -67,12 +67,22 @@ export function readRelayServingRegionalPlacementVersion(input, dependencies = {
     throw new Error('Relay director must have exactly one revision serving 100% traffic')
   }
   const revision = run(gcloudArguments('revisions', input, serving[0].revisionName))
+  const cohortSettings = (revision.spec?.containers ?? []).flatMap((container) =>
+    (container.env ?? []).filter((environment) =>
+      environment.name === 'ORCA_RELAY_REGION_CORRECTION_COHORT_PERCENT')
+  )
+  if (cohortSettings.length > 1 || (cohortSettings.length === 1 &&
+      (typeof cohortSettings[0].value !== 'string' ||
+       !/^(?:[0-9]|[1-9][0-9]|100)$/.test(cohortSettings[0].value)))) {
+    throw new Error('serving region correction cohort is invalid')
+  }
+  const cohort_percent = cohortSettings[0]?.value ?? '0'
   const references = (revision.spec?.containers ?? []).flatMap((container) =>
     (container.env ?? []).filter(
       (environment) => environment.name === 'ORCA_RELAY_REGIONAL_PLACEMENT_ENABLED'
     )
   )
-  if (references.length === 0) return { version: input.bootstrap_version }
+  if (references.length === 0) return { version: input.bootstrap_version, cohort_percent }
   const reference = normalizeSecretReference(references[0])
   if (
     references.length !== 1 ||
@@ -81,7 +91,7 @@ export function readRelayServingRegionalPlacementVersion(input, dependencies = {
   ) {
     throw new Error('serving regional placement secret reference is invalid')
   }
-  return { version: reference.version }
+  return { version: reference.version, cohort_percent }
 }
 
 // Why: the v2 API reports `valueSource.secretKeyRef.{secret,version}`, but

cloud/dev/scripts/read-relay-serving-regional-placement-version.test.mjs

@@ -67,7 +67,7 @@ test('reads the exact version from the sole traffic-serving revision', () => {
     }
   })
 
-  assert.deepEqual(result, { version: '11' })
+  assert.deepEqual(result, { version: '11', cohort_percent: '0' })
   assert.equal(calls[1][3], 'relay-serving')
 })
 
@@ -78,7 +78,7 @@ test('reads the gcloud v1 secret reference shape by bare id and by full resource
   ]) {
     assert.deepEqual(readRelayServingRegionalPlacementVersion(input, {
       run: (args) => args[1] === 'services' ? serving() : v1Revision(name, '1')
-    }), { version: '1' })
+    }), { version: '1', cohort_percent: '0' })
   }
 })
 
@@ -100,12 +100,12 @@ test('falls back only when the service or setting is absent', () => {
   notFound.code = 'NOT_FOUND'
   assert.deepEqual(readRelayServingRegionalPlacementVersion(input, {
     run: () => { throw notFound }
-  }), { version: '7' })
+  }), { version: '7', cohort_percent: '0' })
   assert.deepEqual(readRelayServingRegionalPlacementVersion(input, {
     run: (args) => args[1] === 'services'
       ? { status: { traffic: [{ revisionName: 'relay-serving', percent: 100 }] } }
       : { spec: { containers: [{ env: [] }] } }
-  }), { version: '7' })
+  }), { version: '7', cohort_percent: '0' })
 })
 
 test('classifies real absent-service stderr without weakening revision failures', () => {
@@ -136,3 +136,31 @@ test('rejects ambiguous traffic, malformed references, and read failures', () =>
     run: () => { throw denied }
   }), denied)
 })
+
+
+test('preserves the serving cohort including explicit disable across later Terraform plans', () => {
+  for (const value of ['0', '1', '17', '100']) {
+    const servingRevision = revision()
+    servingRevision.spec.containers[0].env.push({ name: 'ORCA_RELAY_REGION_CORRECTION_COHORT_PERCENT', value })
+    assert.deepEqual(readRelayServingRegionalPlacementVersion(input, {
+      run: (args) => args[1] === 'services' ? serving() : servingRevision
+    }), { version: '11', cohort_percent: value })
+  }
+})
+
+test('fails closed on malformed, secret-backed or duplicate cohorts rather than resetting them', () => {
+  const name = 'ORCA_RELAY_REGION_CORRECTION_COHORT_PERCENT'
+  const cases = [
+    [{ name, value: '101' }], [{ name, value: '-1' }], [{ name, value: '1.5' }],
+    [{ name, value: '' }], [{ name, value: '01' }], [{ name, value: 1 }],
+    [{ name, valueFrom: { secretKeyRef: { name: 'unexpected', key: '1' } } }],
+    [{ name, value: '1' }, { name, value: '2' }]
+  ]
+  for (const settings of cases) {
+    const servingRevision = revision()
+    servingRevision.spec.containers[0].env.push(...settings)
+    assert.throws(() => readRelayServingRegionalPlacementVersion(input, {
+      run: (args) => args[1] === 'services' ? serving() : servingRevision
+    }), /cohort is invalid/)
+  }
+})