usenet-indexer/app/services/batch_store.ts

import db from '@adonisjs/lucid/services/db'
import type { TransactionClientContract } from '@adonisjs/lucid/types/database'
import { UPLOAD_GAP_MS, type EdgeBatch, type ProbeCandidate, type ReleaseChunk } from '#lib/edge/index'
import ChunkPartitions from '#services/chunk_partitions'

export class StaleBatch extends Error {}

export interface StoredBatch {
  releases: number
  newReleases: number
  candidates: number
  fragments: number
}

/**
 * Stores one batch of a leased range (REBUILD.md, parts 3–7), in one
 * transaction: release aggregates in the catalog, one chunk per release,
 * probe candidates, rule counts and samples, edge fragments, and the
 * range's resume point. A batch that doesn't start at the resume point was
 * sent twice or too late, and is refused (StaleBatch), so the sums stay right.
 */
class BatchStore {
  async store(rangeId: number, nodeId: number, batch: EdgeBatch): Promise<StoredBatch> {
    await ChunkPartitions.ensure()
    return db.transaction(async (trx) => {
      const range = await trx.from('ranges').where('id', rangeId).forUpdate().first()
      if (!range || range.status !== 'leased' || Number(range.node_id) !== nodeId)
        throw new StaleBatch(`Range ${rangeId} is not leased to this node`)
      if (Number(range.next_id) !== batch.from)
        throw new StaleBatch(
          `Range ${rangeId} is at ${range.next_id}, the batch starts at ${batch.from}`
        )
      const result: StoredBatch = {
        releases: batch.releases.length,
        newReleases: 0,
        candidates: 0,
        fragments: batch.fragments.length,
      }

      const written = await this.writeReleases(
        trx,
        rangeId,
        batch.from,
        batch.releases,
        batch.candidates
      )
      result.newReleases = written.newReleases
      result.candidates = written.candidates

      if (batch.counts.length > 0) {
        await trx.rawQuery(
          `INSERT INTO rule_counts (range_id, rule, key, articles, files)
           SELECT ?, v.rule, v.key, v.articles, v.files FROM jsonb_to_recordset(?::jsonb) AS v(rule text, key text, articles bigint, files bigint)
           ON CONFLICT (range_id, rule, key) DO UPDATE SET articles = rule_counts.articles + excluded.articles, files = rule_counts.files + excluded.files`,
          [rangeId, JSON.stringify(batch.counts)]
        )
      }

      for (const s of batch.samples) {
        await trx.insertQuery().table('rule_samples').insert({
          range_id: rangeId,
          rule: s.rule,
          rule_key: s.ruleKey,
          release_key: s.key,
          name: s.name,
          files: s.files,
          segments: s.segments,
          payload: s.payload,
        })
      }

      for (let i = 0; i < batch.fragments.length; i += 2000) {
        const rows = batch.fragments.slice(i, i + 2000).map((f) => ({
          range_id: rangeId,
          file_key: f.fileKey,
          article_number: Math.min(...f.segments.map((s) => s[3])),
          segments: f.segments.length,
          fragment: JSON.stringify(f),
        }))
        await trx.insertQuery().table('fragments').multiInsert(rows)
      }

      await trx
        .from('ranges')
        .where('id', rangeId)
        .update({
          next_id: batch.next,
          headers: db.raw('headers + ?', [batch.headers]),
          heartbeat_at: new Date(),
        })
      return result
    })
  }

