usenet-indexer/app/services/SeedOrchestrator.ts

import { createHash, randomBytes } from 'node:crypto'
import db from '@adonisjs/lucid/services/db'
import NntpService from '#services/NntpService'
import SeedService from '#services/SeedService'
import BackfillService from '#services/BackfillService'
import SeedFetchPool from '#services/SeedFetchPool'
import { articleAtDate } from '#lib/nntp/date_search'
import { normalizeSeedFile, type SeedBatch } from '#lib/seed/wire'

/** A lease without a batch for this long is given to another node; the range resumes at next_id. */
export const LEASE_TTL_SECONDS = 15 * 60
/** Leases per range before it counts as failed (a range that keeps breaking nodes). */
const MAX_ATTEMPTS = 5
/** Articles per XOVER probe when looking for an article by date. */
const PROBE = 50

export interface SeedLease {
  id: number
  group: string
  /** Resume point: first article not reported yet. */
  start: number
  end: number
}

export class LeaseLost extends Error {}

const hash = (token: string) => createHash('sha256').update(token).digest('hex')

/**
 * The seed orchestrator (CONCEPT.md D38): seeding nodes lease article
 * ranges over HTTPS, report what they aggregated in batches (D37), and
 * finish or give back the range. Nodes only make outgoing requests; a lease
 * that goes quiet for LEASE_TTL_SECONDS passes to the next node that asks.
 */
class SeedOrchestrator {
  /** Creates a node, or gives an existing one a new token. Returns the token (stored hashed only). */
  async issueToken(name: string): Promise<string> {
    const token = `sn_${randomBytes(24).toString('base64url')}`
    await db
      .table('seed_nodes')
      .insert({ name, token_hash: hash(token), disabled: false })
      .onConflict('name')
      .merge({ token_hash: hash(token), disabled: false })
    return token
  }

  async nodeForToken(token: string): Promise<{ id: number; name: string } | null> {
    const node = await db.from('seed_nodes').where('token_hash', hash(token)).where('disabled', false).first()
    if (!node) return null
    if (!node.last_seen_at || Date.now() - new Date(node.last_seen_at).getTime() > 60_000) {
      await db.from('seed_nodes').where('id', node.id).update({ last_seen_at: new Date() })
    }
    return { id: Number(node.id), name: node.name }
  }

  /**
   * The next range for this node, oldest plan first; a quiet lease is taken
   * over. Nothing while the database is over the backfill's disk limits: the
   * node asks again a minute later.
   */
  async lease(nodeId: number): Promise<SeedLease | null> {
    if ((await BackfillService.room()) !== true) return null
    // Fully reported but never finished (the node went away at the very end): nothing left to fetch.
    await db.rawQuery(
      `UPDATE seed_ranges SET status = 'done', done_at = now(), node_id = NULL
       WHERE next_id > end_id AND (status = 'queued' OR (status = 'leased' AND heartbeat_at < now() - make_interval(secs => ?)))`,
      [LEASE_TTL_SECONDS]
    )
    // Ranges that broke node after node stop being handed out.
    await db.rawQuery(
      `UPDATE seed_ranges SET status = 'failed', error = coalesce(error, 'lease expired too often')
       WHERE status = 'leased' AND heartbeat_at < now() - make_interval(secs => ?) AND attempts >= ?`,
      [LEASE_TTL_SECONDS, MAX_ATTEMPTS]
    )
    const { rows } = await db.rawQuery(
      `UPDATE seed_ranges SET status = 'leased', node_id = ?, leased_at = now(), heartbeat_at = now(), attempts = attempts + 1
       WHERE id = (
         SELECT id FROM seed_ranges
         WHERE status = 'queued' OR (status = 'leased' AND heartbeat_at < now() - make_interval(secs => ?))
         ORDER BY id LIMIT 1 FOR UPDATE SKIP LOCKED)
       RETURNING id, group_name, next_id, end_id`,
      [nodeId, LEASE_TTL_SECONDS]
    )
    const row = rows[0]
    return row ? { id: Number(row.id), group: row.group_name, start: Number(row.next_id), end: Number(row.end_id) } : null
  }

  /** Stores a node's batch and moves the range's resume point; LeaseLost when the range is no longer this node's. */
  async batch(nodeId: number, rangeId: number, batch: SeedBatch) {
    const range = await this.leased(nodeId, rangeId)
    const stored = await SeedService.store(batch.files.map(normalizeSeedFile))
    await SeedService.storeFragments(batch.fragments, { start: Number(range.start_id), end: Number(range.end_id) })
    await db.rawQuery(
      `UPDATE seed_ranges SET next_id = greatest(next_id, least(?, end_id + 1)), heartbeat_at = now(),
         headers = headers + ?, files = files + ?, fragments = fragments + ?
       WHERE id = ?`,
      [batch.next, batch.headers, batch.files.length, batch.fragments.length, rangeId]
    )
    return stored
  }

  /**
   * Relay: overviews of `from`..`to` of a leased range, fetched here and handed
   * on as the gzip stream the news server sent. Keeps the lease alive.
   */
  async chunk(nodeId: number, rangeId: number, from: number, to: number): Promise<Buffer> {
    const range = await this.leased(nodeId, rangeId)
    const start = Math.max(from, Number(range.start_id))
    const end = Math.min(to, Number(range.end_id), start + 49_999)
    if (end < start) throw new RangeError(`Articles ${from}–${to} are outside range ${rangeId}`)
    await db.from('seed_ranges').where('id', rangeId).update({ heartbeat_at: new Date() })
    return SeedFetchPool.fetch(range.group_name, start, end)
  }

