harbor-net/test/server/traffic-history-worker.test.js

import assert from 'node:assert/strict';
import fs from 'node:fs';
import os from 'node:os';
import path from 'node:path';
import { Worker } from 'node:worker_threads';
import test from 'node:test';
import { createTrafficHistoryService } from '../../dist/server/services/trafficHistoryService.js';
import { openTrafficHistoryStore } from '../../dist/server/services/trafficHistoryStore.js';
import { emptyTrafficHistory, parseTrafficHistoryQuery } from '../../dist/shared/trafficHistory.js';

const query = parseTrafficHistoryQuery(new URLSearchParams());
async function fixture(t) {
  const directory = fs.mkdtempSync(path.join(os.tmpdir(), 'harbor-history-queue-'));
  const service = createTrafficHistoryService({ filePath: path.join(directory, 'traffic.sqlite'), source: () => 'live' });
  await service.flush();
  const post = Worker.prototype.postMessage;
  const sent = [];
  const waiting = [];
  const held = [];
  t.mock.method(Worker.prototype, 'postMessage', function (message) {
    sent.push(message.kind);
    if (message.kind !== 'query') return post.call(this, message);
    const job = { worker: this, message };
    if (waiting.length) waiting.shift()(job);
    else held.push(job);
  });
  t.mock.timers.enable({ apis: ['setTimeout'] });
  t.after(async () => {
    t.mock.restoreAll();
    t.mock.timers.reset();
    await service.close();
    fs.rmSync(directory, { recursive: true, force: true });
  });
  return { service, sent, file: path.join(directory, 'traffic.sqlite'),
    next: () => held.length ? Promise.resolve(held.shift()) : new Promise((resolve) => waiting.push(resolve)),
    answer: ({ worker, message }) => worker.emit('message', { id: message.id, result: emptyTrafficHistory(message.query, 'live') }),
  };
}

test('queued reads get their own execution budget and do not reset the collector', async (t) => {
  const f = await fixture(t);
  const results = Array.from({ length: 3 }, () => f.service.query(query));
  const first = await f.next();
  const worker = first.worker;
  assert.equal(f.sent.filter((kind) => kind === 'query').length, 1);
  t.mock.timers.tick(4_000);
  f.answer(first);
  for (let index = 1; index < 3; index++) {
    const next = await f.next();
    assert.equal(next.worker, worker);
    t.mock.timers.tick(4_000);
    f.answer(next);
  }
  assert.deepEqual((await Promise.all(results)).map((result) => result.storage.status), ['ready', 'ready', 'ready']);
});

test('a slow read expires alone, retains its active slot, and lets pending writes go next', async (t) => {
  const f = await fixture(t);
  const firstResult = f.service.query(query);
  const otherResult = f.service.query(query);
  const first = await f.next();
  t.mock.timers.tick(5_000);
  assert.equal((await firstResult).storage.status, 'error');
  const write = f.service.flush();
  assert.equal(f.sent.filter((kind) => kind === 'query').length, 1);
  const before = f.sent.length;
  f.answer(first);
  await write;
  const second = await f.next();
  assert.deepEqual(f.sent.slice(before), ['ingest', 'query']);
  assert.equal(first.worker, second.worker);
  f.answer(second);
  const result = await otherResult;
  assert.equal(result.storage.status, 'ready');
  assert.equal(result.coverage.partial, false);
});

test('a full read queue reserves write capacity and shutdown rejects queued and future reads', async (t) => {
  const f = await fixture(t);
  const results = Array.from({ length: 32 }, () => f.service.query(query));
  const first = await f.next();
  assert.equal((await f.service.query(query)).storage.status, 'error');
  const write = f.service.flush();
  const closing = f.service.close();
  f.answer(first);
  await write;
  await closing;
  const snapshots = await Promise.all(results);
  assert.equal(snapshots[0].storage.status, 'ready');
  assert.ok(snapshots.slice(1).every((result) => result.storage.status === 'error'));
  const sent = f.sent.length;
  assert.equal((await f.service.query(query)).storage.status, 'error');
  await f.service.flush();
  assert.equal(f.sent.length, sent);
});

test('hard recovery waits for the previous worker to exit before starting a replacement', async (t) => {
  const f = await fixture(t);
  const firstResult = f.service.query(query);
  const otherResult = f.service.query(query);
  const first = await f.next();
  const terminate = Worker.prototype.terminate;
  let release;
  let stopping = 0;
  const gate = new Promise((resolve) => { release = resolve; });
  t.mock.method(Worker.prototype, 'terminate', function () {
    stopping++;
    return gate.then(() => terminate.call(this));
  });
  t.mock.timers.tick(60_000);
  assert.equal((await firstResult).storage.status, 'error');
  assert.equal(stopping, 1);
  assert.equal(f.sent.filter((kind) => kind === 'query').length, 1);
  release();
  const second = await f.next();
  assert.notEqual(second.worker, first.worker);
  f.answer(second);
  assert.equal((await otherResult).storage.status, 'ready');
});

test('shutdown persists real counters queued behind an occupied read and survives reopen', async (t) => {
  const f = await fixture(t);
  const reading = f.service.query(query);
  const first = await f.next();
  const at = new Date(Date.now() - 120_000).toISOString();
  f.service.enqueue({ epoch: 'shutdown', observedAt: at, reset: false, closedIds: [], connections: [{
    id: 'persist-on-close', startedAt: at, closedAt: at,
    traffic: { uploadBytes: '17', downloadBytes: '29' },
    destination: { domain: 'example.org', ip: '203.0.113.1' },
    source: { ip: '192.0.2.1' }, inbound: { tag: 'tproxy-in' },
    origin: { kind: 'device', id: 'device-1', label: 'Laptop' },
    route: { kind: 'vpn', outbound: 'vpn-one' },
  }] });
  const closing = f.service.close();
  f.answer(first);
  await reading;
  await closing;
  const store = openTrafficHistoryStore(f.file);
  try {
    assert.deepEqual(store.query(query).totals, { uploadBytes: '17', downloadBytes: '29' });
  } finally { store.close(); }
});