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
}

/** File keys per round. */
const BATCH = 2000

/**
 * 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 {
  async run(): Promise<MergeResult> {
    const result: MergeResult = { files: 0, completed: 0, incomplete: 0, waiting: 0 }
    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' LIMIT ?`,
      [BATCH]
    )
    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 }[] = []
    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') && (await this.neighboursDone(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
  }

  /**
   * Whether the ranges right before and after these fragments' ranges (where
   * a missing part could be) are done. One that isn't planned yet counts as
   * not done (newer articles come later), except before the group's oldest article.
   */
  private async neighboursDone(list: { group_id: number; start_id: number | string; end_id: number | string }[]) {
    const groupId = list[0].group_id
    const first = Math.min(...list.map((r) => Number(r.start_id)))
    const last = Math.max(...list.map((r) => Number(r.end_id)))
    const [before, after, group] = await Promise.all([
      db.from('ranges').where('group_id', groupId).where('end_id', first - 1).select('status').first(),
      db.from('ranges').where('group_id', groupId).where('start_id', last + 1).select('status').first(),
      db.from('groups').where('id', groupId).select('low').first(),
    ])
    const beforeDone = before ? before.status === 'done' : group?.low === null || first <= Number(group?.low ?? 0)
    return beforeDone && after?.status === 'done'
  }
}

export default new FragmentMerger()