MikrotikManager/backend/src/services/traffic-flow-ingest.ts

import { deleteTrafficFlowData, executeTrafficFlowPurge } from "./traffic-flow-purge-execution.js"
import { withTrafficFlowPurgeLock } from "../db/operation-lock.js"
import { Worker } from "node:worker_threads"
import { waitForCollectorShutdown } from "./traffic-flow-worker-shutdown.js"
import { db, dbAll, pool } from "../db/index.js"
import { dropExpiredPartitions } from "../db/partitions.js"
import { servers } from "../db/schema.js"
import type { FlowPurgeDto, FlowStatsDto, FlowTalkerDto } from "@mmapp/contracts/traffic-flow"
import { protoName, type ParsedFlowInput } from "./traffic-flow-parse.js"
import type { CollectorHeartbeat, ExporterMapPayload, MainToWorker, WorkerToMain } from "./traffic-flow-collector-ipc.js"
import {
  applyRingSnapshot,
  configureEngine,
  flushPending,
  getEngineStats,
  getRingMbps as engineGetRingMbps,
  ingestParsedFlowsForServerForTests as engineIngestForServer,
  lastFlushUsedTransactionForTests as engineLastFlushTx,
  listLiveFlowRows as engineListLive,
  peekPendingFlows,
  queueParsedFlows,
  resetEngineForTests,
  resolveServerId,
  rollFlowRings,
  setExporterResolveCtx,
  type PendingFlowRow,
} from "./traffic-flow-engine.js"
import {
  getTrafficFlowSettingsRow,
  listHostPeers,
  invalidateTrafficFlowSettingsCache,
} from "./traffic-flow-settings.js"
import { applicationName } from "./traffic-flow-apps.js"
import { resolveIfaceName } from "./traffic-flow-ifaces.js"
import { canonicalIp } from "./traffic-flow-ip.js"
import { getServerCatalog } from "./traffic-flow-topology.js"

export type { PendingFlowRow }

export interface FlowListenerState {
  bound: boolean
  address: string | null
}

export interface FlowWorkerHealth {
  alive: boolean
  bound: boolean
  pendingSize: number
  dropped: number
  packetsReceived: number
}

let worker: Worker | null = null
let restartTimer: ReturnType<typeof setTimeout> | null = null
let restartAttempts = 0
let lastHeartbeat: CollectorHeartbeat | null = null
let state: FlowListenerState = { bound: false, address: null }
let wantListen = false
let databaseStop: Promise<void> | null = null

export function trafficFlowWorkerFileUrl(): URL {
  const ts = new URL(import.meta.url).pathname.endsWith(".ts")
  return new URL(
    ts ? "./traffic-flow-worker-bootstrap.mjs" : "./traffic-flow-collector-worker.js",
    import.meta.url,
  )
}

export async function buildExporterMapPayload(): Promise<ExporterMapPayload> {
  const settings = await getTrafficFlowSettingsRow()
  const rows = await db.select({
    id: servers.id,
    host: servers.host,
    mgmtTunnelIp: servers.mgmtTunnelIp,
  }).from(servers)
  const byTunnelIp: Array<[string, number]> = []
  const hostIps: Array<[string, number]> = []
  for (const row of rows) {
    if (row.mgmtTunnelIp) byTunnelIp.push([row.mgmtTunnelIp, row.id])
    if (/^\d{1,3}(?:\.\d{1,3}){3}$/.test(row.host)) hostIps.push([row.host, row.id])
  }
  return {
    overlayPrefix: settings.prefix,
    byTunnelIp,
    peers: await listHostPeers(),
    hostIps,
  }
}

async function applyExporterCtxFromDb(): Promise<void> {
  const payload = await buildExporterMapPayload()
  setExporterResolveCtx({
    overlayPrefix: payload.overlayPrefix,
    byTunnelIp: new Map(payload.byTunnelIp),
    peers: payload.peers,
    hostIps: new Map(payload.hostIps),
  })
}

function postToWorker(msg: MainToWorker): void {
  worker?.postMessage(msg)
}

