usenet-indexer/app/services/prefetcher.ts

import db from '@adonisjs/lucid/services/db'
import logger from '@adonisjs/core/services/logger'
import coordinatorConfig from '#config/coordinator'
import NntpPool from '#services/nntp_pool'
import Spool from '#services/spool'

interface Target {
  id: number
  group: string
  start: number
  end: number
  next: number
  leased: boolean
}

/**
 * Keeps the news server busy ahead of the workers (REBUILD.md, prefetch):
 * the missing chunks of leased ranges first, then the next planned ranges,
 * about SPOOL_AHEAD_PER_LEASE per active lease (at least SPOOL_MIN_AHEAD),
 * within SPOOL_MAX_GB. One range per connection at a time, in order.
 */
export class Prefetcher {
  private busy = new Set<number>()
  private stopped = false

  /** One round: starts work for free connections; returns how many ranges it started. */
  async tick(): Promise<number> {
    const free = NntpPool.size - this.busy.size
    if (free <= 0) return 0
    const targets = await this.targets(free)
    for (const target of targets) {
      this.busy.add(target.id)
      this.spoolRange(target)
        .catch((error) => logger.warn(`Spooling range ${target.id}: ${error.message}`))
        .finally(() => this.busy.delete(target.id))
    }
    return targets.length
  }

  stop() {
    this.stopped = true
  }

  get active() {
    return this.busy.size
  }

  private async targets(limit: number): Promise<Target[]> {
    const busy = [...this.busy]
    const select = (q: any) =>
      q
        .from('ranges as r')
        .join('groups as g', 'g.id', 'r.group_id')
        .whereNull('r.spooled_at')
        .whereNotIn('r.id', busy.length > 0 ? busy : [0])
        .select('r.id', 'g.name as group', 'r.start_id', 'r.end_id', 'r.next_id', 'r.status')
    // Leased ranges whose chunks aren't all there: a worker is waiting for them.
    const leased = await select(db.query()).where('r.status', 'leased').orderBy('r.leased_at').limit(limit)
    const out: Target[] = leased.map(toTarget)
    if (out.length >= limit) return out

    const [{ leases }] = await db.from('ranges').where('status', 'leased').count('* as leases')
    const [{ ahead }] = await db.from('ranges').where('status', 'planned').where('spooled_chunks', '>', 0).count('* as ahead')
    const wanted = Math.max(coordinatorConfig.spool.minAhead, coordinatorConfig.spool.aheadPerLease * Number(leases)) - Number(ahead)
    if (wanted <= 0 || (await Spool.size()) >= coordinatorConfig.spool.maxBytes) return out
    const planned = await select(db.query()).where('r.status', 'planned').orderBy([{ column: 'r.priority' }, { column: 'r.id' }]).limit(Math.min(limit - out.length, wanted))
    return [...out, ...planned.map(toTarget)]
  }

  private async spoolRange(target: Target) {
    let count = 0
    for (const span of Spool.chunksOf({ start: target.start, end: target.end }, target.next)) {
      if (this.stopped) return
      if (!(await Spool.has(target.id, span))) {
        const data = await NntpPool.use(async (client) => {
          await client.group(target.group)
          return client.xoverRaw(span.from, span.to)
        })
        await Spool.write(target.id, span, data)
      }
      count++
      await db.from('ranges').where('id', target.id).update({ spooled_chunks: count })
    }
    await db.from('ranges').where('id', target.id).update({ spooled_at: new Date() })
  }
}

function toTarget(row: any): Target {
  return { id: Number(row.id), group: row.group, start: Number(row.start_id), end: Number(row.end_id), next: Number(row.next_id), leased: row.status === 'leased' }
}