diff --git a/cloud/apps/relay/src/assignment-store.ts b/cloud/apps/relay/src/assignment-store.ts index 8c12872effc..17b25ad1ad2 100644 --- a/cloud/apps/relay/src/assignment-store.ts +++ b/cloud/apps/relay/src/assignment-store.ts @@ -3340,6 +3340,17 @@ export class RelayAssignmentStore { "SELECT enabled, not_before FROM relay_region_rehome_control WHERE control_id = 'global'" ))[0] if (!control || Number(control.enabled) !== 1 || Number(control.not_before) > now) return [] + // The dispatch budget is durable and global, but until now only + // `commitIdleRegionalRehome` consulted it -- after the join had already run and + // the worker had already POSTed every candidate to its source cell. An absent + // row means the budget has never been spent, so it opens the gate. + const worker = (await this.database.query( + `SELECT paused_until, next_dispatch_at FROM relay_region_rehome_worker_state + WHERE worker_id = 'global'` + ))[0] + if (worker && (Number(worker.paused_until) > now || Number(worker.next_dispatch_at) > now)) { + return [] + } const fleetSafety = await this.readRegionalRehomeFleetSafety(this.database, now) if (regionalRehomeFleetSafetyFailure(processSafety, fleetSafety, now)) return [] const candidates = await selectIdleRegionalRehomes({ diff --git a/cloud/apps/relay/src/idle-regional-rehome-store.test.ts b/cloud/apps/relay/src/idle-regional-rehome-store.test.ts index 92ced43b510..489f4ce0b53 100644 --- a/cloud/apps/relay/src/idle-regional-rehome-store.test.ts +++ b/cloud/apps/relay/src/idle-regional-rehome-store.test.ts @@ -182,6 +182,53 @@ describe('constrained idle regional assignment transaction', () => { }) }) + it.each(['next_dispatch_at', 'paused_until'] as const)( + 'skips the candidate join while %s holds the durable dispatch budget closed', + async (column) => { + const { store, database, safety } = await setup() + const query = vi.spyOn(database, 'query') + // One assignment only: setup leaves both fields at 0, and naming the other + // one too would assign this column twice, which Postgres rejects. + await database.query( + `UPDATE relay_region_rehome_worker_state SET ${column} = ? WHERE worker_id = 'global'`, + [safety.observedAt + 1] + ) + for (let tick = 0; tick < 3; tick++) { + query.mockClear() + expect(await store.selectIdleRegionalRehomeCandidates(safety)).toEqual([]) + expect(query).toHaveBeenCalledTimes(2) + expect(query.mock.calls[1]![0]).toMatch(/FROM relay_region_rehome_worker_state/s) + } + await database.query( + `UPDATE relay_region_rehome_worker_state SET ${column} = ? WHERE worker_id = 'global'`, + [safety.observedAt] + ) + query.mockClear() + expect(await store.selectIdleRegionalRehomeCandidates(safety)).toHaveLength(1) + expect(query.mock.calls.length).toBeGreaterThan(2) + } + ) + + it('polls when the worker state row has never been written', async () => { + const { store, database, safety } = await setup() + await database.query('DELETE FROM relay_region_rehome_worker_state') + expect(await store.selectIdleRegionalRehomeCandidates(safety)).toHaveLength(1) + }) + + it('leaves the candidate page offset untouched across a closed dispatch budget', async () => { + const { store, database, safety } = await setup() + const first = await store.selectIdleRegionalRehomeCandidates(safety) + await database.query( + `UPDATE relay_region_rehome_worker_state SET next_dispatch_at = ? WHERE worker_id = 'global'`, + [safety.observedAt + 1] + ) + expect(await store.selectIdleRegionalRehomeCandidates(safety)).toEqual([]) + await database.query( + `UPDATE relay_region_rehome_worker_state SET next_dispatch_at = 0 WHERE worker_id = 'global'` + ) + expect(await store.selectIdleRegionalRehomeCandidates(safety)).toEqual(first) + }) + it('progresses past a full page of busy candidates without writing eligibility state', async () => { const { store, database, safety } = await setup() for (const table of [