function handleWorkerMessage(msg: WorkerToMain): void {
  if (msg.type === "heartbeat") {
    lastHeartbeat = msg.payload
    state = { bound: msg.payload.bound, address: msg.payload.address }
    applyRingSnapshot(msg.payload.rings)
    restartAttempts = 0
    return
  }
  if (msg.type === "error") {
    lastHeartbeat = lastHeartbeat
      ? { ...lastHeartbeat, lastError: msg.payload.message, workerAlive: true }
      : null
  }
}

async function spawnWorker(): Promise<void> {
  if (databaseStop) return
  stopWorkerProcess()
  const settings = await getTrafficFlowSettingsRow()
  if (databaseStop) return
  configureEngine({ topN: settings.topN, retentionHours: settings.retentionHours })
  await applyExporterCtxFromDb()
  const exporterMap = await buildExporterMapPayload()
  if (databaseStop) return
  const w = new Worker(trafficFlowWorkerFileUrl(), { execArgv: process.execArgv })
  w.on("message", (msg: WorkerToMain) => handleWorkerMessage(msg))
  w.on("error", (err) => {
    state = { bound: false, address: null }
    lastHeartbeat = lastHeartbeat
      ? { ...lastHeartbeat, workerAlive: false, lastError: err.message, bound: false }
      : {
          bound: false,
          address: null,
          packetsReceived: 0,
          lastExporterIp: null,
          lastError: err.message,
          lastDatagramAt: null,
          pendingSize: 0,
          dropped: 0,
          rowsStored: 0,
          workerAlive: false,
          rings: [],
        }
  })
  w.on("exit", (code) => {
    if (worker !== w) return
    worker = null
    state = { bound: false, address: null }
    if (!wantListen) return
    const delay = Math.min(30_000, 1000 * 2 ** restartAttempts)
    restartAttempts += 1
    restartTimer = setTimeout(() => {
      if (wantListen) void spawnWorker()
    }, delay)
    void code
  })
  worker = w
  const host = process.env.FLOW_LISTEN_HOST?.trim() || settings.collectorIp || "127.0.0.1"
  postToWorker({
    type: "start",
    payload: {
      listenHost: host,
      listenPort: settings.flowListenPort,
      topN: settings.topN,
      retentionHours: settings.retentionHours,
      exporterMap,
    },
  })
}

function stopWorkerProcess(): void {
  if (restartTimer) {
    clearTimeout(restartTimer)
    restartTimer = null
  }
  if (worker) {
    try {
      postToWorker({ type: "stop" })
      void worker.terminate()
    } catch {
      /* ignore */
    }
    worker = null
  }
}

async function stopWorkerProcessAsync(): Promise<void> {
  if (restartTimer) {
    clearTimeout(restartTimer)
    restartTimer = null
  }
  if (!worker) return
  const current = worker
  worker = null
  try {
    current.postMessage({ type: "stop" })
    await current.terminate()
  } catch {
    /* ignore */
  }
}

export function reattachFlowSqlite(): void {
  /* no-op: engine uses PostgreSQL pool */
}

export function applyHeartbeatForTests(payload: CollectorHeartbeat): void {
  handleWorkerMessage({ type: "heartbeat", payload })
}

export function simulateWorkerExitForTests(): number {
  worker = null
  state = { bound: false, address: null }
  lastHeartbeat = lastHeartbeat ? { ...lastHeartbeat, workerAlive: false, bound: false } : null
  if (!wantListen) return restartAttempts
  restartAttempts += 1
  return restartAttempts
}

export function setWantListenForTests(value: boolean): void {
  wantListen = value
}

export function getFlowListenerState(): FlowListenerState {
  return state
}

export function getFlowWorkerHealth(): FlowWorkerHealth {
  const hb = lastHeartbeat
  const mem = getEngineStats()
  return {
    alive: Boolean(worker) && (hb?.workerAlive ?? false),
    bound: state.bound,
    pendingSize: hb?.pendingSize ?? mem.pendingSize,
    dropped: hb?.dropped ?? mem.dropped,
    packetsReceived: hb?.packetsReceived ?? mem.packetsReceived,
  }
}

