diff --git a/cloud/apps/relay/src/assignment/assignment-store.ts b/cloud/apps/relay/src/assignment/assignment-store.ts index 90c440fd72c..19ca602855d 100644 --- a/cloud/apps/relay/src/assignment/assignment-store.ts +++ b/cloud/apps/relay/src/assignment/assignment-store.ts @@ -75,7 +75,6 @@ import { type CellConnectionCap } from '../cell-connection-cap.js' import { AssignmentIdentityQueue } from './assignment-identity-queue.js' -import { recordRegionPreference } from './region-preference.js' import { CONTROL_RENEWAL_BATCH_SQL, CONTROL_RENEWAL_STATEMENT_OUTCOMES, @@ -1145,7 +1144,7 @@ export class RelayAssignmentStore { ) if (!existing) return null const activityLeases = await this.lockAssignmentActivities(transaction, identity, true) - await recordRegionPreference(transaction, identity, preferredRegion, now) + await this.recordRegionPreference(transaction, identity, preferredRegion, now) // A live reserve cell's hosts hold no activity here, so idle is not evidence they left. if ( !(await this.liveReserveCell(text(existing, 'cell_id'), transaction, now)) && @@ -1371,7 +1370,7 @@ export class RelayAssignmentStore { throw new AssignmentInventoryScopeChanged() } const activityLeases = await this.lockAssignmentActivities(transaction, identity, true) - await recordRegionPreference(transaction, identity, preferredRegion, now) + await this.recordRegionPreference(transaction, identity, preferredRegion, now) let forcedDeadReassignment = false let connectionHeadroomReassignment = false let strandedReassignment = false @@ -8280,6 +8279,32 @@ export class RelayAssignmentStore { }) } + private async recordRegionPreference( + database: RelayDatabase, + identity: AssignmentIdentity, + preferredRegion: RelayRegion | undefined, + observedAt: number + ): Promise { + if (!preferredRegion) return + await database.query( + `INSERT INTO relay_assignment_region_preferences + (user_id, relay_host_id, preferred_region, observed_at) + VALUES (?, ?, ?, ?) + ON CONFLICT (user_id, relay_host_id) DO UPDATE SET + preferred_region = CASE + WHEN excluded.observed_at >= relay_assignment_region_preferences.observed_at + THEN excluded.preferred_region + ELSE relay_assignment_region_preferences.preferred_region + END, + observed_at = CASE + WHEN excluded.observed_at >= relay_assignment_region_preferences.observed_at + THEN excluded.observed_at + ELSE relay_assignment_region_preferences.observed_at + END`, + [identity.userId, identity.relayHostId, preferredRegion, observedAt] + ) + } + private async cellRegion(database: RelayDatabase, cellId: string): Promise { const row = ( await database.query(`SELECT region FROM relay_cell_regions WHERE cell_id = ?`, [cellId]) diff --git a/cloud/apps/relay/src/assignment/region-preference.ts b/cloud/apps/relay/src/assignment/region-preference.ts deleted file mode 100644 index 098cecb68a1..00000000000 --- a/cloud/apps/relay/src/assignment/region-preference.ts +++ /dev/null @@ -1,30 +0,0 @@ -import type { RelayRegion } from '@orca-cloud/relay-contract/relay-regions' -import type { RelayDatabase } from '../database.js' - -// The region a host last asked for (relay_assignment_region_preferences); an older observation -// never overwrites a newer one. Written by the cell path and the relay edge's moves. -export async function recordRegionPreference( - database: RelayDatabase, - identity: { userId: string; relayHostId: string }, - preferredRegion: RelayRegion | undefined, - observedAt: number -): Promise { - if (!preferredRegion) return - await database.query( - `INSERT INTO relay_assignment_region_preferences - (user_id, relay_host_id, preferred_region, observed_at) - VALUES (?, ?, ?, ?) - ON CONFLICT (user_id, relay_host_id) DO UPDATE SET - preferred_region = CASE - WHEN excluded.observed_at >= relay_assignment_region_preferences.observed_at - THEN excluded.preferred_region - ELSE relay_assignment_region_preferences.preferred_region - END, - observed_at = CASE - WHEN excluded.observed_at >= relay_assignment_region_preferences.observed_at - THEN excluded.observed_at - ELSE relay_assignment_region_preferences.observed_at - END`, - [identity.userId, identity.relayHostId, preferredRegion, observedAt] - ) -} diff --git a/cloud/apps/relay/src/database.ts b/cloud/apps/relay/src/database.ts index 5349992092f..e3e18f0e711 100644 --- a/cloud/apps/relay/src/database.ts +++ b/cloud/apps/relay/src/database.ts @@ -274,6 +274,18 @@ CREATE TABLE IF NOT EXISTS relay_edge_fence_pending ( PRIMARY KEY (user_id, relay_host_id) ); +-- Region each relay edge host was last moved or kept there for (edge-relay.ts), read by the +-- evacuation sweep for hosts whose assignment row names the edge. Written only by edge moves, and +-- deleted when the host is placed off the edge. No CHECK on region: a new region must not need a +-- constraint swap before edge moves can record it (readers validate it). +CREATE TABLE IF NOT EXISTS relay_edge_host_regions ( + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + region TEXT NOT NULL, + recorded_at BIGINT NOT NULL, + PRIMARY KEY (user_id, relay_host_id) +); + CREATE TABLE IF NOT EXISTS relay_region_decisions ( user_id TEXT NOT NULL, relay_host_id TEXT NOT NULL, generation BIGINT NOT NULL, expires_at BIGINT NOT NULL, diff --git a/cloud/apps/relay/src/edge-relay-control.ts b/cloud/apps/relay/src/edge-relay-control.ts index 893c065083c..83eb52f05a8 100644 --- a/cloud/apps/relay/src/edge-relay-control.ts +++ b/cloud/apps/relay/src/edge-relay-control.ts @@ -52,7 +52,8 @@ const EdgeControlFlagsSchema = z.object({ leaveCeilingPermille: permille.optional(), regionRolloutPermille: z.partialRecord(regionKey, permille).optional(), regionExcludeCountries: z - .partialRecord(regionKey, z.array(z.string().regex(/^[A-Z]{2}$/)).max(250)) + // Upper-cased first, so a lower-case code is not a typo that voids the whole file. + .partialRecord(regionKey, z.array(z.string().trim().toUpperCase().regex(/^[A-Z]{2}$/)).max(250)) .optional(), regionExcludeUnknownCountry: z.partialRecord(regionKey, z.boolean()).optional() }) @@ -141,8 +142,8 @@ export function edgeRolloutPermille(rolloutPercent: number, flags: EdgeControlFl return flags.rolloutPermille ?? rolloutPercent * 10 } -// Any regional share set: edge moves record their region and the sweep judges recorded regions. -// None: neither, so a dark rollout adds no database writes. +// Any regional share set: the evacuation sweep judges each edge host by its recorded region +// (relay_edge_host_regions). None: it judges every host 'any' with today's query. export function edgeRegionalRollout(flags: EdgeControlFlags): boolean { return Object.keys(flags.regionRolloutPermille ?? {}).length > 0 } diff --git a/cloud/apps/relay/src/edge-relay-evacuation.ts b/cloud/apps/relay/src/edge-relay-evacuation.ts index 0d94a9e20b8..6ce6f7e2263 100644 --- a/cloud/apps/relay/src/edge-relay-evacuation.ts +++ b/cloud/apps/relay/src/edge-relay-evacuation.ts @@ -118,21 +118,23 @@ export class EdgeEvacuation { private async leavingHosts(): Promise { const leaving: Identity[] = [] this.pass ??= this.newPass() + const regional = edgeRegionalRollout(this.config.control?.() ?? EDGE_CONTROL_DEFAULTS) for (let pages = 0; pages < 2 && this.pass !== undefined; pages += 1) { const pass: Pass = this.pass - const rows = await this.readPage(pass) + const rows = await this.readPage(pass, regional) let filled = false for (const row of rows) { const identity = { userId: text(row, 'user_id'), relayHostId: text(row, 'relay_host_id') } pass.cursor = identity // Region scope is the assign path's call: it decides with the region the desktop asks for. - // Without regional shares the sweep judges 'any': judging a lagging cell-path preference - // drained hosts that the next ask put straight back. With them, edge asks record their - // region (moveToEdge) and a host with a recorded region is judged by that region's - // effective share, so a regional abort or shrink drains just that region; a host with none - // is judged 'any'. Country exclusion needs the ask's address, so an excluded host leaves - // on its next ask. - if (inEdgeCohort(this.config, identity.relayHostId, this.judgedRegion(row), identity.userId)) { continue } + // Without regional shares the sweep judges 'any' (the cell path's stored preference can + // lag: judging it drained hosts that the next ask put straight back). With them, a host + // is judged by the effective share of the region its last edge ask recorded + // (relay_edge_host_regions), so a regional abort or shrink drains just that region; a + // host with no record is judged 'any'. Country exclusion needs the ask's address, so an + // excluded host leaves on its next ask. + const region = regional ? recordedRegion(row) : 'any' + if (inEdgeCohort(this.config, identity.relayHostId, region, 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. @@ -158,21 +160,19 @@ export class EdgeEvacuation { return leaving } - private judgedRegion(row: SqlRow): RelayRegion | 'any' { - if (!edgeRegionalRollout(this.config.control?.() ?? EDGE_CONTROL_DEFAULTS)) { return 'any' } - const recorded = RelayRegionSchema.safeParse(row.preferred_region) - return recorded.success ? recorded.data : 'any' - } - - // Keys after the cursor; once wrapped, only up to the pass's origin. With the region the host's - // last ask recorded (primary key join; none: NULL). - private async readPage(pass: Pass) { + // Keys after the cursor; once wrapped, only up to the pass's origin. `regional`: with the region + // the host's last edge ask recorded (primary key join; none: NULL). + private async readPage(pass: Pass, regional: boolean) { const { cursor, origin, wrapped } = pass return await this.database.query( - `SELECT a.user_id, a.relay_host_id, p.preferred_region - FROM relay_assignments a - LEFT JOIN relay_assignment_region_preferences p - ON p.user_id = a.user_id AND p.relay_host_id = a.relay_host_id + `SELECT a.user_id, a.relay_host_id${regional ? ', r.region' : ''} + FROM relay_assignments a${ + regional + ? ` + LEFT JOIN relay_edge_host_regions r + ON r.user_id = a.user_id AND r.relay_host_id = a.relay_host_id` + : '' + } WHERE a.cell_id = ? AND (a.user_id > ? OR (a.user_id = ? AND a.relay_host_id > ?)) ${wrapped ? 'AND (a.user_id < ? OR (a.user_id = ? AND a.relay_host_id <= ?))' : ''} ORDER BY a.user_id, a.relay_host_id @@ -223,6 +223,11 @@ export class EdgeEvacuation { } } +function recordedRegion(row: SqlRow): RelayRegion | 'any' { + const recorded = RelayRegionSchema.safeParse(row.region) + return recorded.success ? recorded.data : 'any' +} + function drainKey(identity: Identity): string { return `${identity.userId}\n${identity.relayHostId}` } diff --git a/cloud/apps/relay/src/edge-relay-postgres.test.ts b/cloud/apps/relay/src/edge-relay-postgres.test.ts index 146d146f4f0..8c1b68f1504 100644 --- a/cloud/apps/relay/src/edge-relay-postgres.test.ts +++ b/cloud/apps/relay/src/edge-relay-postgres.test.ts @@ -39,7 +39,7 @@ describePostgres('relay edge on PostgreSQL', () => { await database.query(`DELETE FROM relay_assignments WHERE user_id = ?`, ['edge-pg-user']) await database.query(`DELETE FROM relay_devices WHERE user_id = ?`, ['edge-pg-user']) await database.query(`DELETE FROM relay_edge_fence_pending WHERE user_id = ?`, ['edge-pg-user']) - await database.query(`DELETE FROM relay_assignment_region_preferences WHERE user_id = ?`, ['edge-pg-user']) + await database.query(`DELETE FROM relay_edge_host_regions WHERE user_id = ?`, ['edge-pg-user']) // Leave the shared selector as other files expect it: without this file's cells. await database.transaction(async (transaction) => { const selector = ( @@ -295,7 +295,8 @@ describePostgres('relay edge on PostgreSQL', () => { return { outcome: 'idle' as const } } } - const edge = new EdgeAssignments(database, config, Date.now, admin) + let now = Date.now() + const edge = new EdgeAssignments(database, config, () => now, admin) await edge.ensureEdgeCell() const [asia, us, failed] = ['edgepgregion0001', 'edgepgregion0002', 'edgepgregion0003'].map((relayHostId) => ({ userId: 'edge-pg-user', @@ -304,16 +305,31 @@ describePostgres('relay edge on PostgreSQL', () => { await edge.moveToEdge(asia!, 0, 'asia-east2') await edge.moveToEdge(us!, 0, 'us-central1') await expect(edge.moveToEdge(failed!, 0, 'asia-east2')).rejects.toBeInstanceOf(EdgeMoveDeferredError) - // A kept host re-records (its row was already on the edge: no push). + const recordedAt = async (host: { relayHostId: string }) => + Number( + ( + await database.query(`SELECT recorded_at FROM relay_edge_host_regions WHERE user_id = ? AND relay_host_id = ?`, [ + 'edge-pg-user', + host.relayHostId + ]) + )[0]?.recorded_at + ) + // A kept host with the same region writes nothing; a changed region is rewritten. + const first = await recordedAt(us!) + now += 5_000 + await edge.moveToEdge(us!, 0, 'us-central1') + expect(await recordedAt(us!)).toBe(first) + await edge.moveToEdge(us!, 0, 'asia-east2') + expect(await recordedAt(us!)).toBe(now) await edge.moveToEdge(us!, 0, 'us-central1') const recorded = await database.query( - `SELECT relay_host_id, preferred_region FROM relay_assignment_region_preferences + `SELECT relay_host_id, region FROM relay_edge_host_regions WHERE user_id = ? AND relay_host_id LIKE 'edgepgregion%' ORDER BY relay_host_id`, ['edge-pg-user'] ) expect(recorded).toEqual([ - { relay_host_id: asia!.relayHostId, preferred_region: 'asia-east2' }, - { relay_host_id: us!.relayHostId, preferred_region: 'us-central1' } + { relay_host_id: asia!.relayHostId, region: 'asia-east2' }, + { relay_host_id: us!.relayHostId, region: 'us-central1' } ]) // Abort Asia: only its recorded host drains; us-central1 keeps the global share (100%). flags.regionRolloutPermille = { 'asia-east2': 0 } diff --git a/cloud/apps/relay/src/edge-relay-region.test.ts b/cloud/apps/relay/src/edge-relay-region.test.ts index ee4feec2778..d0c4a3f861b 100644 --- a/cloud/apps/relay/src/edge-relay-region.test.ts +++ b/cloud/apps/relay/src/edge-relay-region.test.ts @@ -18,7 +18,12 @@ import { } from './edge-relay.js' import { EdgeEvacuation } from './edge-relay-evacuation.js' import { EDGE_REGION_COUNTRY_EVENT, EdgeRegionCountryTally } from './edge-region-country-log.js' -import type { EdgeAdmin, EdgeDrainResult } from './edge-admin-client.js' +import { + EdgeAdminUnavailableError, + type EdgeAdmin, + type EdgeAssignBody, + type EdgeDrainResult +} from './edge-admin-client.js' function config(flags: EdgeControlFlags, overrides: Partial = {}): EdgeRelayConfig { return { @@ -214,6 +219,13 @@ describe('edge control regional flags', () => { expect(parseEdgeControl({ v: 1, flags })).toEqual({ flags, ignoredKeys: [] }) }) + it('upper-cases country codes rather than voiding the file', () => { + expect(parseEdgeControl({ v: 1, flags: { mode: 'on', regionExcludeCountries: { 'asia-east2': ['cn', ' Hk '] } } })).toEqual({ + flags: { mode: 'on', regionExcludeCountries: { 'asia-east2': ['CN', 'HK'] } }, + ignoredKeys: [] + }) + }) + it('voids a file whose leave ceiling is below any regional share', () => { expect( parseEdgeControl({ v: 1, flags: { mode: 'on', leaveCeilingPermille: 50, regionRolloutPermille: { 'asia-east2': 51 } } }) @@ -231,7 +243,8 @@ describe('edge control regional flags', () => { { regionRolloutPermille: { 'asia-east2': 1001 } }, { regionRolloutPermille: { 'asia-east2': 1.5 } }, { regionRolloutPermille: ['asia-east2'] }, - { regionExcludeCountries: { 'asia-east2': ['cn'] } }, + { regionExcludeCountries: { 'asia-east2': ['C'] } }, + { regionExcludeCountries: { 'asia-east2': ['CHN'] } }, { regionExcludeCountries: { 'asia-east2': 'CN' } }, { regionExcludeCountries: { 'mars-north1': ['CN'] } }, { regionExcludeUnknownCountry: { 'asia-east2': 'yes' } } @@ -244,69 +257,89 @@ describe('edge control regional flags', () => { describe('edge move records the region', () => { const databases: RelayDatabase[] = [] afterEach(async () => { + vi.restoreAllMocks() for (const database of databases.splice(0)) await database.close() }) const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } - const regional: EdgeControlFlags = { mode: 'on', regionRolloutPermille: { 'asia-east2': 0 } } async function recorded(database: RelayDatabase) { return ( await database.query( - `SELECT preferred_region, observed_at FROM relay_assignment_region_preferences - WHERE user_id = ? AND relay_host_id = ?`, + `SELECT region, recorded_at FROM relay_edge_host_regions WHERE user_id = ? AND relay_host_id = ?`, [identity.userId, identity.relayHostId] ) - ).map((row) => ({ region: row.preferred_region, at: Number(row.observed_at) })) + ).map((row) => ({ region: row.region, at: Number(row.recorded_at) })) } - it('on a first move, a kept host and a move off a cell; no region asked is the default region', async () => { + async function setup(admin?: EdgeAdmin) { const database = await openInMemoryRelayDatabase() databases.push(database) - let now = 1_000 - const edge = new EdgeAssignments(database, config(regional), () => now) + const clock = { now: 1_000 } + // Dark: no regional flag at all. The record does not depend on flags. + const edge = new EdgeAssignments(database, config({ mode: 'on' }), () => clock.now, admin) await edge.ensureEdgeCell() + return { database, clock, edge } + } + + it('on a first move, and on every move off a cell; no region asked is the default region', async () => { + const { database, clock, edge } = await setup() await edge.moveToEdge(identity, 0, 'asia-east2') expect(await recorded(database)).toEqual([{ region: 'asia-east2', at: 1_000 }]) - now = 2_000 - await edge.moveToEdge(identity, 0) - expect(await recorded(database)).toEqual([{ region: RELAY_DEFAULT_REGION, at: 2_000 }]) - // Off to a cell row, then moved back on. + // Off to a cell row, then moved back on with the same region: rewritten. await database.query(`UPDATE relay_assignments SET cell_id = 'us-c1' WHERE user_id = ?`, [identity.userId]) - now = 3_000 + clock.now = 2_000 await edge.moveToEdge(identity, 0, 'asia-east2') - expect(await recorded(database)).toEqual([{ region: 'asia-east2', at: 3_000 }]) + expect(await recorded(database)).toEqual([{ region: 'asia-east2', at: 2_000 }]) + await database.query(`UPDATE relay_assignments SET cell_id = 'us-c1' WHERE user_id = ?`, [identity.userId]) + clock.now = 3_000 + await edge.moveToEdge(identity, 0) + expect(await recorded(database)).toEqual([{ region: RELAY_DEFAULT_REGION, at: 3_000 }]) }) - it('writes nothing without a regional share (dark), in every branch', async () => { - const database = await openInMemoryRelayDatabase() - databases.push(database) - const flags: EdgeControlFlags = { mode: 'on', rolloutPermille: 1000, regionRolloutPermille: {} } - const edge = new EdgeAssignments(database, config(flags), () => 1_000) - await edge.ensureEdgeCell() + it('on a kept host only when its region changed or was never recorded', async () => { + const { database, clock, edge } = await setup() + await edge.moveToEdge(identity, 0, 'asia-east2') + clock.now = 2_000 + await edge.moveToEdge(identity, 0, 'asia-east2') + expect(await recorded(database)).toEqual([{ region: 'asia-east2', at: 1_000 }]) + clock.now = 3_000 await edge.moveToEdge(identity, 0, 'us-central1') + expect(await recorded(database)).toEqual([{ region: 'us-central1', at: 3_000 }]) + // A host moved before records existed is recorded on its next kept ask. + await database.query(`DELETE FROM relay_edge_host_regions`) + clock.now = 4_000 await edge.moveToEdge(identity, 0, 'us-central1') - await database.query(`UPDATE relay_assignments SET cell_id = 'us-c1' WHERE user_id = ?`, [identity.userId]) - await edge.moveToEdge(identity, 0) - expect(await recorded(database)).toEqual([]) - // An empty map counts as absent; the first regional key turns recording on. - flags.regionRolloutPermille = { 'asia-east2': 0 } - await edge.moveToEdge(identity, 0, 'us-central1') - expect(await recorded(database)).toEqual([{ region: 'us-central1', at: 1_000 }]) + expect(await recorded(database)).toEqual([{ region: 'us-central1', at: 4_000 }]) + }) + + it('never writes the cell path region preferences', async () => { + const { database, edge } = await setup() + await edge.moveToEdge(identity, 0, 'asia-east2') + expect(await database.query(`SELECT 1 FROM relay_assignment_region_preferences`)).toEqual([]) }) it('records nothing for a move that does not commit', async () => { - const database = await openInMemoryRelayDatabase() - databases.push(database) const admin = { assign: vi.fn(async () => { - throw new (await import('./edge-admin-client.js')).EdgeAdminUnavailableError('status_503') + throw new EdgeAdminUnavailableError('status_503') }) } as unknown as EdgeAdmin - const edge = new EdgeAssignments(database, config(regional), () => 1_000, admin) - await edge.ensureEdgeCell() + const { database, edge } = await setup(admin) await expect(edge.moveToEdge(identity, 0, 'asia-east2')).rejects.toBeInstanceOf(EdgeMoveDeferredError) expect(await recorded(database)).toEqual([]) }) + + it('forgets the record once the host is placed off the edge', async () => { + const admin = { + // Moves and fences are both assign calls. + assign: async (body: EdgeAssignBody) => ({ outcome: 'applied' as const, assignmentEpoch: body.assignmentEpoch }) + } as unknown as EdgeAdmin + const { database, edge } = await setup(admin) + await edge.moveToEdge(identity, 0, 'asia-east2') + vi.spyOn(console, 'warn').mockImplementation(() => {}) + await edge.fenceAway(identity, 'us-c1', 9) + expect(await recorded(database)).toEqual([]) + }) }) describe('edge evacuation by recorded region', () => { @@ -339,13 +372,19 @@ describe('edge evacuation by recorded region', () => { ) if (region) { await database.query( - `INSERT INTO relay_assignment_region_preferences (user_id, relay_host_id, preferred_region, observed_at) - VALUES ('user-1', ?, ?, 0)`, + `INSERT INTO relay_edge_host_regions (user_id, relay_host_id, region, recorded_at) VALUES ('user-1', ?, ?, 0)`, [relayHostId, region] ) } } - return { evacuation, drained, onEdge } + // What the cell path stores: never read by the sweep. + const prefer = async (relayHostId: string, region: RelayRegion) => + await database.query( + `INSERT INTO relay_assignment_region_preferences (user_id, relay_host_id, preferred_region, observed_at) + VALUES ('user-1', ?, ?, 0)`, + [relayHostId, region] + ) + return { database, evacuation, drained, onEdge, prefer } } it('drains an aborted region by its recorded hosts, and keeps the rest', async () => { @@ -359,7 +398,7 @@ describe('edge evacuation by recorded region', () => { expect(h.drained).toEqual([asia]) }) - it('judges a recorded region by its effective share: global when env-scoped, out when not', async () => { + it('judges a recorded region by its effective share: global when env-scoped', async () => { // A us-central1 shrink (global 0) while Asia is at 1000: recorded us-central1 hosts drain; // an unrecorded host is judged 'any' (Asia's 1000) and stays until its next ask records it. const h = await harness({ mode: 'on', rolloutPermille: 0, regionRolloutPermille: { 'asia-east2': 1000 } }) @@ -379,12 +418,31 @@ describe('edge evacuation by recorded region', () => { expect(h.drained).toEqual([hosts[40]]) }) - it('without regional flags ignores recorded regions entirely (a lagging cell-path preference)', async () => { + it('ignores a stale cell path preference: only the edge record counts', async () => { + const h = await harness({ mode: 'on', rolloutPermille: 1000, regionRolloutPermille: { 'asia-east2': 0 } }) + // Asked Asia on a cell weeks ago; on the edge now for us-central1. + await h.onEdge(hosts[50]!, 'us-central1') + await h.prefer(hosts[50]!, 'asia-east2') + await h.onEdge(hosts[51]!) + await h.prefer(hosts[51]!, 'asia-east2') + await h.evacuation.sweep() + expect(h.drained).toEqual([]) + }) + + it('without regional flags ignores records and sends the same SQL as before', async () => { for (const flags of [{ mode: 'on', rolloutPermille: 1000 }, { mode: 'on', rolloutPermille: 1000, regionRolloutPermille: {} }] as EdgeControlFlags[]) { const h = await harness(flags) await h.onEdge(hosts[30]!, 'asia-east2') + const query = vi.spyOn(h.database, 'query') await h.evacuation.sweep() expect(h.drained).toEqual([]) + const page = query.mock.calls.map(([sql]) => sql).find((sql) => sql.includes('FROM relay_assignments a')) + expect(page).toBe(`SELECT a.user_id, a.relay_host_id + FROM relay_assignments a + WHERE a.cell_id = ? AND (a.user_id > ? OR (a.user_id = ? AND a.relay_host_id > ?)) + + ORDER BY a.user_id, a.relay_host_id + LIMIT ?`) } }) }) diff --git a/cloud/apps/relay/src/edge-relay.ts b/cloud/apps/relay/src/edge-relay.ts index 89ded320458..9051c8b0405 100644 --- a/cloud/apps/relay/src/edge-relay.ts +++ b/cloud/apps/relay/src/edge-relay.ts @@ -2,7 +2,6 @@ import { createHash } from 'node:crypto' import { ASSIGNMENT_LIMITS } from '@orca-cloud/relay-contract' import { RELAY_DEFAULT_REGION, type RelayRegion } from '@orca-cloud/relay-contract/relay-regions' import type { RelayAssignment } from './assignment/assignment-store.js' -import { recordRegionPreference } from './assignment/region-preference.js' import { lockedSelector, registerExistingOnlyCell, requireSelectorMatchesAdmission } from './cell-admission-selector.js' import type { RelayDatabase, SqlRow } from './database.js' import { fenceEdge } from './edge-relay-fence.js' @@ -20,7 +19,6 @@ import { EDGE_CONTROL_DEFAULTS, edgeAnyRegionRolloutPermille, edgeCountryExcluded, - edgeRegionalRollout, edgeRegionRolloutPermille, edgeRolloutPermille, type EdgeControlFlags @@ -288,16 +286,22 @@ export class EdgeAssignments { // The host's own ask: moves its row to the edge, minting a strictly newer epoch only when the cell // changes (phones accept a move only with a newer epoch; re-asks must not flap the epoch). - // While any regional share is set, every committed answer also records the region the cohort was - // judged for, so the evacuation sweep can drain a region's hosts (relay_assignment_region_ - // preferences); with none, nothing extra is written. Written last, under the assignment row lock - // (the cell path's order), so its row lock is held only to the commit. + // Every committed answer also records the region the cohort was judged for (relay_edge_host_ + // regions), so the evacuation sweep can drain one region's hosts: always on a move, and on a kept + // host only when the region changed (or was never recorded). Written last, under the assignment + // row lock, so its row lock is held only to the commit. async moveToEdge(identity: Identity, mapFloor: number, region?: RelayRegion): Promise { const now = this.now() const leaseExpiresAt = now + ASSIGNMENT_LIMITS.activityLeaseMs - const recordRegion = async (transaction: RelayDatabase) => { - if (!edgeRegionalRollout(this.config.control?.() ?? EDGE_CONTROL_DEFAULTS)) return - await recordRegionPreference(transaction, identity, region ?? RELAY_DEFAULT_REGION, now) + const recordRegion = async (transaction: RelayDatabase, onlyIfChanged = false) => { + await transaction.query( + `INSERT INTO relay_edge_host_regions (user_id, relay_host_id, region, recorded_at) + VALUES (?, ?, ?, ?) + ON CONFLICT (user_id, relay_host_id) DO UPDATE + SET region = excluded.region, recorded_at = excluded.recorded_at + ${onlyIfChanged ? 'WHERE relay_edge_host_regions.region IS DISTINCT FROM excluded.region' : ''}`, + [identity.userId, identity.relayHostId, region ?? RELAY_DEFAULT_REGION, now] + ) } return await this.database.transaction(async (transaction) => { let rows: SqlRow[] @@ -347,7 +351,7 @@ export class EdgeAssignments { WHERE user_id = ? AND relay_host_id = ?`, [leaseExpiresAt, now, identity.userId, identity.relayHostId] ) - await recordRegion(transaction) + await recordRegion(transaction, true) return this.answer(identity, integer(row, 'assignment_epoch'), leaseExpiresAt) } if (requestUnits(row) > 0) throw new EdgeMoveDeferredError('units_held') @@ -513,6 +517,20 @@ export class EdgeAssignments { // is recorded for the evacuation sweep to retry (edge-relay-fence.ts). async fenceAway(identity: Identity, cellId: string, assignmentEpoch: number): Promise { if (!this.admin) return true + // Off the edge: its region record goes (the sweep reads only hosts on the edge anyway). + await this.database + .query(`DELETE FROM relay_edge_host_regions WHERE user_id = ? AND relay_host_id = ?`, [ + identity.userId, + identity.relayHostId + ]) + .catch((error: unknown) => + console.warn( + JSON.stringify({ + event: 'orca_relay_edge_region_forget_failed', + reason: error instanceof Error ? error.message : 'unknown' + }) + ) + ) // An edge whose exports are failing (or that is marked down) gets one try, so a skipped export // is not slowed further. const down = this.exportFailingSince !== undefined || this.config.control?.().mode === 'down' diff --git a/cloud/apps/relay/src/ip-country.test.ts b/cloud/apps/relay/src/ip-country.test.ts index 8284856c9ce..806a27b256e 100644 --- a/cloud/apps/relay/src/ip-country.test.ts +++ b/cloud/apps/relay/src/ip-country.test.ts @@ -48,6 +48,11 @@ describe('IP country table', () => { expect(table.country('192.0.2.128')).toBe('HK') expect(table.country('198.51.100.200')).toBe('US') expect(table.country('::ffff:192.0.2.200')).toBe('HK') + // The hex spelling of the same mapped addresses, and a fully written-out one. + expect(table.country('::ffff:c000:2c8')).toBe('HK') + expect(table.country('::FFFF:C000:0201')).toBe('CN') + expect(table.country('0:0:0:0:0:ffff:c633:6401')).toBe('US') + expect(table.country('2001:db8::192.0.2.1')).toBe('CN') expect(table.country('2001:db8::1')).toBe('CN') expect(table.country('2001:db8:0:ffff::1')).toBe('CN') expect(table.country('2001:db8:1:2::')).toBe('JP') diff --git a/cloud/apps/relay/src/ip-country.ts b/cloud/apps/relay/src/ip-country.ts index a66238a84fb..a48647383e0 100644 --- a/cloud/apps/relay/src/ip-country.ts +++ b/cloud/apps/relay/src/ip-country.ts @@ -62,7 +62,7 @@ export function parseIpCountryTable(text: string): IpCountryLookup { if (isIPv4(first) && isIPv4(last)) { v4.push([ipv4Number(first), ipv4Number(last), index]) } else if (isIPv6(first) && isIPv6(last)) { - v6.push([ipv6Prefix64(first), ipv6Prefix64(last), index]) + v6.push([ipv6Prefix64(ipv6Groups(first)), ipv6Prefix64(ipv6Groups(last)), index]) } } if (v4.length + v6.length === 0) throw new Error('ip_country_table_empty') @@ -76,13 +76,21 @@ export function parseIpCountryTable(text: string): IpCountryLookup { const v6Code = Uint16Array.from(v6, (range) => range[2]) return { country(ip) { - const mapped = /^::ffff:(\d+\.\d+\.\d+\.\d+)$/i.exec(ip)?.[1] - if (isIPv4(mapped ?? ip)) { - const at = rangeAt(v4First, v4Last, ipv4Number(mapped ?? ip)) + let v4 = isIPv4(ip) ? ipv4Number(ip) : undefined + let groups: number[] | undefined + if (v4 === undefined) { + if (!isIPv6(ip)) return undefined + groups = ipv6Groups(ip) + // IPv4-mapped (::ffff:a.b.c.d, or its hex spelling ::ffff:abcd:efgh) is the IPv4 address. + if (groups.slice(0, 5).every((group) => group === 0) && groups[5] === 0xffff) { + v4 = groups[6]! * 0x1_0000 + groups[7]! + } + } + if (v4 !== undefined) { + const at = rangeAt(v4First, v4Last, v4) return at === undefined ? undefined : codes[v4Code[at]!] } - if (!isIPv6(ip)) return undefined - const at = rangeAt(v6First, v6Last, ipv6Prefix64(ip)) + const at = rangeAt(v6First, v6Last, ipv6Prefix64(groups!)) return at === undefined ? undefined : codes[v6Code[at]!] } } @@ -113,18 +121,22 @@ function ipv4Number(ip: string): number { return ip.split('.').reduce((value, octet) => value * 256 + Number(octet), 0) } -// The first 64 bits of an IPv6 address (a trailing dotted IPv4 part lies in the last 32). -function ipv6Prefix64(ip: string): bigint { +// The eight 16-bit groups of a valid IPv6 address (a trailing dotted IPv4 part is the last two). +function ipv6Groups(ip: string): number[] { const [head = '', tail] = ip.toLowerCase().split('::') - const headGroups = head ? head.split(':') : [] - const tailGroups = tail ? tail.split(':') : [] - // A dotted IPv4 tail is one token for two groups. - const width = (tail === undefined ? headGroups : tailGroups).at(-1)?.includes('.') ? 7 : 8 - const groups = - tail === undefined - ? headGroups - : [...headGroups, ...Array(width - headGroups.length - tailGroups.length).fill('0'), ...tailGroups] - return groups - .slice(0, 4) - .reduce((value, group) => (value << 16n) | BigInt(Number.parseInt(group, 16)), 0n) + const tokens = (part: string) => + (part ? part.split(':') : []).flatMap((token) => + token.includes('.') + ? [ipv4Number(token) >>> 16, ipv4Number(token) & 0xffff] + : [Number.parseInt(token, 16)] + ) + const headGroups = tokens(head) + if (tail === undefined) return headGroups + const tailGroups = tokens(tail) + return [...headGroups, ...Array(8 - headGroups.length - tailGroups.length).fill(0), ...tailGroups] +} + +// The first 64 bits (country data is allocated far coarser than /64). +function ipv6Prefix64(groups: number[]): bigint { + return groups.slice(0, 4).reduce((value, group) => (value << 16n) | BigInt(group), 0n) }