usenet-indexer/app/services/StagingService.ts

import { createHash } from 'node:crypto'
import db from '@adonisjs/lucid/services/db'
import indexerConfig from '#config/indexer'
import IngestService, { type IngestSegment } from '#services/IngestService'
import { counterFileKey } from '#lib/subject/index'
import { fileStem } from '#lib/release/index'

/** Rows per statement (a key and a JSON array each). */
const CHUNK = 5000

/**
 * Staging of lone header segments (CONCEPT.md D35). Uploads that give every
 * article its own random subject turned each article into a binary, a part
 * and a release, which GC deleted two hours later: ~3.9M binaries an hour
 * against ~123k real files. A segment whose file is new and alone in its
 * fetch batch now waits in `staged_segments` (unlogged, one index) until a
 * second segment of that file arrives; only then do both go into the real
 * tables. GC drops what waited longer than STAGING_TTL_SECONDS (30 min).
 *
 * Counter grouping (D36): those random-subject uploads still keep poster,
 * file counter `[i/M]` and part count per file. Staged segments also carry
 * that key (`counter_key`); once two staged segments from different subjects
 * share it, they become one file under that key, and later segments with it
 * join that file directly.
 */
class StagingService {
  enabled(groupName: string): boolean {
    const groups = indexerConfig.stagingGroups
    return groups.includes('*') || groups.includes(groupName)
  }

  /** Whether counter grouping (D36) applies to this group. */
  countersEnabled(groupName: string): boolean {
    const groups = indexerConfig.counterGroupingGroups
    return groups.includes('*') || groups.includes(groupName)
  }

  /**
   * The segments of a header batch that should be ingested now: the batch's
   * own, minus lone segments of new files (staged instead), plus staged
   * segments whose file this batch shows to be real. `staged` counts the
   * segments that went to staging.
   */
  async filter(segments: IngestSegment[], { counters = false } = {}): Promise<{ ingest: IngestSegment[]; staged: number }> {
    const byKey = new Map<string, IngestSegment[]>()
    for (const segment of segments) {
      const list = byKey.get(segment.fileKey)
      if (list) list.push(segment)
      else byKey.set(segment.fileKey, [segment])
    }

    const ingest: IngestSegment[] = []
    const lone: IngestSegment[] = []
    for (const list of byKey.values()) {
      // Single-segment files are complete; two different articles of a file in one batch make it real.
      // (One article twice is a crosspost, LOG.md 2026-10-02.)
      const articles = new Set(list.map((s) => s.messageId)).size
      if (list[0].totalParts <= 1 || articles >= 2 || list[0].weakIdentity || list[0].size != null) ingest.push(...list)
      else lone.push(list[0])
    }

    const counterOf = new Map<IngestSegment, string>()
    if (counters) {
      for (const segment of lone) {
        const key = counterFileKey(segment)
        if (key) counterOf.set(segment, key)
      }
    }

    // Files we already have a binary for (by subject, or by counters) take their segments directly.
    const known = new Set<string>()
    const lookup = [...new Set([...lone.map((s) => s.fileKey), ...counterOf.values()])]
    for (let i = 0; i < lookup.length; i += CHUNK) {
      const rows = await db.from('binaries').whereIn('file_key', lookup.slice(i, i + CHUNK)).select('file_key')
      for (const row of rows) known.add(row.file_key)
    }
    const toStage: IngestSegment[] = []
    const knownCounters = new Set<string>()
    for (const segment of lone) {
      const counter = counterOf.get(segment)
      if (known.has(segment.fileKey)) {
        ingest.push(segment)
      } else if (counter && known.has(counter)) {
        ingest.push(asCounterFile(segment, counter))
        knownCounters.add(counter)
      } else {
        toStage.push(segment)
      }
    }

    // Stage the rest; a file that now has two staged segments comes out again.
    const promoted: string[] = []
    // Key order, so concurrent jobs lock rows in the same order.
    toStage.sort((a, b) => a.fileKey.localeCompare(b.fileKey))
    for (let i = 0; i < toStage.length; i += CHUNK) {
      const rows = toStage.slice(i, i + CHUNK).map((s) => ({
        file_key: s.fileKey,
        counter_key: counterOf.get(s) ?? null,
        segments: [serialize(s)],
        staged_at: new Date(),
      }))
      const { rows: result } = await retryDeadlock(() => db.rawQuery(
        `INSERT INTO staged_segments (file_key, counter_key, segments, staged_at)
         SELECT file_key, counter_key, segments, staged_at
         FROM jsonb_to_recordset(?::jsonb) AS v(file_key text, counter_key text, segments jsonb, staged_at timestamptz)
         ON CONFLICT (file_key) DO UPDATE SET segments = CASE
           -- The same article again (a crosspost, or a re-fetched range) is no second segment.
           WHEN EXISTS (SELECT 1 FROM jsonb_array_elements(staged_segments.segments) e
                        WHERE e->>'messageId' = EXCLUDED.segments->0->>'messageId') THEN staged_segments.segments
           ELSE staged_segments.segments || EXCLUDED.segments END
         RETURNING file_key, jsonb_array_length(segments) AS n`,
        [JSON.stringify(rows)]
      ))
      for (const row of result) if (Number(row.n) >= 2) promoted.push(row.file_key)
    }

    // Counter files: two staged rows (two subjects) with one counter key make a file.
    const stagedCounters = [...new Set(toStage.map((s) => counterOf.get(s)).filter((k): k is string => Boolean(k)))]
    const readyCounters: string[] = []
    for (let i = 0; i < stagedCounters.length; i += CHUNK) {
      const { rows } = await db.rawQuery(
        `SELECT counter_key FROM staged_segments WHERE counter_key = ANY(?) GROUP BY counter_key HAVING count(*) >= 2`,
        [stagedCounters.slice(i, i + CHUNK)]
      )
      for (const row of rows) readyCounters.push(row.counter_key)
    }

    // Promoted files, staged segments of files this batch ingests anyway, and ready counter files.
    const fileKeys = [...new Set([...promoted, ...ingest.map((s) => s.fileKey)])]
    const counterKeys = [...new Set([...readyCounters, ...knownCounters])]
    const seen = new Set(ingest.map((s) => s.messageId))
    for (const segment of await this.take(fileKeys, counterKeys)) {
      if (!seen.has(segment.messageId)) ingest.push(segment)
      seen.add(segment.messageId)
    }
    const staged = toStage.filter((s) => !seen.has(s.messageId)).length
    return { ingest, staged }
  }