export async function getFlowRuntimeCounters() {
  const hb = lastHeartbeat
  if (hb) {
    return {
      packetsReceived: hb.packetsReceived,
      lastExporterIp: hb.lastExporterIp ?? null,
      lastError: hb.lastError || null,
      lastDatagramAt: hb.lastDatagramAt ?? null,
      dropped: hb.dropped ?? 0,
    }
  }
  const settings = await getTrafficFlowSettingsRow()
  return {
    packetsReceived: settings.packetsReceived,
    lastExporterIp: settings.lastExporterIp ?? null,
    lastError: settings.lastError || null,
    lastDatagramAt: settings.lastDatagramAt ?? null,
    dropped: 0,
  }
}

export async function startTrafficFlowListener() {
  if (databaseStop) throw new Error("Приём NetFlow приостановлен для операции с базой")
  stopTrafficFlowListener()
  const settings = await getTrafficFlowSettingsRow()
  if (databaseStop) throw new Error("Приём NetFlow приостановлен для операции с базой")
  if (!settings.enabled) {
    wantListen = false
    state = { bound: false, address: null }
    return
  }
  wantListen = true
  await spawnWorker()
}

export function stopTrafficFlowListener() {
  wantListen = false
  stopWorkerProcess()
  void flushPending({ force: true }).catch(() => { /* ignore */ })
  state = { bound: false, address: null }
}

export async function stopTrafficFlowListenerForDatabaseOperation(): Promise<void> {
  if (databaseStop) return databaseStop
  wantListen = false
  if (restartTimer) clearTimeout(restartTimer)
  restartTimer = null
  const current = worker
  databaseStop = (async () => {
    if (current) await waitForCollectorShutdown(current)
    await flushPending({ force: true, prune: false })
    state = { bound: false, address: null }
    lastHeartbeat = null
  })()
  // Keep a failed stop latched: restart must not kill a worker whose writes are unconfirmed.
  await databaseStop
}

export async function restartTrafficFlowListenerAfterDatabaseOperation(): Promise<void> {
  if (databaseStop) await databaseStop
  databaseStop = null
  await startTrafficFlowListener()
}

export async function refreshFlowExporterMap(): Promise<void> {
  await applyExporterCtxFromDb()
  postToWorker({ type: "updateExporterMap", payload: await buildExporterMapPayload() })
}

export function getRingMbps(serverId: number, iface = "__all__") {
  return engineGetRingMbps(serverId, iface)
}

function mergeInto(map: Map<string, PendingFlowRow>, row: PendingFlowRow): void {
  const key = `${row.serverId}|${row.bucketAt}|${row.src}|${row.dst}|${row.proto}|${row.srcPort}|${row.dstPort}|${row.inIface}`
  const prev = map.get(key)
  if (prev) {
    prev.bytes += row.bytes
    prev.packets += row.packets
    if (row.outIface && !prev.outIface) prev.outIface = row.outIface
    if (row.nextHop && !prev.nextHop) prev.nextHop = row.nextHop
    if (row.natSrc && !prev.natSrc) prev.natSrc = row.natSrc
    if (row.natDst && !prev.natDst) prev.natDst = row.natDst
    if (row.natSrcPort && !prev.natSrcPort) prev.natSrcPort = row.natSrcPort
    if (row.natDstPort && !prev.natDstPort) prev.natDstPort = row.natDstPort
    if (row.flowStartMs && (!prev.flowStartMs || row.flowStartMs < prev.flowStartMs)) prev.flowStartMs = row.flowStartMs
    if (row.flowEndMs > (prev.flowEndMs ?? 0)) prev.flowEndMs = row.flowEndMs
    return
  }
  map.set(key, { ...row })
}

/** Top-N разговоров на (server_id, bucket_at), как prune persist — не глобальный ORDER BY bytes. */
export function capFlowRowsPerServerBucket<T extends { serverId: number; bucketAt: string; bytes: number }>(
  rows: T[],
  keep: number,
): T[] {
  const cap = Math.max(20, keep)
  const groups = new Map<string, T[]>()
  for (const row of rows) {
    const k = `${row.serverId}\0${row.bucketAt}`
    const list = groups.get(k)
    if (list) list.push(row)
    else groups.set(k, [row])
  }
  const out: T[] = []
  for (const list of groups.values()) {
    list.sort((a, b) => b.bytes - a.bytes)
    out.push(...list.slice(0, cap))
  }
  return out
}

