usenet-indexer/app/services/fragment_merger.ts

import db from '@adonisjs/lucid/services/db'
import BatchStore from '#services/batch_store'
import ChunkPartitions from '#services/chunk_partitions'
import RuleSets from '#services/rule_sets'
import { RuleSet } from '#lib/rules/index'
import { candidatesOf, chunkOf, mergeFragments, releaseKeyOf, type Fragment, type ProbeCandidate, type ReleaseChunk } from '#lib/edge/index'
import type { AggregatedFile } from '#lib/files/index'
import type { EdgeSegment } from '#lib/edge/index'

export interface MergeResult {
  files: number
  completed: number
  incomplete: number
  waiting: number
  /** This round reached the end of the file keys (the next one starts over). */
  wrapped: boolean
}

/** File keys per round. */
const BATCH = 20_000

/**
 * Files cut by a range's edge (REBUILD.md, fragments), put together once
 * the ranges around them are done: a complete file is judged by the rules
 * like any other and written as a release chunk (with probe candidates); a
 * file still missing parts while every neighbouring range is done is
 * incomplete, and dropped (counted on its ranges). Others wait.
 */
class FragmentMerger {
  /** Where the next round goes on: waiting files must not take every round's places. */
  private cursor = ''

  async run(): Promise<MergeResult> {
    const result: MergeResult = { files: 0, completed: 0, incomplete: 0, waiting: 0, wrapped: false }
    const { rows: keys } = await db.rawQuery(
      `SELECT DISTINCT f.file_key FROM fragments f JOIN ranges r ON r.id = f.range_id
       WHERE r.status = 'done' AND f.file_key > ? ORDER BY f.file_key LIMIT ?`,
      [this.cursor, BATCH]
    )
    // At the end of the keys: the next round starts over.
    result.wrapped = keys.length < BATCH
    this.cursor = result.wrapped ? '' : keys.at(-1).file_key
    if (keys.length === 0) return result
    const rules = new RuleSet(await RuleSets.current())
    const rows = await db
      .from('fragments as f')
      .join('ranges as r', 'r.id', 'f.range_id')
      .whereIn('f.file_key', keys.map((k: { file_key: string }) => k.file_key))
      .select('f.id', 'f.file_key', 'f.range_id', 'f.fragment', 'r.group_id', 'r.start_id', 'r.end_id', 'r.status')
    const byKey = new Map<string, typeof rows>()
    for (const row of rows) {
      const list = byKey.get(row.file_key)
      if (list) list.push(row)
      else byKey.set(row.file_key, [row])
    }

    const completed: { file: AggregatedFile<EdgeSegment>; rangeId: number; ids: number[] }[] = []
    const incomplete: { rangeId: number; ids: number[]; segments: number }[] = []
    const neighbours = await this.neighbourStatus([...byKey.values()])
    for (const list of byKey.values()) {
      result.files++
      const file = mergeFragments(list.map((r) => r.fragment as Fragment))
      const ids = list.map((r) => Number(r.id))
      const rangeId = Math.max(...list.map((r) => Number(r.range_id)))
      if (file) {
        completed.push({ file, rangeId, ids })
        continue
      }
      if (list.every((r) => r.status === 'done') && neighbours(list)) {
        incomplete.push({ rangeId, ids, segments: list.reduce((n, r) => n + (r.fragment as Fragment).segments.length, 0) })
      } else {
        result.waiting++
      }
    }

    // Completed files, grouped into their releases, judged and written like a batch's.
    const byRelease = new Map<string, { files: AggregatedFile<EdgeSegment>[]; rangeId: number }>()
    for (const { file, rangeId } of completed) {
      const key = releaseKeyOf(file.segments[0])
      const entry = byRelease.get(key)
      if (entry) entry.files.push(file)
      else byRelease.set(key, { files: [file], rangeId })
    }
    const releases: { rangeId: number; chunk: ReleaseChunk; candidates: ProbeCandidate[] }[] = []
    for (const [key, { files, rangeId }] of byRelease) {
      const verdict = rules.judge(key, files.map((f) => ({ name: f.segments[0].file.name, totalParts: f.segments[0].totalParts, groups: f.segments[0].file.groups, firstMessageId: f.segments[0].messageId })))
      if (verdict.rule !== null) continue
      releases.push({ rangeId, chunk: chunkOf(key, files), candidates: candidatesOf(key, files) })
    }
    await ChunkPartitions.ensure()
    await db.transaction(async (trx) => {
      for (const { rangeId, chunk, candidates } of releases) await BatchStore.writeReleases(trx, rangeId, 0, [chunk], candidates)
      const done = [...completed.flatMap((c) => c.ids), ...incomplete.flatMap((c) => c.ids)]
      for (let i = 0; i < done.length; i += 5000) await trx.from('fragments').whereIn('id', done.slice(i, i + 5000)).delete()
      for (const { rangeId, segments } of incomplete) {
        await trx.rawQuery(
          `INSERT INTO rule_counts (range_id, rule, key, articles, files) VALUES (?, 'incomplete', 'edge', ?, 1)
           ON CONFLICT (range_id, rule, key) DO UPDATE SET articles = rule_counts.articles + excluded.articles, files = rule_counts.files + 1`,
          [rangeId, segments]
        )
      }
    })
    result.completed = completed.length
    result.incomplete = incomplete.length
    return result
  }

  /**
   * For each file's fragments: whether the ranges right before and after
   * their ranges (where a missing part could be) are done, read in one
   * query for the whole round (per file it was three, ~2,000 files a
   * minute). A neighbour not planned yet counts as not done (newer articles
   * come later), except before the group's oldest article.
   */
  private async neighbourStatus(lists: { group_id: number; start_id: number | string; end_id: number | string }[][]) {
    const span = (list: (typeof lists)[number]) => ({
      group: Number(list[0].group_id),
      first: Math.min(...list.map((r) => Number(r.start_id))),
      last: Math.max(...list.map((r) => Number(r.end_id))),
    })
    const spans = lists.map(span)
    const groups = [...new Set(spans.map((s) => s.group))]
    const { rows } = await db.rawQuery(
      `SELECT r.group_id, r.start_id, r.end_id, r.status FROM ranges r
       JOIN unnest(?::int[], ?::bigint[], ?::bigint[]) AS s(group_id, before_end, after_start)
         ON r.group_id = s.group_id AND (r.end_id = s.before_end OR r.start_id = s.after_start)`,
      [spans.map((s) => s.group), spans.map((s) => s.first - 1), spans.map((s) => s.last + 1)]
    )
    const byEnd = new Map<string, string>()
    const byStart = new Map<string, string>()
    for (const r of rows) {
      byEnd.set(`${r.group_id}:${r.end_id}`, r.status)
      byStart.set(`${r.group_id}:${r.start_id}`, r.status)
    }
    const lows = new Map((await db.from('groups').whereIn('id', groups).select('id', 'low')).map((g) => [Number(g.id), g.low === null ? null : Number(g.low)]))
    return (list: (typeof lists)[number]) => {
      const { group, first, last } = span(list)
      const before = byEnd.get(`${group}:${first - 1}`)
      const low = lows.get(group)
      const beforeDone = before !== undefined ? before === 'done' : low === null || low === undefined || first <= low
      return beforeDone && byStart.get(`${group}:${last + 1}`) === 'done'
    }
  }
}

export default new FragmentMerger()