usenet-indexer/commands/edge_bench.ts

import { BaseCommand, flags } from '@adonisjs/core/ace'
import type { CommandOptions } from '@adonisjs/core/types/ace'
import { mkdir, readFile, readdir, writeFile } from 'node:fs/promises'
import { join } from 'node:path'
import { unzipSync } from 'node:zlib'
import env from '#start/env'
import { XoverClient } from '#lib/nntp/index'
import { RangeIndexer } from '#lib/edge/index'
import { DEFAULT_RULES, type Rules } from '#lib/rules/index'
import { decodeChunk, encodeChunk } from '#lib/chunks/index'
import db from '@adonisjs/lucid/services/db'
import BatchStore from '#services/batch_store'

/**
 * Measures the edge worker on a real range (REBUILD.md: TypeScript or Rust):
 * --capture saves a range's compressed XOVER chunks to a directory once;
 * without it the saved chunks are replayed through the range indexer, timing
 * inflate, parse and indexing per chunk, with memory and output sizes.
 */
export default class EdgeBench extends BaseCommand {
  static commandName = 'edge:bench'
  static description = 'Measures the edge worker on a captured range'
  static options: CommandOptions = { startApp: true }

  @flags.string({ description: 'Directory of the captured chunks', required: true })
  declare dir: string

  @flags.boolean({ description: 'Capture the range first (needs NNTP_* in the environment)' })
  declare capture: boolean

  @flags.string({ description: 'Group to capture' })
  declare group?: string

  @flags.number({ description: 'First article to capture' })
  declare start?: number

  @flags.number({ description: 'Articles to capture', default: 2_000_000 })
  declare articles: number

  @flags.number({ description: 'Articles per XOVER', default: 50_000 })
  declare chunk: number

  @flags.number({ description: 'Chunks per batch (results taken every N chunks)', default: 10 })
  declare batchChunks: number

  @flags.number({ description: 'Chunks without a new part before an open file is given up', default: 5 })
  declare idle: number

  @flags.boolean({ description: 'Decode every release chunk and encode it again: the same bytes' })
  declare verify: boolean

  @flags.boolean({ description: 'Also store the batches (BatchStore, the configured database), timed' })
  declare store: boolean

  @flags.string({ description: 'Indexed groups in priority order, comma-separated (owner rule)' })
  declare priorities?: string

  async run() {
    if (this.capture) await this.captureRange()
    await this.replay()
  }

  private async captureRange() {
    if (!this.group || this.start === undefined) throw new Error('--group and --start are needed to capture')
    await mkdir(this.dir, { recursive: true })
    const client = await XoverClient.connect({
      host: env.get('NNTP_HOST')!,
      port: env.get('NNTP_PORT')!,
      user: env.get('NNTP_USER'),
      password: env.get('NNTP_PASSWORD'),
      secure: env.get('NNTP_SECURE', true),
      compress: true,
      timeoutMs: 300_000,
    })
    await client.group(this.group)
    const end = this.start + this.articles - 1
    let index = 0
    for (let from = this.start; from <= end; from += this.chunk, index++) {
      const to = Math.min(from + this.chunk - 1, end)
      const started = Date.now()
      const raw = await client.xoverRaw(from, to)
      await writeFile(join(this.dir, `${String(index).padStart(4, '0')}-${from}-${to}.gz`), raw)
      this.logger.info(`chunk ${index}: ${from}–${to}, ${(raw.length / 1e6).toFixed(1)} MB in ${((Date.now() - started) / 1000).toFixed(1)} s`)
    }
    await writeFile(join(this.dir, 'range.json'), JSON.stringify({ group: this.group, start: this.start, end }))
    client.close()
  }

