Files
electric-cherry/tests/unit/rebalancing-runs.test.ts
T

468 lines
22 KiB
TypeScript

import { randomUUID } from 'node:crypto';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { newFixture, type Fixture } from '../helpers/fixture.js';
import { prepareNativeHost } from '../helpers/native-host.js';
import type { RebalanceMove } from '../../src/lib/rebalancing.js';
vi.mock('../../src/lib/server/pipeline/worker.js', () => ({
kick: vi.fn(async () => {}),
kickProvisioning: vi.fn(async () => {})
}));
vi.mock('../../src/lib/server/integrations/plutonium.js', () => ({
checkRelease: vi.fn(async () => {}),
freshRelease: () => 5354,
probeVersion: vi.fn(async () => 5354)
}));
vi.mock('../../src/lib/server/integrations/rcon.js', () => ({
rconStatusResult: vi.fn(async () => ({
ok: true,
endpoint: { host: 'agent', port: 28960 },
status: { players: 0, map: 'zombie_theater', gametype: 'zom', roster: [] }
})),
rconGetPlayerCount: vi.fn(async () => 0),
rconPlutoniumVersion: vi.fn(async () => ({ revision: 5354, error: null }))
}));
let fixture: Fixture, sourceId: string, targetId: string, libraryId: string;
const names = Array.from({ length: 4 }, (_, index) => `t5_zombie_theater_${index + 1}`);
const image = 'example.test/plutainer:rebalance';
const gib = 1024 ** 3;
const clockStep = () => vi.setSystemTime(Date.now() + 1000);
const settings = () => ({ hostIds: [sourceId, targetId], game: 'all' as const, maxMoves: 5 });
async function heartbeat(hostId: string, cpuUsagePercent: number) {
const { db } = await import('../../src/lib/server/db/index.js');
const { heartbeat } = await import('../../src/lib/server/hosts/index.js');
const instances = db()
.prepare(
'SELECT id,generation FROM native_instances WHERE host_id=? AND retired_at IS NULL'
)
.all(hostId) as { id: string; generation: number }[];
heartbeat(hostId, {
protocolVersion: 1,
cpuCount: 16,
cpuUsagePercent,
memoryTotal: 32 * gib,
memoryAvailable: 32 * gib - (2 + instances.length * 2) * gib,
dockerRunning: true,
directoriesWritable: true,
storageAvailable: 100 * gib,
runtimeImages: [image],
libraries: [{ id: libraryId, state: 'ready' }],
instanceResources: instances.map((instance) => ({
instanceId: instance.id,
generation: instance.generation,
uptimeSeconds: 3600,
cpuCores: 3,
memoryBytes: 2 * gib
}))
});
}
beforeEach(async () => {
vi.resetModules();
vi.clearAllMocks();
vi.useFakeTimers({ toFake: ['Date'] });
vi.setSystemTime(new Date('2026-10-08T16:00:00.000Z'));
fixture = await newFixture([]);
vi.stubEnv('PLUTAINER_IMAGE', image);
vi.stubEnv('RCON_PASSWORD', 'preserved-password');
const { db } = await import('../../src/lib/server/db/index.js');
const at = new Date().toISOString(),
launched = new Date(Date.now() - 600_000).toISOString();
for (const [index, name] of names.entries())
db()
.prepare(
`INSERT INTO servers(name,game,map,game_mode,map_slug,
config_file,game_id,image,game_port,container_state,container_checked_at,last_players,last_map,last_seen_at,last_checked_at,last_roster)
VALUES(?,'t5','zombie_theater','zombie','theater','dedicated_sp.cfg','t5sp',?,?,'running',?,0,'zombie_theater',?,?,'[]')`
)
.run(name, image, 28960 + index, at, at, at);
const prepared = await prepareNativeHost(names);
sourceId = prepared.host.id;
libraryId = prepared.versions.get('t5')!;
const hosts = await import('../../src/lib/server/hosts/index.js');
targetId = hosts.createHost({ name: 'Empty destination', portMin: 29000, portMax: 29020 }).id;
hosts.enrollHost(hosts.createEnrollment(targetId).token, {
protocolVersion: 1,
cpuCount: 16,
cpuUsagePercent: 6.25,
memoryTotal: 32 * gib,
memoryAvailable: 30 * gib,
dockerRunning: true,
directoriesWritable: true,
storageAvailable: 100 * gib,
runtimeImages: [image],
libraries: [{ id: libraryId, state: 'ready' }]
});
const libraries = await import('../../src/lib/server/game-files/index.js');
libraries.updateHostLibraryCaches(targetId, [
{
libraryVersionId: libraryId,
status: 'ready',
manifestSha256: libraries.getLibraryVersion(libraryId)!.manifestSha256
}
]);
const { nativeInstance } = await import('../../src/lib/server/backends/native.js');
const { beginDeployment } = await import('../../src/lib/server/provisioning/history.js');
const { captureRecipe, recipeStorage } =
await import('../../src/lib/server/fleet/launch-recipes.js');
const { reserveServerKey } = await import('../../src/lib/server/fleet/server-keys.js');
const { recordScript } = await import('../../src/lib/server/fleet/inventory.js');
for (const name of names) {
const instance = nativeInstance(name)!;
reserveServerKey(name, 't5', `saved-key-${name}`);
for (const protocol of ['udp', 'tcp'] as const)
hosts.reserveHostPort(sourceId, instance.id, {
port: instance.port,
bindAddress: '0.0.0.0',
protocol
});
const id = beginDeployment(name, 'create', 'fixture');
const operation = hosts.queueHostOperation(
sourceId,
'instance.create',
{ name, rconPassword: 'preserved-password' },
{ instanceId: instance.id, generation: 1, idempotencyKey: randomUUID() }
);
hosts.reportHostOperation(sourceId, operation.id, {
state: 'completed',
result: { state: 'running', healthy: true }
});
db()
.prepare(
"UPDATE server_deployments SET state='done',submitted_at=?,finished_at=? WHERE id=?"
)
.run(launched, launched, id);
db()
.prepare(
'UPDATE native_provisioning SET operation_id=?,verified_at=? WHERE deployment_id=?'
)
.run(operation.id, launched, id);
const recipe = captureRecipe('t5', 'theater', 'preserved-password');
const stored = recipeStorage({
...recipe,
image,
native: { hostId: sourceId, libraryId, image }
});
db()
.prepare(
'INSERT INTO server_launch_recipes(deployment_id,recipe,recipe_summary) VALUES(?,?,?)'
)
.run(id, stored.recipe, stored.recipeSummary);
for (const file of recipe.files)
recordScript(name, file.path, file.versionId, file.sha256, file.size, 'native-create');
}
await heartbeat(sourceId, 75);
await heartbeat(targetId, 6.25);
const backend = await import('../../src/lib/server/backends/native.js');
vi.spyOn(backend, 'nativeRequest').mockImplementation(async (_instance, kind) => {
if (kind !== 'instance.inspect') throw new Error(`Unexpected host request ${kind}`);
return { state: 'running', healthy: true, exists: true } as never;
});
});
afterEach(async () => {
vi.useRealTimers();
vi.restoreAllMocks();
await fixture.cleanup();
vi.unstubAllEnvs();
});
async function start() {
const service = await import('../../src/lib/server/rebalancing/index.js');
const preview = service.previewRebalance(settings());
expect(preview.moves.length, JSON.stringify(preview.exclusions)).toBeGreaterThanOrEqual(2);
return service.startRebalance(preview.id);
}
async function currentMove(runId: string): Promise<RebalanceMove> {
const { getRebalanceRun } = await import('../../src/lib/server/rebalancing/index.js');
return getRebalanceRun(runId)!.moves.find((move) => move.state === 'moving')!;
}
async function provision() {
const { processProvisioning } = await import('../../src/lib/server/provisioning/worker.js');
await processProvisioning();
}
async function finishMove(move: RebalanceMove) {
const { hostMigration } = await import('../../src/lib/server/provisioning/host-drain.js');
const hosts = await import('../../src/lib/server/hosts/index.js');
const { db } = await import('../../src/lib/server/db/index.js');
const removal = hostMigration(move.operationId!)!.remove_operation_id;
expect(removal).toBeTruthy();
clockStep();
hosts.reportHostOperation(sourceId, removal!, { state: 'completed' });
await provision();
await provision();
const created = db()
.prepare('SELECT operation_id FROM native_provisioning WHERE deployment_id=?')
.get(move.operationId) as { operation_id: string | null };
expect(created.operation_id).toBeTruthy();
clockStep();
hosts.reportHostOperation(targetId, created.operation_id!, {
state: 'completed',
result: { state: 'running', healthy: true }
});
clockStep();
await provision();
expect(
db().prepare('SELECT state FROM server_deployments WHERE id=?').get(move.operationId)
).toEqual({ state: 'done' });
}
describe('manual rebalance runs', () => {
it('previews without reserving capacity and rejects stale plans', async () => {
const { db } = await import('../../src/lib/server/db/index.js');
const { previewRebalance, startRebalance } =
await import('../../src/lib/server/rebalancing/index.js');
const ports = db().prepare('SELECT * FROM ec_host_ports').all();
const commands = db().prepare('SELECT id FROM ec_host_operations').all();
const preview = previewRebalance(settings());
expect(preview.moves.length).toBeGreaterThanOrEqual(2);
expect(db().prepare('SELECT * FROM ec_host_ports').all()).toEqual(ports);
expect(db().prepare('SELECT id FROM ec_host_operations').all()).toEqual(commands);
expect(db().prepare('SELECT * FROM rebalancing_runs').all()).toEqual([]);
db()
.prepare("UPDATE servers SET desired_state='stopped' WHERE name=?")
.run(preview.moves[0].server);
expect(() => startRebalance(preview.id)).toThrow('fleet changed');
});
it('starts idempotently, rejects concurrent runs, and queues exactly one move', async () => {
const service = await import('../../src/lib/server/rebalancing/index.js');
const preview = service.previewRebalance(settings());
const competing = service.previewRebalance(settings());
const run = service.startRebalance(preview.id);
expect(service.startRebalance(preview.id).id).toBe(run.id);
expect(() => service.startRebalance(competing.id)).toThrow('current rebalance');
expect(await service.runRebalancing()).toBe(true);
expect(await service.runRebalancing()).toBe(false);
const moves = service.getRebalanceRun(run.id)!.moves;
expect(moves.filter((move) => move.operationId)).toHaveLength(1);
expect(moves.filter((move) => move.state === 'planned').length).toBeGreaterThan(0);
});
it('rejects expired previews but returns an already-started run on a late retry', async () => {
const service = await import('../../src/lib/server/rebalancing/index.js');
const first = service.previewRebalance(settings());
const expired = service.previewRebalance(settings());
const run = service.startRebalance(first.id);
vi.setSystemTime(Date.now() + 6 * 60_000);
expect(service.startRebalance(first.id).id).toBe(run.id);
expect(() => service.startRebalance(expired.id)).toThrow('Preview expired');
});
it('stops before removal without changing placement or releasing the key, and allows a later host drain', async () => {
const run = await start();
const service = await import('../../src/lib/server/rebalancing/index.js');
await service.runRebalancing();
const move = await currentMove(run.id);
const stopped = service.stopRebalance(run.id);
expect(stopped.state).toBe('stopped');
expect(stopped.moves.every((move) => move.state === 'skipped')).toBe(true);
await provision();
const { nativeInstance } = await import('../../src/lib/server/backends/native.js');
const { getServerKey } = await import('../../src/lib/server/fleet/server-keys.js');
expect(nativeInstance(move.server)?.host_id).toBe(sourceId);
expect(getServerKey(move.server)).toBe(`saved-key-${move.server}`);
const { db } = await import('../../src/lib/server/db/index.js');
expect(
db().prepare("SELECT id FROM ec_host_operations WHERE kind='instance.remove'").all()
).toEqual([]);
const { updateHost } = await import('../../src/lib/server/hosts/index.js');
updateHost(sourceId, { mode: 'draining' });
const { reconcileHostDrains } =
await import('../../src/lib/server/provisioning/host-drain.js');
reconcileHostDrains();
expect(
db()
.prepare(
"SELECT id FROM server_deployments WHERE server=? AND triggered_by='host-drain'"
)
.get(move.server)
).toBeTruthy();
});
it('finishes a committed move after Stop while withdrawing the remaining moves', async () => {
const run = await start();
const service = await import('../../src/lib/server/rebalancing/index.js');
await service.runRebalancing();
await provision();
const move = await currentMove(run.id);
expect(service.stopRebalance(run.id).state).toBe('stopping');
const { cancelProvisioningJob } =
await import('../../src/lib/server/provisioning/cancellation.js');
expect(cancelProvisioningJob(move.operationId!)).toBe(false);
await finishMove(move);
await service.runRebalancing();
const finished = service.getRebalanceRun(run.id)!;
expect(finished.state).toBe('stopped');
expect(finished.moves.filter((move) => move.state === 'done')).toHaveLength(1);
expect(finished.moves.filter((move) => move.state === 'skipped')).toHaveLength(
run.moves.length - 1
);
const { nativeInstance } = await import('../../src/lib/server/backends/native.js');
const { getHost } = await import('../../src/lib/server/hosts/index.js');
expect(nativeInstance(move.server)).toMatchObject({
id: move.instanceId,
generation: 2,
host_id: targetId
});
expect(getHost(sourceId)?.mode).toBe('active');
});
it('treats cancellation from Provisioning as stopping the rest of the run', async () => {
const run = await start();
const service = await import('../../src/lib/server/rebalancing/index.js');
await service.runRebalancing();
const move = await currentMove(run.id);
const { cancelProvisioningJob } =
await import('../../src/lib/server/provisioning/cancellation.js');
expect(cancelProvisioningJob(move.operationId!)).toBe(true);
expect(service.getRebalanceRun(run.id)?.state).toBe('stopping');
await service.runRebalancing();
expect(service.getRebalanceRun(run.id)?.state).toBe('stopped');
expect(
service.getRebalanceRun(run.id)?.moves.every((move) => move.state === 'skipped')
).toBe(true);
});
it('waits for both hosts to report fresh measurements before admitting the next move', async () => {
const run = await start();
const service = await import('../../src/lib/server/rebalancing/index.js');
await service.runRebalancing();
await provision();
await finishMove(await currentMove(run.id));
expect(await service.runRebalancing()).toBe(false);
expect(service.getRebalanceRun(run.id)?.message).toContain('fresh source and destination');
expect(
service.getRebalanceRun(run.id)?.moves.filter((move) => move.operationId)
).toHaveLength(1);
clockStep();
await heartbeat(sourceId, 56.25);
expect(await service.runRebalancing()).toBe(false);
await heartbeat(targetId, 25);
expect(await service.runRebalancing()).toBe(true);
expect(
service.getRebalanceRun(run.id)?.moves.filter((move) => move.operationId)
).toHaveLength(2);
});
it('skips an uncommitted move that becomes occupied and continues other approved work', async () => {
const run = await start();
const service = await import('../../src/lib/server/rebalancing/index.js');
await service.runRebalancing();
const move = await currentMove(run.id);
const { db } = await import('../../src/lib/server/db/index.js');
db().prepare('UPDATE servers SET last_players=2 WHERE name=?').run(move.server);
await service.runRebalancing();
const fresh = service.getRebalanceRun(run.id)!;
expect(fresh.state).toBe('running');
expect(fresh.moves.find((item) => item.id === move.id)?.state).toBe('skipped');
expect(fresh.moves.filter((item) => item.state === 'moving')).toHaveLength(1);
});
it('halts remaining moves after a provisioning failure', async () => {
const run = await start();
const service = await import('../../src/lib/server/rebalancing/index.js');
await service.runRebalancing();
const move = await currentMove(run.id);
const { db } = await import('../../src/lib/server/db/index.js');
db()
.prepare(
"UPDATE server_deployments SET state='failed',error='Recovery needed' WHERE id=?"
)
.run(move.operationId);
expect(await service.runRebalancing()).toBe(false);
const failed = service.getRebalanceRun(run.id)!;
expect(failed.state).toBe('failed');
expect(failed.moves.find((item) => item.id === move.id)).toMatchObject({
state: 'failed',
message: 'Recovery needed'
});
expect(failed.moves.filter((item) => item.state === 'skipped')).toHaveLength(
run.moves.length - 1
);
});
it('keeps the source running until a global startup slot is available', async () => {
const run = await start();
const service = await import('../../src/lib/server/rebalancing/index.js');
await service.runRebalancing();
const move = await currentMove(run.id);
const { db } = await import('../../src/lib/server/db/index.js');
const { beginDeployment } = await import('../../src/lib/server/provisioning/history.js');
const occupied = beginDeployment('another-server', 'create', 'fixture');
await provision();
const { hostMigration } = await import('../../src/lib/server/provisioning/host-drain.js');
expect(hostMigration(move.operationId!)?.remove_operation_id).toBeNull();
db().prepare("UPDATE server_deployments SET state='done' WHERE id=?").run(occupied);
await provision();
expect(hostMigration(move.operationId!)?.remove_operation_id).toEqual(expect.any(String));
});
it('retains its global startup reservation between source removal and destination creation', async () => {
const run = await start();
const service = await import('../../src/lib/server/rebalancing/index.js');
await service.runRebalancing();
await provision();
const move = await currentMove(run.id);
const { hostMigration } = await import('../../src/lib/server/provisioning/migrations.js');
const { reportHostOperation } = await import('../../src/lib/server/hosts/index.js');
reportHostOperation(sourceId, hostMigration(move.operationId!)!.remove_operation_id!, {
state: 'completed'
});
await provision();
const { db } = await import('../../src/lib/server/db/index.js');
expect(
db().prepare('SELECT state FROM server_deployments WHERE id=?').get(move.operationId)
).toEqual({ state: 'waiting' });
const { queueNativeCreation } = await import('../../src/lib/server/provisioning/native.js');
const extra = queueNativeCreation({
game: 't5',
gameMode: 'zombie',
mapSlug: 'theater',
hostId: targetId,
libraryId,
image,
plutoKey: 'another-key'
});
expect(extra.ok).toBe(true);
if (!extra.ok) throw new Error(extra.message);
await provision();
expect(
db()
.prepare('SELECT operation_id FROM native_provisioning WHERE deployment_id=?')
.get(extra.deploymentId)
).toBeUndefined();
expect(
db()
.prepare('SELECT operation_id FROM native_provisioning WHERE deployment_id=?')
.get(move.operationId)
).toEqual({ operation_id: expect.any(String) });
});
it.each([
{ state: 'running', healthy: false },
{ state: 'stopped', healthy: false }
])(
'holds removal when fresh inspection returns $state with healthy=$healthy',
async (health) => {
const run = await start();
const service = await import('../../src/lib/server/rebalancing/index.js');
await service.runRebalancing();
const move = await currentMove(run.id);
const { nativeRequest } = await import('../../src/lib/server/backends/native.js');
vi.mocked(nativeRequest).mockResolvedValueOnce(health);
await provision();
const { hostMigration } =
await import('../../src/lib/server/provisioning/host-drain.js');
expect(hostMigration(move.operationId!)?.remove_operation_id).toBeNull();
const { db } = await import('../../src/lib/server/db/index.js');
expect(
db()
.prepare('SELECT stage FROM server_deployments WHERE id=?')
.get(move.operationId)
).toEqual({ stage: expect.stringContaining('no longer healthy') });
}
);
});