Harden traffic history worker lifecycle and query performance

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

@@ -0,0 +1,142 @@
+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(); }
+});