  /**
   * GC: drops staged segments that waited longer than `maxAgeSeconds`, then
   * ingests those whose file got a binary meanwhile (two fetch jobs racing).
   * Expiry comes first so the table stays bounded whatever the budget; the
   * adoption starts from recently seen binaries (index on last_seen_at,
   * which moves in 10-minute steps) rather than joining every staged row.
   */
  async sweep(maxAgeSeconds: number, budget: number): Promise<{ adopted: number; expired: number }> {
    let adopted = 0
    let expired = 0
    while (Date.now() < budget) {
      const result = await db.rawQuery(
        `DELETE FROM staged_segments WHERE ctid IN (
           SELECT ctid FROM staged_segments WHERE staged_at < now() - make_interval(secs => ?) LIMIT 20000)`,
        [maxAgeSeconds]
      )
      expired += result.rowCount
      if (result.rowCount < 20000) break
    }
    let lastId = 0
    while (Date.now() < budget) {
      const recent = await db
        .from('binaries')
        .where('last_seen_at', '>', new Date(Date.now() - (maxAgeSeconds + 15 * 60) * 1000))
        .where('id', '>', lastId)
        .orderBy('id')
        .limit(CHUNK)
        .select('id', 'file_key')
      if (recent.length === 0) break
      lastId = Number(recent[recent.length - 1].id)
      const keys = recent.map((b) => b.file_key as string)
      // A binary's key is a subject key or a counter key; staged rows are looked up by both.
      const segments = await this.take(keys, keys)
      if (segments.length > 0) await IngestService.ingest(segments)
      adopted += segments.length
    }
    return { adopted, expired }
  }

  /**
   * Removes and returns the staged segments of these files: by subject key,
   * or by counter key (returned as segments of the counter file).
   */
  private async take(fileKeys: string[], counterKeys: string[] = []): Promise<IngestSegment[]> {
    const taken: IngestSegment[] = []
    for (let i = 0; i < fileKeys.length; i += CHUNK) {
      const { rows } = await retryDeadlock(() =>
        db.rawQuery(`DELETE FROM staged_segments WHERE file_key = ANY(?) RETURNING segments`, [fileKeys.slice(i, i + CHUNK)])
      )
      for (const row of rows) taken.push(...(row.segments as any[]).map(deserialize))
    }
    for (let i = 0; i < counterKeys.length; i += CHUNK) {
      const { rows } = await retryDeadlock(() =>
        db.rawQuery(`DELETE FROM staged_segments WHERE counter_key = ANY(?) RETURNING counter_key, segments`, [counterKeys.slice(i, i + CHUNK)])
      )
      for (const row of rows) taken.push(...(row.segments as any[]).map((s) => asCounterFile(deserialize(s), row.counter_key)))
    }
    return taken
  }
}

/**
 * A segment as part of its counter file (D36): the counter key is its file
 * key, it is a weak identity (a repeated part number from another upload
 * marks the binary conflicted), and its release is a weak one of its own,
 * which the upload's par2 set can gather later.
 */
function asCounterFile(segment: IngestSegment, counterKey: string): IngestSegment {
  return {
    ...segment,
    fileKey: counterKey,
    weakIdentity: true,
    release: {
      key: createHash('sha1').update(`counter-release\n${counterKey}`).digest('hex'),
      name: fileStem(segment.name),
      source: 'weak-stem',
    },
  }
}

/**
 * Runs a staging statement again when it lost a deadlock: concurrent fetch
 * jobs insert by file key and delete by counter key, so their row locks can
 * cross. Postgres rolled the statement back, and each one is safe to repeat.
 */
async function retryDeadlock<T>(statement: () => Promise<T>, attempts = 5): Promise<T> {
  for (let attempt = 1; ; attempt++) {
    try {
      return await statement()
    } catch (error: any) {
      if (error?.code !== '40P01' || attempt >= attempts) throw error
      await new Promise((resolve) => setTimeout(resolve, 50 * attempt + Math.random() * 100))
    }
  }
}

function serialize(segment: IngestSegment) {
  return {
    ...segment,
    postedAt: segment.postedAt?.toISOString(),
    articleNumber: segment.articleNumber?.toString(),
  }
}

function deserialize(stored: any): IngestSegment {
  return {
    ...stored,
    postedAt: stored.postedAt ? new Date(stored.postedAt) : undefined,
    articleNumber: stored.articleNumber ?? undefined,
  }
}

export default new StagingService()