orca/cloud/apps/relay/src/control-rebind-inventory-lock-postgres.test.ts

import { afterAll, beforeAll, describe, expect, it } from 'vitest'
import { RelayAssignmentStore } from './assignment-store.js'
import { openRelayDatabase, type RelayDatabase } from './database.js'

const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL
const describePostgres = databaseUrl ? describe : describe.skip

// Three cells: the inventory lock covers more than the rows a move touches, and
// a high-to-low move exposes any lock taken out of cell_id order.
const cells = [
  {
    id: 'rebind-inventory-postgres-a',
    url: 'https://rebind-inventory-postgres-a.example.com',
    capacityRequests: 1_000,
    connectionHardCap: 600 as const,
    connectionUnobservedBound: 50
  },
  {
    id: 'rebind-inventory-postgres-b',
    url: 'https://rebind-inventory-postgres-b.example.com',
    capacityRequests: 1_000,
    connectionHardCap: 600 as const,
    connectionUnobservedBound: 50
  },
  {
    id: 'rebind-inventory-postgres-c',
    url: 'https://rebind-inventory-postgres-c.example.com',
    capacityRequests: 1_000,
    connectionHardCap: 600 as const,
    connectionUnobservedBound: 50
  }
]
const identity = { userId: 'rebind-inventory-postgres-user', relayHostId: 'rebindinvhost001' }

function heartbeat(cell: (typeof cells)[number]) {
  return {
    cellId: cell.id,
    cellUrl: cell.url,
    cellIncarnation: '11111111-1111-4111-8111-111111111111',
    startedAt: 50,
    ready: true,
    observedRequests: 0,
    totalConnections: 0,
    inFlightConnections: 0,
    reservedConnectionUnits: 0,
    enforcedConnectionUnits: 0,
    connectionInclusionWatermark: 1,
    connectionHardCap: 600 as const,
    connectionUnobservedBound: 50
  }
}