  async finish(nodeId: number, rangeId: number) {
    await this.leased(nodeId, rangeId)
    await db.from('seed_ranges').where('id', rangeId).update({ status: 'done', done_at: new Date(), heartbeat_at: new Date(), error: null })
  }

  /**
   * A node gives a range back (it failed, or is shutting down); it's queued
   * again unless it failed too often. `uncounted`: the node never started on
   * it (couldn't reach its news server), so the lease doesn't count as an attempt.
   */
  async release(nodeId: number, rangeId: number, error?: string, uncounted = false) {
    const range = await this.leased(nodeId, rangeId)
    const attempts = uncounted ? Math.max(0, range.attempts - 1) : range.attempts
    await db
      .from('seed_ranges')
      .where('id', rangeId)
      .update({ status: attempts >= MAX_ATTEMPTS ? 'failed' : 'queued', node_id: null, attempts, error: error?.slice(0, 2000) ?? null })
  }

  private async leased(nodeId: number, rangeId: number) {
    const range = await db.from('seed_ranges').where('id', rangeId).first()
    if (!range || range.status !== 'leased' || Number(range.node_id) !== nodeId) throw new LeaseLost(`Range ${rangeId} is not leased to this node`)
    return range
  }

  /**
   * Plans ranges for groups: from below what live indexing and the backfill
   * cover (or below earlier seed ranges) down to `days` back, or to the
   * server's oldest article. A group's backfill stops where its seed plan
   * begins. Ranges are interleaved across groups, newest first.
   */
  async plan(groups: string[], options: { days?: number; size: number }): Promise<{ group: string; from: number; to: number; ranges: number }[]> {
    const perGroup: { group: string; ranges: [number, number][] }[] = []
    const summary: { group: string; from: number; to: number; ranges: number }[] = []
    const conn = await NntpService.acquire()
    let failed = false
    try {
      for (const name of groups) {
        const group = await db.from('groups').where('name', name).first()
        if (!group) throw new Error(`Unknown group ${name}`)
        const { first, last } = await NntpService.timed(NntpService.selectGroup(conn, name), `GROUP ${name}`)
        const dateOf = async (n: bigint) => {
          const { overviews } = await NntpService.timed<any>(conn.xover(n, n + BigInt(PROBE)), `XOVER ${name} ${n}`)
          const header = overviews.find(([number, h]: [number | null, any]) => number !== null && h?.date)?.[1]
          const date = header ? Date.parse(header.date) : Number.NaN
          return Number.isNaN(date) ? null : date
        }
        // Top: where live indexing began, or the backfill's position, or below earlier seed ranges.
        let top = Number(await articleAtDate(first, last, new Date(group.created_at).getTime(), dateOf))
        const backfilling = group.backfill_next_id !== null && group.backfill_until_id !== null && Number(group.backfill_next_id) > Number(group.backfill_until_id)
        if (backfilling) top = Math.min(top, Number(group.backfill_next_id))
        const seeded = await db.from('seed_ranges').where('group_name', name).min('start_id as low').first()
        if (seeded?.low !== null && seeded?.low !== undefined) top = Math.min(top, Number(seeded.low))
        const bottom = options.days
          ? Number(await articleAtDate(first, last, Date.now() - options.days * 24 * 60 * 60 * 1000, dateOf))
          : Number(first)

        const ranges: [number, number][] = []
        for (let end = top - 1; end >= bottom; end -= options.size) ranges.push([Math.max(bottom, end - options.size + 1), end])
        perGroup.push({ group: name, ranges })
        summary.push({ group: name, from: bottom, to: top - 1, ranges: ranges.length })
        // The seed plan takes over the rest of this group's backfill.
        if (backfilling && ranges.length > 0) {
          await db.from('groups').where('id', group.id).update({ backfill_until_id: group.backfill_next_id })
        }
      }
    } catch (error) {
      failed = true
      throw error
    } finally {
      NntpService.release(conn, failed)
    }

    // Round-robin over groups, each newest first, so every group fills in from the present backwards.
    const rows: Record<string, unknown>[] = []
    for (let i = 0; perGroup.some((g) => i < g.ranges.length); i++) {
      for (const g of perGroup) {
        const range = g.ranges[i]
        if (range) rows.push({ group_name: g.group, start_id: range[0], end_id: range[1], next_id: range[0] })
      }
    }
    for (let i = 0; i < rows.length; i += 5000) await db.table('seed_ranges').multiInsert(rows.slice(i, i + 5000))
    return summary
  }

  async status() {
    const { rows: ranges } = await db.rawQuery(
      `SELECT status, count(*)::int AS ranges, sum(end_id - start_id + 1)::bigint AS articles,
              sum(next_id - start_id)::bigint AS reached, sum(headers)::bigint AS headers, sum(files)::bigint AS files
       FROM seed_ranges GROUP BY status ORDER BY status`
    )
    const { rows: nodes } = await db.rawQuery(
      `SELECT n.name, n.last_seen_at, n.disabled,
              (SELECT count(*)::int FROM seed_ranges r WHERE r.node_id = n.id AND r.status = 'leased') AS leases,
              (SELECT coalesce(sum(headers), 0)::bigint FROM seed_ranges r WHERE r.node_id = n.id) AS headers
       FROM seed_nodes n ORDER BY n.name`
    )
    const { rows: recent } = await db.rawQuery(
      `SELECT coalesce(sum(headers), 0)::bigint AS headers FROM seed_ranges WHERE done_at > now() - interval '1 hour'`
    )
    const room = await BackfillService.room()
    return { ranges, nodes, headersDoneLastHour: Number(recent[0].headers), paused: room === true ? null : room }
  }
}

export default new SeedOrchestrator()