usenet-indexer/app/services/group_leases.ts

import db from '@adonisjs/lucid/services/db'
import { createReadStream, createWriteStream } from 'node:fs'
import { mkdir, rename, stat } from 'node:fs/promises'
import { join } from 'node:path'
import type { Readable } from 'node:stream'
import { pipeline } from 'node:stream/promises'
import coordinatorConfig from '#config/coordinator'
import RuleSets from '#services/rule_sets'
import CatalogSpace from '#services/catalog_space'
import { LeaseLost } from '#services/leases'
import type { GroupLease, NextRange } from '#lib/edge/index'

/**
 * Group leases (REBUILD.md, "Voyager 3: window indexing", 1 and 5): a worker thread holds a group and asks for its
 * next range each time it finished one, so a group's ranges are indexed in article order through one carried
 * window. The coordinator keeps one row a group (group_leases): the holder, its heartbeat, the range it works and
 * the snapshot its window stands at.
 *
 * - lease: the free group with the most urgent planned range (the ranges' priority: the tail before history, then
 *   the windows), the longest unheld first; a holder quiet for LEASE_TTL_SECONDS is taken over. Taking a group
 *   resets what the window's snapshot doesn't cover: the ranges from the snapshot's article on are planned again
 *   (the batch store's keys make the redo idempotent); without a snapshot the window is lost and the group
 *   continues at its resume point (one seam).
 * - next: the group's lowest planned range (or the one this node already has), leased to the node; none when the
 *   group is caught up: the worker snapshots and gives the group back, and gets it again when the tail planned
 *   more.
 * - snapshot: the serialized window, kept as a file on bulk (SNAPSHOT_DIR), given to the next holder.
 */
class GroupLeases {
  get dir(): string {
    return coordinatorConfig.snapshotDir
  }

  /** LEASE_CAPTURE_ONLY: only the corpus windows' ranges (the history's first buckets share their priorities). */
  private get captureOnly(): string {
    return coordinatorConfig.leaseCaptureOnly ? 'AND r.capture' : ''
  }

  async lease(nodeId: number): Promise<GroupLease | null> {
    const rules = await RuleSets.current()
    const ttl = coordinatorConfig.leaseTtlSeconds
    const { leasePriorityMin: min, leasePriorityMax: max } = coordinatorConfig
    // The catalog at its cap on the NVMe: the tail only (CatalogSpace), as the range leases did.
    const history = (await CatalogSpace.full()) ? 'AND r.priority < 0' : ''
    return db.transaction(async (trx) => {
      // One lease at a time: the choice reads several tables.
      await trx.rawQuery(`SELECT pg_advisory_xact_lock(hashtextextended('group-lease', 0))`)
      const { rows } = await trx.rawQuery(
        `SELECT g.id, g.name, g.priority, l.node_id AS holder, l.snapshot_next, todo.priority AS todo_priority
         FROM groups g
         JOIN LATERAL (SELECT min(r.priority) AS priority FROM ranges r WHERE r.group_id = g.id AND r.status = 'planned' AND r.next_id <= r.end_id AND r.priority BETWEEN ? AND ? ${history} ${this.captureOnly}) todo ON true
         LEFT JOIN group_leases l ON l.group_id = g.id
         WHERE g.active AND todo.priority IS NOT NULL
           AND (l.node_id IS NULL OR l.heartbeat_at < now() - make_interval(secs => ?))
         ORDER BY todo.priority, l.heartbeat_at NULLS FIRST, g.priority
         LIMIT 1`,
        [min, max, ttl]
      )
      const row = rows[0]
      if (!row) return null
      const groupId = Number(row.id)
      // What the snapshot doesn't cover is read again: the ranges from its article on, the one holding it from there.
      if (row.snapshot_next !== null) {
        await trx.rawQuery(
          `UPDATE ranges SET status = 'planned', node_id = NULL, next_id = greatest(start_id, ?::bigint)
           WHERE group_id = ? AND end_id >= ?::bigint AND (status IN ('leased', 'done') OR next_id > ?::bigint)`,
          [row.snapshot_next, groupId, row.snapshot_next, row.snapshot_next]
        )
      } else {
        await trx.rawQuery(
          `UPDATE ranges SET status = 'planned', node_id = NULL WHERE group_id = ? AND status = 'leased'`,
          [groupId]
        )
      }
      await trx.rawQuery(
        `INSERT INTO group_leases (group_id, node_id, leased_at, heartbeat_at, range_id)
         VALUES (?, ?, now(), now(), NULL)
         ON CONFLICT (group_id) DO UPDATE SET node_id = excluded.node_id, leased_at = now(), heartbeat_at = now(), range_id = NULL`,
        [groupId, nodeId]
      )
      await trx.from('nodes').where('id', nodeId).update({ last_seen_at: new Date() })
      return {
        group: { id: groupId, name: row.name, priority: Number(row.priority) },
        rules,
        snapshot: row.snapshot_next !== null ? { next: Number(row.snapshot_next) } : null,
      }
    })
  }

