diff --git a/cloud/apps/relay/src/assignment-store.ts b/cloud/apps/relay/src/assignment-store.ts index 3e81b37689f..824cb1e0b2f 100644 --- a/cloud/apps/relay/src/assignment-store.ts +++ b/cloud/apps/relay/src/assignment-store.ts @@ -5958,6 +5958,10 @@ export class RelayAssignmentStore { async recordRegionalRehomeDispatchFailure(attemptId: string): Promise { const now = this.now() const disableLog = await this.database.transaction(async (transaction) => { + // Match claim and enable ordering before a spent budget updates the control. + await transaction.queryLocked( + `SELECT * FROM relay_region_rehome_control WHERE control_id = 'global'` + ) const worker = ( await transaction.queryLocked( `SELECT * FROM relay_region_rehome_worker_state WHERE worker_id = 'global'` diff --git a/cloud/apps/relay/src/regional-rehome-postgres.test.ts b/cloud/apps/relay/src/regional-rehome-postgres.test.ts index 44f3b3434af..fdefda54401 100644 --- a/cloud/apps/relay/src/regional-rehome-postgres.test.ts +++ b/cloud/apps/relay/src/regional-rehome-postgres.test.ts @@ -1,4 +1,4 @@ -import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest' +import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' import { RelayAssignmentStore } from './assignment-store.js' import { openRelayDatabase, type RelayDatabase } from './database.js' import { @@ -267,6 +267,91 @@ describePostgres('PostgreSQL regional rehoming', () => { )).toEqual([{ count: '1' }]) }) + it('serializes an enable with a budget-exhausting failure without retries', async () => { + const context = await fixture() + const attempt = await context.store.claimRegionalRehome() + await context.store.recordRegionalRehomeDispatchFailure(attempt!.attemptId) + await context.store.recordRegionalRehomeDispatchFailure(attempt!.attemptId) + const locked = Promise.withResolvers() + const release = Promise.withResolvers() + const primaryTransaction = primary.transaction.bind(primary) + const secondaryTransaction = secondary.transaction.bind(secondary) + let enableTransactions = 0 + let failureTransactions = 0 + let enablePid = 0 + let failurePid = 0 + const enableSpy = vi.spyOn(primary, 'transaction').mockImplementation((operation, options) => + primaryTransaction(async (transaction) => { + enableTransactions++ + enablePid = Number((await transaction.query('SELECT pg_backend_pid() AS pid'))[0]!.pid) + return await operation({ + dialect: 'postgres', + query: transaction.query.bind(transaction), + queryLocked: async (sql, params, lockOptions) => { + const rows = await transaction.queryLocked(sql, params, lockOptions) + if (sql.includes('FROM relay_region_rehome_control')) { + locked.resolve() + await release.promise + } + return rows + }, + transaction: transaction.transaction.bind(transaction), + close: transaction.close.bind(transaction) + }) + }, options) + ) + const failureSpy = vi.spyOn(secondary, 'transaction').mockImplementation((operation, options) => + secondaryTransaction(async (transaction) => { + failureTransactions++ + failurePid = Number((await transaction.query('SELECT pg_backend_pid() AS pid'))[0]!.pid) + return await operation(transaction) + }, options) + ) + const enable = context.store.applyRegionalRehomeControl({ + expectedGeneration: 1, + enabled: true, + notBefore: context.now(), + ratePerMinute: 10, + preferenceMaxAgeMs: 24 * 60 * 60_000, + hostCooldownMs: 7 * 24 * 60 * 60_000, + drainGraceMs: 60_000 + }) + let failure: Promise | undefined + let outcomes: PromiseSettledResult[] = [] + try { + await Promise.race([ + locked.promise, + enable.then(() => { + throw new Error('enable completed before the control lock') + }) + ]) + failure = context.competingStore.recordRegionalRehomeDispatchFailure(attempt!.attemptId) + // Observe the actual PostgreSQL wait before letting enable acquire the worker row. + await vi.waitFor(async () => { + expect(failurePid).not.toBe(0) + const rows = await primary.query('SELECT pg_blocking_pids(?) AS blockers', [failurePid]) + expect(rows[0]!.blockers).toContain(enablePid) + }, { interval: 10, timeout: 800 }) + } finally { + release.resolve() + outcomes = await Promise.allSettled([enable, ...(failure ? [failure] : [])]) + enableSpy.mockRestore() + failureSpy.mockRestore() + } + expect(outcomes.map((outcome) => outcome.status)).toEqual(['fulfilled', 'fulfilled']) + expect({ enableTransactions, failureTransactions }).toEqual({ + enableTransactions: 1, + failureTransactions: 1 + }) + expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ + generation: 2, + enabled: true + }) + expect(await primary.query( + `SELECT consecutive_failures, paused_until FROM relay_region_rehome_worker_state` + )).toEqual([{ consecutive_failures: '1', paused_until: '0' }]) + }) + it('increments the disable generation once across competing directors', async () => { const context = await fixture() const disabled = await Promise.all([