usenet-indexer/app/services/nntp_pool.ts

import env from '#start/env'
import logger from '@adonisjs/core/services/logger'
import coordinatorConfig from '#config/coordinator'
import { XoverClient } from '#lib/nntp/index'
import type { NntpShare } from '#services/nntp_shares'

/** Connections wanted beyond the most in use at once since the last sync (the planner's, a hedge's). */
const SHARE_MARGIN = 2

/**
 * The coordinator's header connections (REBUILD.md, decision e: workers
 * have no NNTP account). At most NNTP_CONNECTIONS, reused; a connection
 * that failed, or is older than NNTP_CONNECTION_MAX_MINUTES, is closed and
 * replaced on the next use. Sharing the account (share()), at most its
 * share: over it, connections close as they come free.
 */
class NntpPool {
  private idle: XoverClient[] = []
  private busy = new Set<XoverClient>()
  private open = 0
  private waiting: ((client: XoverClient) => void)[] = []
  private openedAt = new Map<XoverClient, number>()
  /** When a connection gets closed once free: its max age ±25%, so they don't all reconnect at once. */
  private retireAt = new Map<XoverClient, number>()
  private limit = coordinatorConfig.nntpConnections
  /** Since the last sync: the most connections in use at once, and whether work waited for one (starved()). */
  private peak = 0
  private short = false

  get size() {
    return this.limit
  }

  /** Work that could use another connection didn't get one (the prefetcher with all its loops running). */
  starved() {
    this.short = true
  }

  /**
   * Tells the share what the pool could use and has open, and takes the share's grant as its size. After work waited:
   * twice its size (history going on gets its maximum in a few syncs; a tail round, done in seconds, doesn't take
   * the prober's connections for one); else the most in use at once since the last sync and a margin.
   */
  async share(share: NntpShare) {
    const want = this.short || this.waiting.length > 0 ? Math.max(this.limit * 2, this.limit + SHARE_MARGIN) : this.peak + SHARE_MARGIN
    this.short = false
    this.peak = this.busy.size
    const granted = await share.sync(want, this.open)
    const was = this.limit
    this.limit = granted
    // Over the share: the idle ones close now, the busy ones once free (release()).
    while (this.open > this.limit && this.idle.length > 0) this.retire(this.idle.pop()!)
    // Under it with work waiting: connections for it.
    while (this.waiting.length > 0 && this.open < this.limit) {
      const next = this.waiting.shift()!
      this.acquire().then(next, () => this.waiting.unshift(next))
    }
    return { was, granted }
  }

  async use<T>(work: (client: XoverClient) => Promise<T>): Promise<T> {
    return this.run(await this.acquire(), work)
  }

  /** As use(), if a connection is free or can be opened now; else null, without waiting. */
  tryUse<T>(work: (client: XoverClient) => Promise<T>): Promise<T> | null {
    if (this.idle.length === 0 && this.open >= this.size) return null
    return this.acquire().then((client) => this.run(client, work))
  }

  /** Connections with a command in flight (the prefetcher's, the planner's). */
  get busyCount() {
    return this.busy.size
  }

  /** How long ago the connection was opened, in ms. */
  ageOf(client: XoverClient): number {
    return Date.now() - (this.openedAt.get(client) ?? Date.now())
  }

  private async run<T>(client: XoverClient, work: (client: XoverClient) => Promise<T>): Promise<T> {
    this.busy.add(client)
    this.peak = Math.max(this.peak, this.busy.size)
    let broken = false
    try {
      return await work(client)
    } catch (error) {
      broken = true
      throw error
    } finally {
      this.busy.delete(client)
      this.release(client, broken)
    }
  }

  private async acquire(): Promise<XoverClient> {
    const ready = this.idle.pop()
    if (ready) return ready
    if (this.open < this.size) {
      this.open++
      try {
        const client = await XoverClient.connect({
          host: env.get('NNTP_HOST')!,
          port: env.get('NNTP_PORT', 563),
          user: env.get('NNTP_USER'),
          password: env.get('NNTP_PASSWORD'),
          secure: env.get('NNTP_SECURE', true),
          compress: true,
          timeoutMs: coordinatorConfig.nntpTimeoutMs,
        })
        this.openedAt.set(client, Date.now())
        const maxMs = coordinatorConfig.nntpConnectionMaxMs
        if (maxMs > 0) this.retireAt.set(client, Date.now() + maxMs * (0.75 + Math.random() / 2))
        return client
      } catch (error) {
        this.open--
        throw error
      }
    }
    return new Promise((resolve) => this.waiting.push(resolve))
  }

  private release(client: XoverClient, broken: boolean) {
    if (this.open > this.limit) {
      this.retire(client)
      return
    }
    if (broken || Date.now() >= (this.retireAt.get(client) ?? Infinity)) {
      this.retire(client)
      // Someone waiting gets a fresh connection.
      const next = this.waiting.shift()
      if (next) this.acquire().then(next, () => this.waiting.unshift(next))
      return
    }
    const next = this.waiting.shift()
    if (next) next(client)
    else this.idle.push(client)
  }

  private retire(client: XoverClient) {
    this.open--
    this.openedAt.delete(client)
    this.retireAt.delete(client)
    try {
      client.close()
    } catch (error: any) {
      logger.debug(`Closing an NNTP connection: ${error.message}`)
    }
  }

  /** Closes every connection, also those in use: their commands fail. */
  private plain: Promise<XoverClient> | null = null
  private plainQueue: Promise<unknown> = Promise.resolve()

  /**
   * One connection without XFEATURE COMPRESS GZIP, opened when first needed,
   * one command at a time: for chunks whose compressed reply the server
   * cuts short (Leases.corrupt). Its replies are deflated here.
   */
  async usePlain<T>(work: (client: XoverClient) => Promise<T>): Promise<T> {
    const run = this.plainQueue.then(async () => {
      this.plain ??= XoverClient.connect({
        host: env.get('NNTP_HOST')!,
        port: env.get('NNTP_PORT', 563),
        user: env.get('NNTP_USER'),
        password: env.get('NNTP_PASSWORD'),
        secure: env.get('NNTP_SECURE', true),
        compress: false,
        timeoutMs: coordinatorConfig.nntpTimeoutMs,
      })
      const client = await this.plain
      try {
        return await work(client)
      } catch (error) {
        this.plain = null
        client.close()
        throw error
      }
    })
    this.plainQueue = run.catch(() => {})
    return run
  }

  async close() {
    void this.plain?.then((c) => c.close()).catch(() => {})
    this.plain = null
    for (const client of [...this.idle.splice(0), ...this.busy]) client.close()
    this.busy.clear()
    this.openedAt.clear()
    this.retireAt.clear()
    this.open = 0
  }
}

export default new NntpPool()