  /** The group's lease row, when this node holds it. */
  async held(groupId: number, nodeId: number) {
    const lease = await db.from('group_leases').where('group_id', groupId).first()
    if (!lease || Number(lease.node_id) !== nodeId)
      throw new LeaseLost(`Group ${groupId} is not held by this node`)
    return lease
  }

  /** The group's next range in article order (or the one this node has), leased to the node; null when caught up. */
  async next(groupId: number, nodeId: number): Promise<NextRange | null> {
    await this.held(groupId, nodeId)
    const rules = await RuleSets.current()
    const { leasePriorityMin: min, leasePriorityMax: max } = coordinatorConfig
    const mine = await db
      .from('ranges')
      .where('group_id', groupId)
      .where('status', 'leased')
      .where('node_id', nodeId)
      .orderBy('start_id')
      .first()
    const row =
      mine ??
      (
        await db.rawQuery(
          `UPDATE ranges r SET status = 'leased', node_id = ?, leased_at = now(), heartbeat_at = now(), attempts = attempts + 1, rule_version = ?
           WHERE r.id = (
             SELECT id FROM ranges r WHERE r.group_id = ? AND r.status = 'planned' AND r.next_id <= r.end_id AND r.priority BETWEEN ? AND ? ${this.captureOnly}
             ORDER BY r.start_id LIMIT 1 FOR UPDATE SKIP LOCKED)
           RETURNING r.id, r.start_id, r.end_id, r.next_id`,
          [nodeId, rules.version, groupId, min, max]
        )
      ).rows[0]
    await db
      .from('group_leases')
      .where('group_id', groupId)
      .update({ range_id: row ? Number(row.id) : null, heartbeat_at: new Date() })
    if (!row) return null
    return {
      range: {
        id: Number(row.id),
        start: Number(row.start_id),
        end: Number(row.end_id),
        next: Number(row.next_id),
      },
    }
  }

  async heartbeat(groupId: number, nodeId: number) {
    await this.held(groupId, nodeId)
    await db.from('group_leases').where('group_id', groupId).update({ heartbeat_at: new Date() })
    await db.from('nodes').where('id', nodeId).update({ last_seen_at: new Date() })
  }

  /** A range of the group finished by its holder: counted on the lease (Leases.finish does the range). */
  async rangeDone(groupId: number, headers: number) {
    await db.rawQuery(
      `UPDATE group_leases SET ranges_done = ranges_done + 1, headers = headers + ?, range_id = NULL WHERE group_id = ?`,
      [headers, groupId]
    )
  }

  /** The worker gives the group back (after its snapshot): its range stays at its resume point for the next holder. */
  async release(groupId: number, nodeId: number) {
    await this.held(groupId, nodeId)
    await db
      .from('ranges')
      .where('group_id', groupId)
      .where('status', 'leased')
      .where('node_id', nodeId)
      .update({ status: 'planned', node_id: null })
    await db
      .from('group_leases')
      .where('group_id', groupId)
      .update({ node_id: null, range_id: null })
  }

  private snapshotPath(groupId: number) {
    return join(this.dir, `${groupId}.gz`)
  }

  /**
   * The window's snapshot standing at `next` (the next article to index), streamed to disk (a busy group's is
   * gigabytes) and written whole before it replaces the old one.
   */
  async putSnapshot(groupId: number, nodeId: number, next: number, body: Readable) {
    await this.held(groupId, nodeId)
    await mkdir(this.dir, { recursive: true })
    const path = this.snapshotPath(groupId)
    await pipeline(body, createWriteStream(`${path}.part`))
    await rename(`${path}.part`, path)
    await db
      .from('group_leases')
      .where('group_id', groupId)
      .update({ snapshot_next: next, snapshot_at: new Date(), heartbeat_at: new Date() })
  }

  /** The snapshot's stream and size, or null when there is none. */
  async getSnapshot(
    groupId: number,
    nodeId: number
  ): Promise<{ stream: Readable; size: number; next: number } | null> {
    const lease = await this.held(groupId, nodeId)
    if (lease.snapshot_next === null) return null
    const path = this.snapshotPath(groupId)
    try {
      const { size } = await stat(path)
      return { stream: createReadStream(path), size, next: Number(lease.snapshot_next) }
    } catch {
      return null
    }
  }
}

export default new GroupLeases()