// Same-cell rebind must proceed while another cell row is held; fleet-wide locking formerly waited to COMMIT and hit request 55P03.
describePostgres('PostgreSQL control rebind under a held cell row', () => {
  const databases: RelayDatabase[] = []

  beforeAll(async () => {
    databases.push(
      await openRelayDatabase({ databaseUrl, dataDir: '' }),
      await openRelayDatabase({ databaseUrl, dataDir: '' })
    )
  })

  async function removeTestRows(database: RelayDatabase): Promise<void> {
    await database.query(
      `DELETE FROM relay_control_connection_reservations WHERE user_id = ?`,
      [identity.userId]
    )
    for (const table of [
      'relay_assignment_activity_leases',
      'relay_post_drain_migration_pins',
      'relay_assignment_migration_incarnations',
      'relay_assignment_migrations',
      'relay_assignments'
    ]) {
      await database.query(`DELETE FROM ${table} WHERE user_id = ?`, [identity.userId])
    }
    for (const cell of cells) {
      for (const table of [
        'relay_cell_connection_snapshots',
        'relay_cell_connection_runtime',
        'relay_cell_connection_limits',
        'relay_cell_runtime',
        'relay_cells'
      ]) {
        await database.query(`DELETE FROM ${table} WHERE cell_id = ?`, [cell.id])
      }
    }
  }

  afterAll(async () => {
    if (databases[0]) await removeTestRows(databases[0])
    for (const connection of databases) await connection.close()
  })

  it("rebinds and supersedes a control while another cell's row is held", async () => {
    // A prior aborted run leaves connection snapshots that reject a replayed watermark.
    await removeTestRows(databases[0]!)
    const store = new RelayAssignmentStore(databases[0]!, () => 100)
    await store.reconcileCells(cells)
    for (const cell of cells) await store.recordCellHeartbeat(heartbeat(cell))
    // Pin the host to cell A so placement is deterministic.
    await store.setCellEnabled(cells[1]!.id, false)
    await store.setCellEnabled(cells[2]!.id, false)
    const assignment = await store.assign(identity)
    expect(assignment.cellId).toBe(cells[0]!.id)
    await store.setCellEnabled(cells[1]!.id, true)
    await store.setCellEnabled(cells[2]!.id, true)
    await store.activateControl(identity, {
      cellId: cells[0]!.id,
      assignmentEpoch: assignment.assignmentEpoch,
      generation: 1,
      connectionInclusionWatermark: 10
    })

    // Hold only cell B's row on a second connection, the way a rebind on B
    // does, for longer than the request-path lock bound.
    let releaseInventory!: () => void
    const inventoryReleased = new Promise<void>((resolve) => {
      releaseInventory = resolve
    })
    let inventoryHeld!: () => void
    const inventoryHeldPromise = new Promise<void>((resolve) => {
      inventoryHeld = resolve
    })
    const holder = databases[1]!.transaction(async (transaction) => {
      await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cells[1]!.id])
      inventoryHeld()
      await inventoryReleased
    })
    let deadline: ReturnType<typeof setTimeout> | undefined
    let reboundWork: Promise<string> | undefined
    try {
      await inventoryHeldPromise

      // A generation-2 rebind on cell A supersedes generation 1. It must not
      // wait on cell B's row.
      const startedAt = Date.now()
      const blockedStatement = async (): Promise<string> => {
        const rows = await databases[1]!.query(
          `SELECT left(query, 160) AS q FROM pg_stat_activity
         WHERE datname = current_database() AND wait_event_type = 'Lock'`
        )
        return rows.map((row) => String(row.q)).join(' | ')
      }
      const timeout = new Promise<never>((_, reject) => {
        deadline = setTimeout(
          () =>
            void blockedStatement()
              .then((statement) =>
                reject(new Error(`rebind on cell A blocked behind cell B's row: ${statement}`))
              )
              .catch(reject),
          2_000
        )
      })
      reboundWork = store.activateControl(identity, {
        cellId: cells[0]!.id,
        assignmentEpoch: assignment.assignmentEpoch,
        generation: 2,
        connectionInclusionWatermark: 11
      })
      const rebound = await Promise.race([reboundWork, timeout])
      const elapsedMs = Date.now() - startedAt

      expect(rebound).toBe(`control:${cells[0]!.id}:2`)
      expect(elapsedMs).toBeLessThan(2_000)
      const controls = await databases[0]!.query(
        `SELECT activity_id FROM relay_assignment_activity_leases
       WHERE user_id = ? AND activity_kind = 'control' ORDER BY activity_id`,
        [identity.userId]
      )
      expect(controls).toEqual([{ activity_id: `control:${cells[0]!.id}:2` }])
      const reserved = await databases[0]!.query(
        `SELECT reserved_requests FROM relay_cells WHERE cell_id = ?`,
        [cells[0]!.id]
      )
      expect(Number(reserved[0]!.reserved_requests)).toBe(1)
    } finally {
      clearTimeout(deadline)
      releaseInventory()
      try {
        await holder
      } finally {
        await Promise.allSettled(reboundWork ? [reboundWork] : [])
      }
    }
  }, 15_000)

  // Client activity IDs can migrate across cells; high→low moves expose lock cycles against ascending placement order.
  it('moves an activity from a higher cell to a lower one in cell_id order', async () => {
    await removeTestRows(databases[0]!)
    const [cellA, cellB, cellC] = cells as [(typeof cells)[0], (typeof cells)[0], (typeof cells)[0]]
    const store = new RelayAssignmentStore(databases[0]!, () => 100)
    await store.reconcileCells(cells)
    for (const cell of cells) await store.recordCellHeartbeat(heartbeat(cell))
    await store.setCellEnabled(cellA.id, false)
    await store.setCellEnabled(cellB.id, false)
    const assignment = await store.assign(identity)
    expect(assignment.cellId).toBe(cellC.id)
    await store.setCellEnabled(cellA.id, true)
    await store.setCellEnabled(cellB.id, true)
    const activityId = 'splice:rebind-inventory-postgres'
    await store.acquireActivity(identity, { activityId, kind: 'splice', cellId: cellC.id })
    // The migration makes B authoritative; the lease still sits on C.
    const migration = await store.startEvacuation(identity, cellB.id)
    expect(migration.targetCellId).toBe(cellB.id)

    const movingDatabase = await openRelayDatabase({
      databaseUrl,
      dataDir: '',
      applicationName: 'orca-relay/cloud3-rebind-move'
    })
    const moveStore = new RelayAssignmentStore(movingDatabase, () => 100)
    // Hold B: ordered B→C move waits with no other lock; C-first or fleet-wide locking would hold C or A.
    let releaseRow!: () => void
    const rowReleased = new Promise<void>((resolve) => {
      releaseRow = resolve
    })
    let rowHeld!: () => void
    const rowHeldPromise = new Promise<void>((resolve) => {
      rowHeld = resolve
    })
    const heldWhileMoverWaits: string[] = []
    const holder = databases[1]!.transaction(async (transaction) => {
      await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cellB.id])
      rowHeld()
      await rowReleased
      for (const cell of [cellA, cellC]) {
        try {
          await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cell.id], {
            failIfUnavailable: true
          })
        } catch {
          heldWhileMoverWaits.push(cell.id)
        }
      }
    })
    await rowHeldPromise
    const move = moveStore.acquireActivity(identity, {
      activityId,
      kind: 'splice',
      cellId: cellB.id
    })
    let moved = false
    const observedMove = move.then(() => {
      moved = true
    })
    try {
      await expect
        .poll(
          async () => {
            const rows = await databases[0]!.query(
              `SELECT count(*) AS waiting FROM pg_stat_activity
        WHERE datname = current_database() AND application_name = ? AND wait_event_type = 'Lock'`,
              ['orca-relay/cloud3-rebind-move']
            )
            return Number(rows[0]!.waiting)
          },
          { timeout: 2_000, interval: 10 }
        )
        .toBe(1)
      expect(moved).toBe(false)
      releaseRow()
      await holder
      await observedMove
      expect(heldWhileMoverWaits).toEqual([])
    } finally {
      releaseRow()
      try {
        await holder
      } finally {
        try {
          await observedMove
        } finally {
          await movingDatabase.close()
        }
      }
    }

    const reservations = await databases[0]!.query(
      `SELECT cell_id, reserved_requests FROM relay_cells
       WHERE cell_id IN (?, ?, ?) ORDER BY cell_id ASC`,
      [cellA.id, cellB.id, cellC.id]
    )
    const reserved = reservations.map((row) => [String(row.cell_id), Number(row.reserved_requests)])
    expect(reserved).toEqual([
      [cellA.id, 0],
      // Migration grant plus the moved splice, as in the SQLite origin-scoped
      // reservation case: the lock change did not alter accounting.
      [cellB.id, 6],
      // The sticky grant stays on the source until the migration completes.
      [cellC.id, 1]
    ])
  }, 15_000)
})