  private async replay() {
    const { group, start, end } = JSON.parse(await readFile(join(this.dir, 'range.json'), 'utf8')) as { group: string; start: number; end: number }
    const priorities = Object.fromEntries((this.priorities ?? group).split(',').map((g, i) => [g.trim(), i + 1]))
    if (!(group in priorities)) priorities[group] = Object.keys(priorities).length + 1
    const rules: Rules = { ...DEFAULT_RULES, priorities, droppedSignatures: [] }
    const files = (await readdir(this.dir)).filter((f) => f.endsWith('.gz')).sort()
    const indexer = new RangeIndexer({ group, start, end, rules, idleChunks: this.idle })
    const t = { inflate: 0, parse: 0, index: 0 }
    let compressed = 0
    let text = 0
    let peakRss = 0
    let peakDetail = ''
    const out = { headers: 0, releases: 0, files: 0, segments: 0, payload: 0, candidates: 0, fragments: 0, samples: 0 }
    const counts = new Map<string, number>()
    const keys = new Set<string>()
    let batches = 0
    let storeMs = 0
    let rangeId = 0
    let nodeId = 0
    if (this.store) {
      // Inactive: the coordinator neither plans nor leases it.
      const [g] = await db.table('groups').insert({ name: `bench:${group}:${Date.now()}`, priority: 1_000_000 + (Date.now() % 1_000_000), active: false }).returning('id')
      const [n] = await db.table('nodes').insert({ name: `bench-${Date.now()}` }).returning('id')
      nodeId = Number(n.id)
      const [r] = await db.table('ranges').insert({ group_id: g.id, start_id: start, end_id: end, next_id: start, status: 'leased', node_id: nodeId }).returning('id')
      rangeId = Number(r.id)
    }
    const cpuStart = process.cpuUsage()
    const wall = Date.now()
    for (const file of files) {
      const raw = await readFile(join(this.dir, file))
      compressed += raw.length
      let at = performance.now()
      const data = unzipSync(raw).toString('utf8')
      t.inflate += performance.now() - at
      text += data.length
      at = performance.now()
      const next = Number(file.split('-')[2].replace('.gz', '')) + 1
      // The worker's path: the reply's text straight into the indexer (parsing is part of "index").
      indexer.addText(data, next)
      const last = file === files.at(-1)
      if (last) indexer.finish()
      t.index += performance.now() - at
      if (!last && (files.indexOf(file) + 1) % this.batchChunks !== 0) continue
      at = performance.now()
      const batch = indexer.take()
      t.index += performance.now() - at
      if (this.store) {
        const storedAt = performance.now()
        await BatchStore.store(rangeId, nodeId, batch)
        storeMs += performance.now() - storedAt
      }
      out.headers += batch.headers
      out.releases += batch.releases.length
      batches++
      for (const r of batch.releases) {
        if (this.verify) {
          const decoded = decodeChunk(r.payload)
          if (!encodeChunk(decoded).equals(r.payload) || decoded.reduce((n, f) => n + f.segments.length, 0) !== r.segments) {
            throw new Error(`Release ${r.key}: its chunk doesn't decode to itself`)
          }
        }
        keys.add(r.key)
        out.files += r.files
        out.segments += r.segments
        out.payload += r.payload.length
      }
      out.candidates += batch.candidates.length
      out.fragments += batch.fragments.length
      out.samples += batch.samples.length
      for (const c of batch.counts) counts.set(`${c.rule} ${c.key}`.trim(), (counts.get(`${c.rule} ${c.key}`.trim()) ?? 0) + c.articles)
      const m = process.memoryUsage()
      if (m.rss > peakRss) {
        peakRss = m.rss
        peakDetail = `heap ${(m.heapUsed / 1e6).toFixed(0)}/${(m.heapTotal / 1e6).toFixed(0)} MB, external ${(m.external / 1e6).toFixed(0)} MB, array buffers ${(m.arrayBuffers / 1e6).toFixed(0)} MB`
      }
    }
    // Done, so nothing takes it for a range to work on (a leased one was set back to planned once).
    if (this.store) await db.from('ranges').where('id', rangeId).update({ status: 'done', done_at: new Date() })
    const cpu = process.cpuUsage(cpuStart)
    const cpuS = (cpu.user + cpu.system) / 1e6
    const total = t.inflate + t.parse + t.index
    this.logger.info(`${group} ${start}–${end}: ${files.length} chunks, ${(compressed / 1e6).toFixed(0)} MB compressed, ${(text / 1e6).toFixed(0)} MB text`)
    this.logger.info(`${out.headers.toLocaleString('en')} headers in ${(total / 1000).toFixed(1)} s (${Math.round(out.headers / (total / 1000)).toLocaleString('en')}/s on one core; CPU ${cpuS.toFixed(1)} s, wall ${((Date.now() - wall) / 1000).toFixed(1)} s)`)
    this.logger.info(`  inflate ${(t.inflate / 1000).toFixed(1)} s, parse ${(t.parse / 1000).toFixed(1)} s, index ${(t.index / 1000).toFixed(1)} s; peak RSS ${(peakRss / 1e6).toFixed(0)} MB (${peakDetail})`)
    this.logger.info(`  ${batches} batches of ${this.batchChunks} chunks: ${out.releases} release chunks for ${keys.size} releases (${(out.releases / Math.max(1, keys.size)).toFixed(2)} chunks per release)`)
    this.logger.info(`  kept: ${out.releases} release chunks, ${out.files} files, ${out.segments} segments, payload ${(out.payload / 1e6).toFixed(1)} MB (${(out.payload / Math.max(1, out.segments)).toFixed(1)} B/segment); ${out.candidates} candidates, ${out.fragments} fragments, ${out.samples} samples`)
    for (const [k, v] of [...counts].sort((a, b) => b[1] - a[1]).slice(0, 8)) this.logger.info(`  dropped ${k}: ${v.toLocaleString('en')} articles`)
    if (this.store) this.logger.info(`  stored ${batches} batches in ${(storeMs / 1000).toFixed(1)} s (${(storeMs / Math.max(1, out.releases)).toFixed(2)} ms per release chunk)`)
  }
}