usenet-indexer/app/services/batch_store.ts

import db from '@adonisjs/lucid/services/db'
import type { TransactionClientContract } from '@adonisjs/lucid/types/database'
import type { EdgeBatch, ProbeCandidate, 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).
   * Releases are sorted by key, so concurrent writers lock them in the same order.
   */
  async writeReleases(
    trx: TransactionClientContract,
    rangeId: number,
    batchFrom: number,
    releases: ReleaseChunk[],
    candidates: ProbeCandidate[]
  ) {
    const out = { newReleases: 0, candidates: 0, ids: new Map<string, number>() }
    const ids = out.ids
    // Releases, sorted by key so concurrent batches lock them in the same order.
    const sorted = [...releases].sort((a, b) => a.key.localeCompare(b.key))
    for (let i = 0; i < sorted.length; i += 1000) {
      const chunk = sorted.slice(i, i + 1000)
      const { rows } = await trx.rawQuery(
        `INSERT INTO releases (release_key, name, source, poster, groups, posted_at, file_total, files, segments, bytes, chunks, first_article, last_article, readable, has_par2, complete)
           SELECT v.key, v.name, v.source, v.poster, v.groups, to_timestamp(v.posted_at / 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(key text, name text, source text, poster text, groups jsonb, posted_at bigint, file_total int, files int, segments bigint, bytes bigint, first_article bigint, last_article bigint, readable boolean, has_par2 boolean)
           ON CONFLICT (release_key) DO UPDATE SET
             files = releases.files + excluded.files,
             segments = releases.segments + excluded.segments,
             bytes = releases.bytes + excluded.bytes,
             chunks = releases.chunks + 1,
             posted_at = least(releases.posted_at, excluded.posted_at),
             file_total = coalesce(releases.file_total, excluded.file_total),
             first_article = least(releases.first_article, excluded.first_article),
             last_article = greatest(releases.last_article, excluded.last_article),
             readable = releases.readable OR excluded.readable,
             has_par2 = releases.has_par2 OR excluded.has_par2,
             -- As it stands (the counter's files all found); settling makes it final.
             complete = CASE WHEN coalesce(releases.file_total, excluded.file_total) IS NULL THEN NULL
                             ELSE releases.files + excluded.files >= coalesce(releases.file_total, excluded.file_total) END,
             last_seen_at = now()
           RETURNING id, release_key, (xmax = 0) AS inserted`,
        [
          JSON.stringify(
            chunk.map((r) => ({
              key: r.key,
              name: r.name,
              source: r.source,
              poster: r.poster,
              groups: r.groups,
              posted_at: 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,
            }))
          ),
        ]
      )
      for (const row of rows) {
        ids.set(row.release_key, Number(row.id))
        if (row.inserted) out.newReleases++
      }
    }

    // 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(r.key)!,
        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(c.releaseKey))
    for (let i = 0; i < kept.length; i += 2000) {
      const rows = kept
        .slice(i, i + 2000)
        .map((c) => ({
          release_id: ids.get(c.releaseKey)!,
          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()