function isoBucketAt(v: unknown): string {
  if (v instanceof Date) return v.toISOString()
  return String(v ?? "")
}

export async function listLiveFlowRows(sinceIso: string): Promise<PendingFlowRow[]> {
  if (worker && lastHeartbeat?.workerAlive) {
    return await listStoredFlowRows(sinceIso)
  }
  return engineListLive(sinceIso)
}

interface StoredBucketRow {
  server_id: number
  bucket_at: string | Date
  src: string
  dst: string
  proto: number
  src_port: number
  dst_port: number
  bytes: number
  packets: number
  in_iface: string
  out_iface: string | null
  next_hop: string | null
  flow_start_ms: number | null
  flow_end_ms: number | null
  nat_src: string | null
  nat_dst: string | null
  nat_src_port: number | null
  nat_dst_port: number | null
}

export async function listStoredFlowRows(sinceIso: string): Promise<PendingFlowRow[]> {
  const settings = await getTrafficFlowSettingsRow()
  const keep = Math.max(20, settings.topN)
  const stored = await dbAll<StoredBucketRow>(`
    SELECT
      server_id, bucket_at, COALESCE(host(src), '') AS src, COALESCE(host(dst), '') AS dst, proto,
      src_port, dst_port, bytes, packets, in_iface, out_iface,
      COALESCE(host(next_hop), '') AS next_hop, flow_start_ms, flow_end_ms,
      COALESCE(host(nat_src), '') AS nat_src, COALESCE(host(nat_dst), '') AS nat_dst, nat_src_port, nat_dst_port
    FROM (
      SELECT fb.*,
             ROW_NUMBER() OVER (
               PARTITION BY server_id, bucket_at ORDER BY bytes DESC
             ) AS rn
      FROM flow_buckets fb
      WHERE bucket_at >= $1
    ) ranked
    WHERE rn <= $2
  `, [sinceIso, keep])
  const merged = new Map<string, PendingFlowRow>()
  for (const r of stored) {
    mergeInto(merged, {
      serverId: Number(r.server_id),
      bucketAt: isoBucketAt(r.bucket_at),
      src: canonicalIp(r.src),
      dst: canonicalIp(r.dst),
      proto: Number(r.proto) || 0,
      srcPort: Number(r.src_port) || 0,
      dstPort: Number(r.dst_port) || 0,
      bytes: Number(r.bytes) || 0,
      packets: Number(r.packets) || 0,
      inIface: r.in_iface ?? "",
      outIface: r.out_iface ?? "",
      nextHop: canonicalIp(r.next_hop ?? ""),
      flowStartMs: Number(r.flow_start_ms) || 0,
      flowEndMs: Number(r.flow_end_ms) || 0,
      natSrc: canonicalIp(r.nat_src ?? ""),
      natDst: canonicalIp(r.nat_dst ?? ""),
      natSrcPort: Number(r.nat_src_port) || 0,
      natDstPort: Number(r.nat_dst_port) || 0,
    })
  }
  if (!worker) {
    for (const p of peekPendingFlows()) {
      if (p.bucketAt < sinceIso) continue
      mergeInto(merged, p)
    }
  }
  return capFlowRowsPerServerBucket([...merged.values()], keep)
}

export async function listFlowRowsForWindow(minutes: number): Promise<PendingFlowRow[]> {
  const sinceIso = new Date(Date.now() - minutes * 60_000).toISOString()
  // History remains available while the collector starts or is stopped; pending rows are merged below.
  return await listStoredFlowRows(sinceIso)
}

