usenet-indexer/tests/functional/group_leases.spec.ts

import { test } from '@japa/runner'
import db from '@adonisjs/lucid/services/db'
import testUtils from '@adonisjs/core/services/test_utils'
import { mkdtemp, rm } from 'node:fs/promises'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import coordinatorConfig from '#config/coordinator'
import GroupLeases from '#services/group_leases'
import { LeaseLost } from '#services/leases'
import { Readable } from 'node:stream'

async function node(name: string) {
  return Number((await db.table('nodes').insert({ name }).returning('id'))[0].id)
}

async function group(name: string, priority: number) {
  return Number((await db.table('groups').insert({ name, priority }).returning('id'))[0].id)
}

async function range(groupId: number, start: number, priority: number, extra: Record<string, unknown> = {}) {
  const [row] = await db
    .table('ranges')
    .insert({ group_id: groupId, start_id: start, end_id: start + 99, next_id: start, priority, ...extra })
    .returning('id')
  return Number(row.id)
}

const rangeRow = async (id: number) => (await db.from('ranges').where('id', id).firstOrFail()) as any

test.group('Group leases', (g) => {
  let dir = ''
  let snapshotDir = ''
  g.each.setup(async () => {
    const rollback = await testUtils.db().withGlobalTransaction()
    dir = await mkdtemp(join(tmpdir(), 'snapshots-'))
    snapshotDir = coordinatorConfig.snapshotDir
    coordinatorConfig.snapshotDir = dir
    return rollback
  })
  g.each.teardown(async () => {
    coordinatorConfig.snapshotDir = snapshotDir
    await rm(dir, { recursive: true, force: true })
  })

  test('a lease takes the group with the most urgent planned range; its ranges come in article order; caught up it gives nothing', async ({
    assert,
  }) => {
    const a = await group('alt.binaries.gl.a', 1)
    const b = await group('alt.binaries.gl.b', 2)
    // Group a: history at priority 0 and 1; group b: a tail range (priority -3) and history.
    const a1 = await range(a, 1, 0)
    const a2 = await range(a, 101, 0)
    const a3 = await range(a, 201, 1)
    const b2 = await range(b, 101, 0)
    const b1 = await range(b, 1, -3)
    const n1 = await node('gl-1')
    const n2 = await node('gl-2')

    const first = await GroupLeases.lease(n1)
    assert.equal(first?.group.id, b, 'the tail first')
    assert.isNull(first?.snapshot)
    // Held: the next node gets the other group; a third gets nothing.
    const second = await GroupLeases.lease(n2)
    assert.equal(second?.group.id, a)
    assert.isNull(await GroupLeases.lease(await node('gl-3')))

    // Ranges in article order, whatever their priority.
    assert.deepEqual((await GroupLeases.next(b, n1))?.range, { id: b1, start: 1, end: 100, next: 1 })
    assert.include(await rangeRow(b1), { status: 'leased', node_id: n1 })
    // Asked again before finishing: the same range.
    assert.equal((await GroupLeases.next(b, n1))?.range.id, b1)
    await db.from('ranges').where('id', b1).update({ status: 'done', next_id: 101 })
    assert.equal((await GroupLeases.next(b, n1))?.range.id, b2)
    await db.from('ranges').where('id', b2).update({ status: 'done', next_id: 201 })
    assert.isNull(await GroupLeases.next(b, n1), 'caught up')

    assert.deepEqual([a1, a2, a3].length, 3)
    await assert.rejects(() => GroupLeases.next(a, n1), LeaseLost)
  })

  test('a stale holder is taken over and the ranges after the snapshot are read again; a release keeps the resume point', async ({
    assert,
  }) => {
    const a = await group('alt.binaries.gl.c', 1)
    const r1 = await range(a, 1, 0)
    const r2 = await range(a, 101, 0)
    const r3 = await range(a, 201, 0)
    const r4 = await range(a, 301, 0)
    const n1 = await node('gl-4')
    const n2 = await node('gl-5')
    assert.equal((await GroupLeases.lease(n1))?.group.id, a)
    // The worker did r1 and r2, snapshotted at 201 (r3's start), then got halfway through r3 and went quiet.
    await db.from('ranges').whereIn('id', [r1, r2]).update({ status: 'done', next_id: db.raw('end_id + 1') })
    await GroupLeases.putSnapshot(a, n1, 201, Readable.from([Buffer.from('window at 201')]))
    await GroupLeases.next(a, n1)
    await db.from('ranges').where('id', r3).update({ next_id: 251 })
    await db.from('group_leases').where('group_id', a).update({ heartbeat_at: new Date(Date.now() - 24 * 3600_000) })

    const lease = await GroupLeases.lease(n2)
    assert.equal(lease?.group.id, a)
    assert.deepEqual(lease?.snapshot, { next: 201 })
    // r1 and r2 stay done (before the snapshot); r3 is read again from the snapshot's article, r4 untouched.
    assert.include(await rangeRow(r1), { status: 'done' })
    assert.include(await rangeRow(r2), { status: 'done' })
    assert.include(await rangeRow(r3), { status: 'planned', node_id: null })
    assert.equal(Number((await rangeRow(r3)).next_id), 201)
    assert.include(await rangeRow(r4), { status: 'planned' })
    const snapshot = await GroupLeases.getSnapshot(a, n2)
    assert.equal(snapshot?.next, 201)
    assert.equal(snapshot?.size, 'window at 201'.length)
    await assert.rejects(() => GroupLeases.next(a, n1), LeaseLost)

    // The new holder works r3, snapshots at 301 and gives the group back: r3 done, r4 waits at its start.
    assert.equal((await GroupLeases.next(a, n2))?.range.id, r3)
    await db.from('ranges').where('id', r3).update({ status: 'done', next_id: 301 })
    await GroupLeases.putSnapshot(a, n2, 301, Readable.from([Buffer.from('window at 301')]))
    assert.equal((await GroupLeases.next(a, n2))?.range.id, r4)
    await db.from('ranges').where('id', r4).update({ next_id: 301 })
    await GroupLeases.release(a, n2)
    assert.include(await rangeRow(r4), { status: 'planned', node_id: null })
    assert.isNull((await db.from('group_leases').where('group_id', a).firstOrFail()).node_id)
    // Leased again at once (free, not stale): the snapshot at 301 covers everything done.
    const again = await GroupLeases.lease(n1)
    assert.deepEqual(again?.snapshot, { next: 301 })
    assert.include(await rangeRow(r3), { status: 'done' })
    assert.equal((await GroupLeases.next(a, n1))?.range.id, r4)
  })

  test('a snapshot older than a done range makes that range be read again from its start', async ({ assert }) => {
    const a = await group('alt.binaries.gl.d', 1)
    const r1 = await range(a, 1, 0)
    const r2 = await range(a, 101, 0)
    const n1 = await node('gl-6')
    await GroupLeases.lease(n1)
    // The worker snapshotted at the start, finished r1, and crashed before its next snapshot: r1 is read again.
    await GroupLeases.putSnapshot(a, n1, 1, Readable.from([Buffer.from('window at 1')]))
    await db.from('ranges').where('id', r1).update({ status: 'done', next_id: 101 })
    await db.from('group_leases').where('group_id', a).update({ heartbeat_at: new Date(Date.now() - 24 * 3600_000) })
    await GroupLeases.lease(await node('gl-7'))
    assert.equal((await rangeRow(r1)).status, 'planned')
    assert.equal(Number((await rangeRow(r1)).next_id), 1)
    assert.equal((await rangeRow(r2)).status, 'planned')
    assert.equal(Number((await rangeRow(r2)).next_id), 101)
  })
})