Harden traffic history worker lifecycle and query performance

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

@@ -0,0 +1,65 @@
+import assert from 'node:assert/strict';
+import fs from 'node:fs';
+import os from 'node:os';
+import path from 'node:path';
+import { DatabaseSync } from 'node:sqlite';
+import { performance } from 'node:perf_hooks';
+import test from 'node:test';
+import { openTrafficHistoryStore } from '../../dist/server/services/trafficHistoryStore.js';
+import { createTrafficHistoryService } from '../../dist/server/services/trafficHistoryService.js';
+import { parseTrafficHistoryQuery } from '../../dist/shared/trafficHistory.js';
+
+test('large history preserves totals and serves concurrent readers through cleanup and filters', async (t) => {
+  const directory = fs.mkdtempSync(path.join(os.tmpdir(), 'harbor-history-load-'));
+  const file = path.join(directory, 'traffic.sqlite');
+  const at = Math.floor(Date.now() / 60_000) * 60_000;
+  const minutes = process.env.HARBOR_HISTORY_LOAD === '1' ? 3000 : 300;
+  let service;
+  t.after(async () => {
+    await service?.close();
+    fs.rmSync(directory, { recursive: true, force: true });
+  });
+  openTrafficHistoryStore(file).close();
+  const db = new DatabaseSync(file);
+  try {
+    db.exec(`BEGIN;
+      WITH RECURSIVE n(i) AS (VALUES(1) UNION ALL SELECT i+1 FROM n WHERE i<2000)
+      INSERT INTO dimensions SELECT i,'dimension-'||i,'device-'||(i%20),'Device '||(i%20),
+        '192.0.2.1','tproxy-in','service-'||(i%250),'domain-'||(i%250)||'.test',
+        'host-'||i||'.test','203.0.113.'||(i%250),'vpn','vpn-one' FROM n;`);
+    db.prepare(`WITH RECURSIVE minutes(i) AS (VALUES(1) UNION ALL SELECT i+1 FROM minutes WHERE i<?)
+      INSERT INTO buckets SELECT ?-i*60000,60000,d.id,100,200 FROM minutes
+      CROSS JOIN dimensions d WHERE d.id<=1000`).run(minutes, at);
+    db.prepare('INSERT INTO buckets SELECT ?,60000,id,100,200 FROM dimensions WHERE id>1000').run(at - 91 * 86_400_000);
+    db.exec('COMMIT');
+    // Exercise the additive index upgrade on a populated pre-fix database.
+    db.exec('DROP INDEX buckets_dimension_time');
+  } finally { db.close(); }
+
+  service = createTrafficHistoryService({ filePath: file, source: () => 'live' });
+  const query = parseTrafficHistoryQuery(new URLSearchParams(`range=7d&until=${at}`));
+  const start = performance.now();
+  const concurrent = await Promise.all(Array.from({ length: 4 }, () => service.query(query)));
+  const concurrentMs = performance.now() - start;
+  for (const result of concurrent) {
+    assert.equal(result.storage.status, 'ready');
+    assert.equal(result.totals.downloadBytes, String(minutes * 1000 * 200));
+    assert.equal(result.origins.length, 20);
+    assert.equal(result.nextOffset, 100);
+  }
+  const timings = [];
+  for (const [filter, dimensions] of [
+    [{ originId: 'device-1' }, 50],
+    [{ search: 'host-99' }, 11],
+    [{ level: 'ip', service: 'service-1', domain: 'domain-1.test', hostname: 'host-1.test' }, 1],
+    [{ offset: 300 }, 1000],
+  ]) {
+    const begin = performance.now();
+    const result = await service.query({ ...query, ...filter });
+    timings.push({ filter, ms: performance.now() - begin });
+    assert.equal(result.storage.status, 'ready');
+    assert.equal(result.totals.downloadBytes, String(minutes * dimensions * 200));
+    if (filter.offset) assert.equal(result.rows.length, 0);
+  }
+  t.diagnostic(JSON.stringify({ buckets: minutes * 1000, concurrentMs, timings }));
+});

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(); }
+});

test/server/traffic-history.test.js

@@ -297,3 +297,35 @@ test('high-churn sample measures closed-lifecycle storage and epoch reclamation'
   db.close();
   t.diagnostic(JSON.stringify({ closedLifecycles: 100_000, writeMs, queryMs, bytes, reusableBytes }));
 });
+
+test('history searches each destination once, and cached periods stay exact after late data and rollback', (t) => {
+  const f = fixture(t); let clock = base;
+  let searched = 0;
+  const registerFunction = DatabaseSync.prototype.function;
+  t.mock.method(DatabaseSync.prototype, 'function', function (name, options, callback) {
+    return registerFunction.call(this, name, options, name === 'lower_unicode'
+      ? (value) => { searched++; return callback(value); } : callback);
+  });
+  const store = f.register(openTrafficHistoryStore(f.file, () => clock));
+  store.ingest([batch(clock, [], true)], 'live');
+  for (let minute = 0; minute < 60; minute++) {
+    clock = base + minute * 60_000 + 1_000;
+    store.ingest([batch(clock, [connection('a', (minute + 1) * 10, (minute + 1) * 20)])], 'live');
+  }
+  clock = base + 60 * 60_000;
+  const searchedQuery = query({ search: 'яндекс' });
+  assert.deepEqual(store.query(searchedQuery).totals, { uploadBytes: '600', downloadBytes: '1200' });
+  assert.ok(searched <= 2, `one destination must not be searched per time bucket (${searched} calls)`);
+  store.ingest([batch(clock + 1, [connection('a', 601, 1202)])], 'live');
+  assert.equal(store.query(searchedQuery).totals.uploadBytes, '600');
+  store.ingest([batch(base + 30_000, [connection('late', 7, 9)])], 'live');
+  assert.deepEqual(store.query(searchedQuery).totals, { uploadBytes: '607', downloadBytes: '1209' });
+  assert.throws(() => store.ingest([batch(base + 40_000, [
+    connection('rolled-back', 50, 80), connection('invalid', 'invalid', 1),
+  ])], 'live'));
+  assert.equal(store.query(searchedQuery).totals.uploadBytes, '607');
+  clock += 60_000;
+  assert.deepEqual(store.query(searchedQuery).totals, { uploadBytes: '608', downloadBytes: '1211' });
+  store.maintain();
+  assert.equal(store.query(searchedQuery).totals.uploadBytes, '608');
+});