  /**
   * Release aggregates in the catalog, one write-once chunk per release, and
   * probe candidates (a batch's, or files the fragment merger completed).
   * A chunk joins the release of its key whose articles the range's group
   * owns and whose posting window is within UPLOAD_GAP_MS of its own, else
   * starts one: re-uploads are releases of their own (decision d). Keys are
   * locked in order, so concurrent writers don't deadlock or both start one.
   */
  async writeReleases(
    trx: TransactionClientContract,
    rangeId: number,
    batchFrom: number,
    releases: ReleaseChunk[],
    candidates: ProbeCandidate[]
  ) {
    const out = { newReleases: 0, candidates: 0 }
    const ids = new Map<string, number>()
    const slot = (key: string, upload: number) => `${key}\n${upload}`
    const { group_id: groupId } = await trx.from('ranges').where('id', rangeId).select('group_id').firstOrFail()
    const sorted = [...releases].sort((a, b) => a.key.localeCompare(b.key) || a.upload - b.upload)
    for (let i = 0; i < sorted.length; i += 1000) {
      const part = sorted.slice(i, i + 1000)
      const keys = [...new Set(part.map((r) => r.key))]
      await trx.rawQuery(`SELECT pg_advisory_xact_lock(hashtextextended(k, 0)) FROM unnest(?::text[]) AS k ORDER BY k`, [keys])
      const windows = part.map((r, idx) => ({ idx, key: r.key, first: r.postedAt, last: r.postedLast ?? r.postedAt }))
      const { rows: matches } = await trx.rawQuery(
        `SELECT DISTINCT ON (v.idx) v.idx, r.id
         FROM jsonb_to_recordset(?::jsonb) AS v(idx int, key text, first bigint, last bigint)
         JOIN releases r ON r.identity_key = v.key AND r.group_id = ?
         WHERE v.first IS NULL OR r.posted_at IS NULL
            OR (r.posted_at <= to_timestamp((v.last + ?) / 1000.0) AND coalesce(r.last_posted_at, r.posted_at) >= to_timestamp((v.first - ?) / 1000.0))
         ORDER BY v.idx, r.last_posted_at DESC NULLS LAST, r.id DESC`,
        [JSON.stringify(windows), groupId, UPLOAD_GAP_MS, UPLOAD_GAP_MS]
      )
      const joined = new Map<number, number>(matches.map((m: { idx: number; id: string }) => [m.idx, Number(m.id)]))
      const row = (r: ReleaseChunk) => ({
        name: r.name,
        source: r.source,
        poster: r.poster,
        groups: r.groups,
        first: r.postedAt,
        last: r.postedLast ?? r.postedAt,
        file_total: r.fileTotal,
        files: r.files,
        segments: r.segments,
        bytes: r.bytes,
        first_article: r.firstArticle,
        last_article: r.lastArticle,
        readable: r.readable,
        has_par2: r.hasPar2,
      })
      const updates = part.flatMap((r, idx) => (joined.has(idx) ? [{ id: joined.get(idx)!, ...row(r) }] : []))
      if (updates.length > 0) {
        await trx.rawQuery(
          `WITH v AS (SELECT * FROM jsonb_to_recordset(?::jsonb) AS v(id bigint, first bigint, last bigint, file_total int, files int, segments bigint, bytes bigint, first_article bigint, last_article bigint, readable boolean, has_par2 boolean)),
                a AS (SELECT id, sum(files) AS files, sum(segments) AS segments, sum(bytes) AS bytes, count(*) AS chunks, min(first) AS first, max(last) AS last,
                             max(file_total) AS file_total, min(first_article) AS first_article, max(last_article) AS last_article, bool_or(readable) AS readable, bool_or(has_par2) AS has_par2
                      FROM v GROUP BY id)
           UPDATE releases r SET
             files = r.files + a.files,
             segments = r.segments + a.segments,
             bytes = r.bytes + a.bytes,
             chunks = r.chunks + a.chunks,
             posted_at = least(r.posted_at, to_timestamp(a.first / 1000.0)),
             last_posted_at = greatest(r.last_posted_at, to_timestamp(a.last / 1000.0)),
             file_total = coalesce(r.file_total, a.file_total),
             first_article = least(r.first_article, a.first_article),
             last_article = greatest(r.last_article, a.last_article),
             readable = r.readable OR a.readable,
             has_par2 = r.has_par2 OR a.has_par2,
             -- As it stands (the counter's files all found); settling makes it final.
             complete = CASE WHEN coalesce(r.file_total, a.file_total) IS NULL THEN NULL ELSE r.files + a.files >= coalesce(r.file_total, a.file_total) END,
             last_seen_at = now()
           FROM a WHERE r.id = a.id`,
          [JSON.stringify(updates)]
        )
        part.forEach((r, idx) => joined.has(idx) && ids.set(slot(r.key, r.upload), joined.get(idx)!))
      }
      const fresh = part.filter((_, idx) => !joined.has(idx))
      if (fresh.length > 0) {
        // The first upload of a key keeps the key; later ones are key:2, key:3, …
        const { rows: counts } = await trx.rawQuery(`SELECT identity_key, count(*) AS n FROM releases WHERE identity_key = ANY(?) GROUP BY 1`, [[...new Set(fresh.map((r) => r.key))]])
        const taken = new Map<string, number>(counts.map((c: { identity_key: string; n: string }) => [c.identity_key, Number(c.n)]))
        const inserts = fresh.map((r) => {
          const n = (taken.get(r.key) ?? 0) + 1
          taken.set(r.key, n)
          return { release_key: n === 1 ? r.key : `${r.key}:${n}`, key: r.key, upload: r.upload, ...row(r) }
        })
        const { rows } = await trx.rawQuery(
          `INSERT INTO releases (release_key, identity_key, group_id, name, source, poster, groups, posted_at, last_posted_at, file_total, files, segments, bytes, chunks, first_article, last_article, readable, has_par2, complete)
           SELECT v.release_key, v.key, ?, v.name, v.source, v.poster, v.groups, to_timestamp(v.first / 1000.0), to_timestamp(v.last / 1000.0), v.file_total, v.files, v.segments, v.bytes, 1,
                  v.first_article, v.last_article, v.readable, v.has_par2, CASE WHEN v.file_total IS NULL THEN NULL ELSE v.files >= v.file_total END
           FROM jsonb_to_recordset(?::jsonb) AS v(release_key text, key text, name text, source text, poster text, groups jsonb, first bigint, last bigint, file_total int, files int, segments bigint, bytes bigint, first_article bigint, last_article bigint, readable boolean, has_par2 boolean)
           RETURNING id, release_key`,
          [groupId, JSON.stringify(inserts)]
        )
        const byReleaseKey = new Map<string, number>(rows.map((x: { id: string; release_key: string }) => [x.release_key, Number(x.id)]))
        for (const r of inserts) ids.set(slot(r.key, r.upload), byReleaseKey.get(r.release_key)!)
        out.newReleases += rows.length
      }
    }

    // One write-once chunk per release.
    for (let i = 0; i < sorted.length; i += 500) {
      const rows = sorted.slice(i, i + 500).map((r) => ({
        release_id: ids.get(slot(r.key, r.upload))!,
        range_id: rangeId,
        batch_from: batchFrom,
        files: r.files,
        segments: r.segments,
        payload: r.payload,
      }))
      if (rows.length > 0) await trx.insertQuery().table('release_chunks').multiInsert(rows)
    }

    const kept = candidates.filter((c) => ids.has(slot(c.releaseKey, c.upload ?? 0)))
    for (let i = 0; i < kept.length; i += 2000) {
      const rows = kept
        .slice(i, i + 2000)
        .map((c) => ({
          release_id: ids.get(slot(c.releaseKey, c.upload ?? 0))!,
          kind: c.kind,
          message_id: c.messageId,
          file: c.file,
        }))
      const added = await trx
        .knexQuery()
        .table('probe_candidates')
        .insert(rows)
        .onConflict(['release_id', 'kind', 'message_id'])
        .ignore()
        .returning('id')
      out.candidates += added.length
    }

    return out
  }
}

export default new BatchStore()