export async function listFlowTalkers(minutes = 5): Promise<FlowStatsDto> {
  const settings = await getTrafficFlowSettingsRow()
  const runtime = await getFlowRuntimeCounters()
  const rows = await listFlowRowsForWindow(minutes)
  const catalog = await getServerCatalog()
  const nameById = new Map([...catalog.byId].map(([id, s]) => [id, s.name]))
  const agg = new Map<string, FlowTalkerDto & { rawBytes: number }>()
  const protoBytes = new Map<number, number>()
  const srcs = new Set<string>()
  const dsts = new Set<string>()
  const exporters = new Set<number>()
  let totalBytes = 0
  for (const r of rows) {
    const resolved = resolveIfaceName(r.serverId, r.inIface)
    const key = `${r.serverId}|${r.src}|${r.dst}|${r.proto}|${r.srcPort}|${r.dstPort}|${r.inIface}`
    const prev = agg.get(key)
    const bytes = r.bytes
    totalBytes += bytes
    srcs.add(r.src)
    dsts.add(r.dst)
    exporters.add(r.serverId)
    protoBytes.set(r.proto, (protoBytes.get(r.proto) ?? 0) + bytes)
    if (prev) {
      prev.rawBytes += bytes
      prev.bytes += bytes
      prev.packets += r.packets
    } else {
      agg.set(key, {
        serverId: String(r.serverId),
        serverName: nameById.get(r.serverId) ?? String(r.serverId),
        src: r.src,
        dst: r.dst,
        proto: r.proto,
        protoName: protoName(r.proto),
        srcPort: r.srcPort,
        dstPort: r.dstPort,
        bytes,
        packets: r.packets,
        bps: 0,
        inIface: resolved.name,
        inIfaceIndex: resolved.index,
        application: applicationName(r.proto, r.dstPort, r.srcPort),
        rawBytes: bytes,
      })
    }
  }
  const windowSec = Math.max(60, minutes * 60)
  const talkers = [...agg.values()]
    .map((t) => ({ ...t, bps: (t.rawBytes * 8) / windowSec }))
    .sort((a, b) => b.bytes - a.bytes)
    .slice(0, settings.topN)
    .map(({ rawBytes: _raw, ...rest }) => rest)
  let topProto = "—"
  let topProtoBytes = 0
  for (const [p, b] of protoBytes) {
    if (b > topProtoBytes) {
      topProtoBytes = b
      topProto = protoName(p)
    }
  }
  return {
    exportersOnline: exporters.size,
    bytesPerMin: minutes > 0 ? totalBytes / minutes : totalBytes,
    uniqueSrc: srcs.size,
    uniqueDst: dsts.size,
    topProto,
    talkers,
    lastExporterIp: runtime.lastExporterIp,
    lastError: runtime.lastError,
    packetsReceived: runtime.packetsReceived,
    lastDatagramAt: runtime.lastDatagramAt,
    listenerBound: state.bound,
    listenerAddress: state.address,
  }
}

export async function ingestParsedFlowsForTests(exporterIp: string, flows: ParsedFlowInput[]) {
  await applyExporterCtxFromDb()
  const serverId = resolveServerId(exporterIp)
  if (serverId == null) return
  queueParsedFlows(serverId, flows)
  rollFlowRings()
  await flushPending({ force: true })
}

export function ingestParsedFlowsForServerForTests(serverId: number, flows: ParsedFlowInput[]) {
  engineIngestForServer(serverId, flows)
}

export function resetFlowRingsForTests() {
  resetEngineForTests()
  lastHeartbeat = null
  wantListen = false
  restartAttempts = 0
}

export function lastFlushUsedTransactionForTests(): boolean {
  return engineLastFlushTx()
}

export async function flushPendingForTests(): Promise<void> {
  await flushPending({ force: true })
}

const PURGE_LOCK_KEY = 8_723_103

/** Удаляет данные одним transaction; ключи и пиры не меняются. */
export async function purgeTrafficFlowStore(): Promise<FlowPurgeDto> {
  return withTrafficFlowPurgeLock(() => pool.connect(), PURGE_LOCK_KEY, () => executeTrafficFlowPurge({
    stop: async () => {
      wantListen = false
      await stopWorkerProcessAsync()
      resetEngineForTests()
      lastHeartbeat = null
      state = { bound: false, address: null }
    },
    deleteData: () => deleteTrafficFlowData(() => pool.connect()),
    invalidate: invalidateTrafficFlowSettingsCache,
    maintain: () => dropExpiredPartitions(pool),
    restart: startTrafficFlowListener,
  }))
}

export { peekPendingFlows }
export { setPendingCapForTests } from "./traffic-flow-engine.js"
export { maybeRefreshIfaces, setRefreshIfacesForTests } from "./traffic-flow-ifaces.js"