Files
qinglong/test/back/runDispatchLease.test.cjs

1047 lines
32 KiB
JavaScript

require('ts-node/register/transpile-only');
const assert = require('node:assert/strict');
const fs = require('node:fs/promises');
const os = require('node:os');
const path = require('node:path');
const { test } = require('node:test');
const { Sequelize } = require('sequelize');
const {
defineSchemaMigrationModel,
} = require('../../back/data/schemaMigration');
const { runSchemaMigration } = require('../../back/migrations/0002-run-schema');
const {
runCancellationRequestMigration,
} = require('../../back/migrations/0004-run-cancellation-request');
const {
runAttemptDeadlineMigration,
} = require('../../back/migrations/0006-run-attempt-deadline');
const {
workerRegistryMigration,
} = require('../../back/migrations/0008-worker-registry');
const {
runDispatchLeaseMigration,
} = require('../../back/migrations/0009-run-dispatch-lease');
const { runMigrations } = require('../../back/migrations/runner');
const {
LegacySequelizeRunDispatchLeaseRepository,
} = require('../../back/runtime/adapters/legacy-sequelize/runDispatchLeaseRepository');
const {
LegacySequelizeRunDispatchLeaseExpirySource,
} = require('../../back/runtime/adapters/legacy-sequelize/runDispatchLeaseExpirySource');
const {
LegacySequelizeRunRepository,
} = require('../../back/runtime/adapters/legacy-sequelize/runRepository');
const {
LegacySequelizeWorkerRegistryRepository,
} = require('../../back/runtime/adapters/legacy-sequelize/workerRegistryRepository');
const {
RemoteRunCompletionService,
} = require('../../back/runtime/application/remoteRunCompletionService');
const {
PrimaryRunCreator,
} = require('../../back/runtime/application/primaryRunCreator');
const {
RunCommandService,
} = require('../../back/runtime/application/runCommandService');
const {
RunDispatchLeaseService,
} = require('../../back/runtime/application/runDispatchLeaseService');
const {
RunDispatchLeaseExpiryService,
} = require('../../back/runtime/application/runDispatchLeaseExpiryService');
const {
RunDispatchLeaseExpiryScanner,
} = require('../../back/runtime/application/runDispatchLeaseExpiryScanner');
const {
BoundWorkerControlPlaneClient,
WorkerControlService,
} = require('../../back/runtime/application/workerControlService');
const {
RunDispatchLeaseFenceRejectedError,
} = require('../../back/runtime/domain/runDispatchLease');
const START = 1_750_000_000_000;
const SESSION_A = '019f7800-0000-7000-8000-000000000001';
const SESSION_B = '019f7800-0000-7000-8000-000000000002';
const TOKEN_A = 'lease_token_a_abcdefghijklmnopqrstuvwxyz0123456789';
const TOKEN_B = 'lease_token_b_abcdefghijklmnopqrstuvwxyz0123456789';
let idSequence = 1_000;
function nextId() {
idSequence += 1;
return `019f7800-0000-7000-8000-${String(idSequence).padStart(12, '0')}`;
}
function capabilities() {
return {
architecture: 'arm64',
operatingSystem: 'linux',
executors: ['remote_worker'],
runtimes: [{ name: 'node', version: '24.14.0' }],
labels: { tier: 'worker' },
capacity: { cpuCores: 2, memoryBytes: 512 * 1024 * 1024 },
features: ['direct_file_log'],
};
}
const migrationChain = [
runSchemaMigration,
runCancellationRequestMigration,
runAttemptDeadlineMigration,
workerRegistryMigration,
runDispatchLeaseMigration,
];
async function setupDatabase(database) {
await runMigrations({
database,
migrationModel: defineSchemaMigrationModel(database),
migrations: migrationChain,
logger: { info() {} },
});
}
async function fixture(t) {
const database = new Sequelize({
dialect: 'sqlite',
storage: ':memory:',
logging: false,
});
t.after(() => database.close());
await setupDatabase(database);
let nowMs = START;
const workers = new WorkerControlService(
new LegacySequelizeWorkerRegistryRepository(database),
{ leaseDurationMs: 60_000, clock: { now: () => nowMs } },
);
const leases = new LegacySequelizeRunDispatchLeaseRepository(database);
const leaseService = new RunDispatchLeaseService(leases, {
leaseDurationMs: 5_000,
clock: { now: () => nowMs },
createEventId: nextId,
});
const runs = new LegacySequelizeRunRepository(database);
return {
database,
workers,
leases,
leaseService,
runs,
setNow(value) {
nowMs = value;
},
getNow() {
return nowMs;
},
};
}
async function registerWorker(context, workerId, sessionId, max = 2) {
return new BoundWorkerControlPlaneClient(context.workers, {
workerId,
}).register({
workerId,
sessionId,
capabilities: capabilities(),
maxConcurrentRuns: max,
availableSlots: max,
});
}
async function createRun(context) {
return new PrimaryRunCreator(context.runs, nextId).create(
{
projectId: 'default',
taskId: `remote-task:${nextId()}`,
taskRevision: 'remote-revision-1',
triggerType: 'manual',
executionOrigin: 'manual',
acceptedAtMs: context.getNow(),
actor: { type: 'user', id: 'user:1' },
},
'remote_worker',
);
}
function claimRequest(reference, worker, token) {
return {
runId: reference.run.id,
attemptId: reference.attempt.id,
workerId: worker.id,
workerSessionId: worker.sessionId,
workerGeneration: worker.generation,
leaseToken: token,
};
}
function fenceRequest(reference, worker, lease) {
return {
attemptId: reference.attempt.id,
workerId: worker.id,
workerSessionId: worker.sessionId,
workerGeneration: worker.generation,
leaseGeneration: lease.leaseGeneration,
leaseToken: lease.leaseToken,
expectedVersion: lease.version,
};
}
async function moveToRunning(context, reference, claimed) {
const commands = new RunCommandService(context.runs, nextId);
const starting = await commands.transitionRunAttempt({
runId: reference.run.id,
attemptId: reference.attempt.id,
to: 'starting',
expectedRunVersion: claimed.event.sequence,
atMs: context.getNow() + 1,
actor: { type: 'worker', id: claimed.lease.workerId },
});
const runningAttempt = await commands.transitionRunAttempt({
runId: reference.run.id,
attemptId: reference.attempt.id,
to: 'running',
expectedRunVersion: starting.run.version,
atMs: context.getNow() + 2,
executorHandle: 'remote-worker-handle',
actor: { type: 'worker', id: claimed.lease.workerId },
});
return commands.transitionRun({
runId: reference.run.id,
to: 'running',
expectedVersion: runningAttempt.run.version,
atMs: context.getNow() + 2,
actor: { type: 'worker', id: claimed.lease.workerId },
});
}
test('claims idempotently, renews, releases, and fences an old generation', async (t) => {
const context = await fixture(t);
const workerA = await registerWorker(context, 'worker-a', SESSION_A);
const workerB = await registerWorker(context, 'worker-b', SESSION_B);
const reference = await createRun(context);
const requestA = claimRequest(reference, workerA, TOKEN_A);
const claimed = await context.leaseService.claim(
{ workerId: 'worker-a' },
requestA,
);
assert.equal(claimed.status, 'claimed');
assert.equal(claimed.lease.leaseGeneration, 1);
assert.equal(claimed.lease.version, 0);
assert.equal(claimed.event.type, 'run.dispatching');
assert.equal(
(await context.leaseService.claim({ workerId: 'worker-a' }, requestA))
.status,
'idempotent',
);
assert.equal(
(
await context.leaseService.claim(
{ workerId: 'worker-b' },
claimRequest(reference, workerB, TOKEN_B),
)
).status,
'leased',
);
context.setNow(START + 1_000);
const renewed = await context.leaseService.renew(
{ workerId: 'worker-a' },
fenceRequest(reference, workerA, claimed.lease),
);
assert.equal(renewed.version, 1);
assert.equal(renewed.expiresAtMs, START + 6_000);
await assert.rejects(
context.leaseService.renew(
{ workerId: 'worker-a' },
fenceRequest(reference, workerA, claimed.lease),
),
(error) =>
error instanceof RunDispatchLeaseFenceRejectedError &&
error.reason === 'version_mismatch',
);
context.setNow(START + 2_000);
const released = await context.leaseService.release(
{ workerId: 'worker-a' },
{
...fenceRequest(reference, workerA, renewed),
runId: reference.run.id,
reason: 'shutdown',
},
);
assert.equal(released.lease.status, 'released');
assert.equal(released.event.type, 'run.dispatch_released');
const eventsAfterRelease = await context.runs.listEvents(reference.run.id);
context.setNow(START + 2_001);
const releaseReplay = await context.leaseService.release(
{ workerId: 'worker-a' },
{
...fenceRequest(reference, workerA, renewed),
runId: reference.run.id,
reason: 'shutdown',
},
);
assert.equal(releaseReplay.lease.status, 'released');
assert.equal(releaseReplay.event, undefined);
assert.equal(
(await context.runs.listEvents(reference.run.id)).length,
eventsAfterRelease.length,
);
context.setNow(START + 2_002);
const reclaimed = await context.leaseService.claim(
{ workerId: 'worker-b' },
claimRequest(reference, workerB, TOKEN_B),
);
assert.equal(reclaimed.status, 'claimed');
assert.equal(reclaimed.lease.leaseGeneration, 2);
assert.equal(reclaimed.event.type, 'run.dispatch_reclaimed');
await assert.rejects(
context.leaseService.renew(
{ workerId: 'worker-a' },
fenceRequest(reference, workerA, renewed),
),
(error) =>
error instanceof RunDispatchLeaseFenceRejectedError &&
error.reason === 'worker_mismatch',
);
});
test('bounds active claims by Worker capacity and frees capacity after expiry', async (t) => {
const context = await fixture(t);
const worker = await registerWorker(context, 'worker-capacity', SESSION_A, 1);
const first = await createRun(context);
const second = await createRun(context);
assert.equal(
(
await context.leaseService.claim(
{ workerId: worker.id },
claimRequest(first, worker, TOKEN_A),
)
).status,
'claimed',
);
assert.equal(
(
await context.leaseService.claim(
{ workerId: worker.id },
claimRequest(second, worker, TOKEN_B),
)
).status,
'capacity_exhausted',
);
context.setNow(START + 5_001);
assert.equal(
(
await context.leaseService.claim(
{ workerId: worker.id },
claimRequest(second, worker, TOKEN_B),
)
).status,
'claimed',
);
});
test('keeps exact claim replay idempotent after capacity drops or drain begins', async (t) => {
const context = await fixture(t);
const worker = await registerWorker(context, 'worker-replay', SESSION_A, 1);
const first = await createRun(context);
const second = await createRun(context);
const request = claimRequest(first, worker, TOKEN_A);
const claimed = await context.leaseService.claim(
{ workerId: worker.id },
request,
);
assert.equal(claimed.status, 'claimed');
context.setNow(START + 100);
const client = new BoundWorkerControlPlaneClient(context.workers, {
workerId: worker.id,
});
const noCapacity = await client.heartbeat({
workerId: worker.id,
sessionId: worker.sessionId,
generation: worker.generation,
expectedVersion: worker.version,
availableSlots: 0,
});
assert.equal(
(await context.leaseService.claim({ workerId: worker.id }, request)).status,
'idempotent',
);
assert.equal(
(
await context.leaseService.claim(
{ workerId: worker.id },
claimRequest(second, worker, TOKEN_B),
)
).status,
'worker_unavailable',
);
context.setNow(START + 200);
await client.drain({
workerId: worker.id,
sessionId: worker.sessionId,
generation: worker.generation,
expectedVersion: noCapacity.version,
});
assert.equal(
(await context.leaseService.claim({ workerId: worker.id }, request)).status,
'idempotent',
);
const renewed = await context.leaseService.renew(
{ workerId: worker.id },
fenceRequest(first, worker, claimed.lease),
);
assert.equal(renewed.version, claimed.lease.version + 1);
});
test('reclaims an expired unstarted Attempt and permanently fences the old token', async (t) => {
const context = await fixture(t);
const workerA = await registerWorker(context, 'worker-expired-a', SESSION_A);
const workerB = await registerWorker(context, 'worker-expired-b', SESSION_B);
const reference = await createRun(context);
const first = await context.leaseService.claim(
{ workerId: workerA.id },
claimRequest(reference, workerA, TOKEN_A),
);
assert.equal(first.status, 'claimed');
context.setNow(START + 5_001);
const replacement = await context.leaseService.claim(
{ workerId: workerB.id },
claimRequest(reference, workerB, TOKEN_B),
);
assert.equal(replacement.status, 'claimed');
assert.equal(replacement.lease.leaseGeneration, 2);
assert.equal(replacement.event.payload.previous_state, 'expired');
await assert.rejects(
context.leaseService.renew(
{ workerId: workerA.id },
fenceRequest(reference, workerA, first.lease),
),
(error) =>
error instanceof RunDispatchLeaseFenceRejectedError &&
error.reason === 'worker_mismatch',
);
});
test('atomically releases an expired running lease and marks Attempt and Run lost', async (t) => {
const context = await fixture(t);
const worker = await registerWorker(context, 'worker-expiry', SESSION_A);
const reference = await createRun(context);
const claimed = await context.leaseService.claim(
{ workerId: worker.id },
claimRequest(reference, worker, TOKEN_A),
);
await moveToRunning(context, reference, claimed);
context.setNow(START + 5_001);
const service = new RunDispatchLeaseExpiryService(context.leases, {
clock: { now: () => context.getNow() },
createEventId: nextId,
});
const expired = await service.reconcile(
reference.run.id,
reference.attempt.id,
);
assert.equal(expired.status, 'lost');
assert.equal(expired.run.status, 'lost');
assert.equal(expired.attempt.status, 'lost');
assert.equal(expired.lease.status, 'released');
assert.equal(expired.lease.releaseReason, 'lease_expired');
const events = await context.runs.listEvents(reference.run.id);
assert.deepEqual(
events.slice(-2).map((event) => [
event.type,
event.actorType,
event.payload.lease_generation,
]),
[
['attempt.lost', 'reconciler', claimed.lease.leaseGeneration],
['run.lost', 'reconciler', claimed.lease.leaseGeneration],
],
);
const replay = await service.reconcile(
reference.run.id,
reference.attempt.id,
);
assert.equal(replay.status, 'already_expired');
assert.equal(
(await context.runs.listEvents(reference.run.id)).length,
events.length,
);
});
test('does not expire a live lease and rolls back when lost reconciliation fails', async (t) => {
const context = await fixture(t);
const worker = await registerWorker(
context,
'worker-expiry-rollback',
SESSION_A,
);
const reference = await createRun(context);
const claimed = await context.leaseService.claim(
{ workerId: worker.id },
claimRequest(reference, worker, TOKEN_A),
);
await moveToRunning(context, reference, claimed);
const service = new RunDispatchLeaseExpiryService(context.leases, {
clock: { now: () => context.getNow() },
createEventId: nextId,
});
assert.equal(
(await service.reconcile(reference.run.id, reference.attempt.id)).status,
'not_due',
);
context.setNow(START + 5_001);
await assert.rejects(
context.leases.expireWithLease(
{
runId: reference.run.id,
attemptId: reference.attempt.id,
observedAtMs: context.getNow(),
},
async (transaction) => {
const run = await transaction.findRunById(reference.run.id);
await transaction.compareAndSetRun(
{ ...run, version: run.version + 1 },
run.version,
);
throw new Error('simulated expiry crash');
},
),
/simulated expiry crash/,
);
assert.equal(
(await context.leases.findByAttemptId(reference.attempt.id)).status,
'leased',
);
assert.equal(
(await context.runs.findRunById(reference.run.id)).status,
'running',
);
assert.equal(
(await service.reconcile(reference.run.id, reference.attempt.id)).status,
'lost',
);
});
test('releases expired authority but leaves an accepted cancellation to its own lifecycle', async (t) => {
const context = await fixture(t);
const worker = await registerWorker(
context,
'worker-expiry-cancel',
SESSION_A,
);
const reference = await createRun(context);
const claimed = await context.leaseService.claim(
{ workerId: worker.id },
claimRequest(reference, worker, TOKEN_A),
);
await moveToRunning(context, reference, claimed);
const commands = new RunCommandService(context.runs, nextId);
await commands.requestCancellation({
runId: reference.run.id,
attemptId: reference.attempt.id,
atMs: START + 3,
reason: 'user',
actor: { type: 'user', id: 'user:1' },
});
context.setNow(START + 5_001);
const result = await new RunDispatchLeaseExpiryService(context.leases, {
clock: { now: () => context.getNow() },
createEventId: nextId,
}).reconcile(reference.run.id, reference.attempt.id);
assert.equal(result.status, 'cancellation_pending');
assert.equal(result.run.status, 'running');
assert.equal(result.attempt.status, 'running');
assert.equal(result.lease.releaseReason, 'lease_expired');
assert.equal(
(await context.runs.findRunById(reference.run.id)).cancelReason,
'user',
);
});
test('releases an expired unstarted authority without marking the Attempt lost', async (t) => {
const context = await fixture(t);
const workerA = await registerWorker(
context,
'worker-expiry-unstarted-a',
SESSION_A,
);
const workerB = await registerWorker(
context,
'worker-expiry-unstarted-b',
SESSION_B,
);
const reference = await createRun(context);
await context.leaseService.claim(
{ workerId: workerA.id },
claimRequest(reference, workerA, TOKEN_A),
);
context.setNow(START + 5_001);
const result = await new RunDispatchLeaseExpiryService(context.leases, {
clock: { now: () => context.getNow() },
createEventId: nextId,
}).reconcile(reference.run.id, reference.attempt.id);
assert.equal(result.status, 'unstarted_released');
assert.equal(result.run.status, 'dispatching');
assert.equal(result.attempt.status, 'claimed');
assert.equal(result.lease.releaseReason, 'lease_expired');
context.setNow(START + 5_002);
const replacement = await context.leaseService.claim(
{ workerId: workerB.id },
claimRequest(reference, workerB, TOKEN_B),
);
assert.equal(replacement.status, 'claimed');
assert.equal(replacement.lease.leaseGeneration, 2);
});
test('scans expired leases in bounded stable pages without touching a live lease', async (t) => {
const context = await fixture(t);
const worker = await registerWorker(
context,
'worker-expiry-scanner',
SESSION_A,
3,
);
const references = [];
for (let index = 0; index < 3; index += 1) {
references.push(await createRun(context));
}
const claims = [];
for (const [index, reference] of references.entries()) {
claims.push(
await context.leaseService.claim(
{ workerId: worker.id },
claimRequest(
reference,
worker,
`scanner_lease_${index}_abcdefghijklmnopqrstuvwxyz0123456789`,
),
),
);
}
context.setNow(START + 100);
const live = await context.leaseService.renew(
{ workerId: worker.id },
fenceRequest(references[2], worker, claims[2].lease),
);
context.setNow(START + 5_001);
const scanner = new RunDispatchLeaseExpiryScanner(
new LegacySequelizeRunDispatchLeaseExpirySource(context.database),
new RunDispatchLeaseExpiryService(context.leases, {
clock: { now: () => context.getNow() },
createEventId: nextId,
}),
{ clock: { now: () => context.getNow() } },
);
const first = await scanner.scan({ limit: 1 });
assert.equal(first.scanned, 1);
assert.equal(first.counts.unstarted_released, 1);
assert.equal(first.truncated, true);
assert.ok(first.nextCursor);
const second = await scanner.scan({ after: first.nextCursor, limit: 1 });
assert.equal(second.scanned, 1);
assert.equal(second.counts.unstarted_released, 1);
const final = await scanner.scan({
after: second.nextCursor,
limit: 1,
});
assert.equal(final.scanned, 0);
assert.equal(final.truncated, false);
assert.equal(
(await context.leases.findByAttemptId(references[2].attempt.id)).version,
live.version,
);
assert.equal(
(await context.leases.findByAttemptId(references[2].attempt.id)).status,
'leased',
);
});
test('commits remote completion and lease completion atomically with replay', async (t) => {
const context = await fixture(t);
const worker = await registerWorker(context, 'worker-complete', SESSION_A);
const reference = await createRun(context);
const claimed = await context.leaseService.claim(
{ workerId: worker.id },
claimRequest(reference, worker, TOKEN_A),
);
assert.equal(claimed.status, 'claimed');
await moveToRunning(context, reference, claimed);
context.setNow(START + 10);
const service = new RemoteRunCompletionService(context.runs, context.leases, {
clock: { now: () => context.getNow() },
createEventId: nextId,
});
const command = {
runId: reference.run.id,
attemptId: reference.attempt.id,
callbackSequence: 1,
result: {
outcome: 'succeeded',
startedAtMs: START + 2,
finishedAtMs: START + 9,
exitCode: 0,
},
executorType: 'remote_worker',
workerId: worker.id,
workerSessionId: worker.sessionId,
workerGeneration: worker.generation,
leaseGeneration: claimed.lease.leaseGeneration,
leaseToken: claimed.lease.leaseToken,
expectedLeaseVersion: claimed.lease.version,
};
const applied = await service.complete({ workerId: worker.id }, command);
assert.equal(applied.status, 'applied');
assert.equal(applied.run.status, 'succeeded');
assert.equal(
(await context.leases.findByAttemptId(reference.attempt.id)).status,
'completed',
);
const events = await context.runs.listEvents(reference.run.id);
assert.deepEqual(
events
.slice(-2)
.map((event) => [event.type, event.actorType, event.actorId]),
[
['attempt.succeeded', 'worker', worker.id],
['run.succeeded', 'worker', worker.id],
],
);
assert.equal(
(await service.complete({ workerId: worker.id }, command)).status,
'already_terminal',
);
assert.equal(
(await context.runs.listEvents(reference.run.id)).length,
events.length,
);
});
test('accepts only an exact completion replay after the Worker session is replaced', async (t) => {
const context = await fixture(t);
const worker = await registerWorker(
context,
'worker-completion-replay',
SESSION_A,
);
const reference = await createRun(context);
const claimed = await context.leaseService.claim(
{ workerId: worker.id },
claimRequest(reference, worker, TOKEN_A),
);
await moveToRunning(context, reference, claimed);
const command = {
runId: reference.run.id,
attemptId: reference.attempt.id,
callbackSequence: 1,
result: {
outcome: 'succeeded',
startedAtMs: START + 2,
finishedAtMs: START + 9,
exitCode: 0,
},
executorType: 'remote_worker',
workerId: worker.id,
workerSessionId: worker.sessionId,
workerGeneration: worker.generation,
leaseGeneration: claimed.lease.leaseGeneration,
leaseToken: claimed.lease.leaseToken,
expectedLeaseVersion: claimed.lease.version,
};
const service = new RemoteRunCompletionService(context.runs, context.leases, {
clock: { now: () => context.getNow() },
createEventId: nextId,
});
context.setNow(START + 20);
assert.equal(
(await service.complete({ workerId: worker.id }, command)).status,
'applied',
);
const events = await context.runs.listEvents(reference.run.id);
context.setNow(START + 21);
await registerWorker(context, worker.id, SESSION_B);
context.setNow(START + 10);
assert.equal(
(await service.complete({ workerId: worker.id }, command)).status,
'already_terminal',
);
assert.equal(
(await context.runs.listEvents(reference.run.id)).length,
events.length,
);
await assert.rejects(
service.complete(
{ workerId: worker.id },
{ ...command, callbackSequence: 2 },
),
/Primary completion callback sequence is invalid/,
);
await assert.rejects(
service.complete(
{ workerId: worker.id },
{
...command,
result: { ...command.result, outcome: 'failed', exitCode: 9 },
},
),
/Primary completion target state is inconsistent/,
);
await assert.rejects(
service.complete(
{ workerId: worker.id },
{ ...command, leaseToken: TOKEN_B },
),
(error) =>
error instanceof RunDispatchLeaseFenceRejectedError &&
error.reason === 'lease_token_mismatch',
);
});
test('rolls back Run and lease when completion crashes before commit', async (t) => {
const context = await fixture(t);
const worker = await registerWorker(context, 'worker-crash', SESSION_A);
const reference = await createRun(context);
const claimed = await context.leaseService.claim(
{ workerId: worker.id },
claimRequest(reference, worker, TOKEN_A),
);
await moveToRunning(context, reference, claimed);
context.setNow(START + 10);
let fail = true;
const failingLeases = {
findByAttemptId: (...args) => context.leases.findByAttemptId(...args),
claim: (...args) => context.leases.claim(...args),
renew: (...args) => context.leases.renew(...args),
release: (...args) => context.leases.release(...args),
completeWithLease(command, work) {
return context.leases.completeWithLease(
command,
async (transaction, lease) => {
const value = await work(transaction, lease);
if (fail) {
fail = false;
throw new Error('simulated crash before commit');
}
return value;
},
);
},
};
const command = {
runId: reference.run.id,
attemptId: reference.attempt.id,
callbackSequence: 1,
result: {
outcome: 'succeeded',
startedAtMs: START + 2,
finishedAtMs: START + 9,
exitCode: 0,
},
executorType: 'remote_worker',
workerId: worker.id,
workerSessionId: worker.sessionId,
workerGeneration: worker.generation,
leaseGeneration: claimed.lease.leaseGeneration,
leaseToken: claimed.lease.leaseToken,
expectedLeaseVersion: claimed.lease.version,
};
const failing = new RemoteRunCompletionService(context.runs, failingLeases, {
clock: { now: () => context.getNow() },
createEventId: nextId,
});
await assert.rejects(
failing.complete({ workerId: worker.id }, command),
/simulated crash/,
);
assert.equal(
(await context.runs.findRunById(reference.run.id)).status,
'running',
);
assert.equal(
(await context.leases.findByAttemptId(reference.attempt.id)).status,
'leased',
);
assert.equal(
(
await new RemoteRunCompletionService(context.runs, context.leases, {
clock: { now: () => context.getNow() },
createEventId: nextId,
}).complete({ workerId: worker.id }, command)
).status,
'applied',
);
});
test('rejects a wrong completion token without mutating Run or lease', async (t) => {
const context = await fixture(t);
const worker = await registerWorker(context, 'worker-token', SESSION_A);
const reference = await createRun(context);
const claimed = await context.leaseService.claim(
{ workerId: worker.id },
claimRequest(reference, worker, TOKEN_A),
);
await moveToRunning(context, reference, claimed);
context.setNow(START + 10);
const beforeEvents = await context.runs.listEvents(reference.run.id);
const service = new RemoteRunCompletionService(context.runs, context.leases, {
clock: { now: () => context.getNow() },
createEventId: nextId,
});
await assert.rejects(
service.complete(
{ workerId: worker.id },
{
runId: reference.run.id,
attemptId: reference.attempt.id,
callbackSequence: 1,
result: {
outcome: 'succeeded',
startedAtMs: START + 2,
finishedAtMs: START + 9,
exitCode: 0,
},
executorType: 'remote_worker',
workerId: worker.id,
workerSessionId: worker.sessionId,
workerGeneration: worker.generation,
leaseGeneration: claimed.lease.leaseGeneration,
leaseToken: TOKEN_B,
expectedLeaseVersion: claimed.lease.version,
},
),
(error) =>
error instanceof RunDispatchLeaseFenceRejectedError &&
error.reason === 'lease_token_mismatch',
);
assert.equal(
(await context.runs.findRunById(reference.run.id)).status,
'running',
);
assert.equal(
(await context.leases.findByAttemptId(reference.attempt.id)).status,
'leased',
);
assert.equal(
(await context.runs.listEvents(reference.run.id)).length,
beforeEvents.length,
);
});
test('rejects completion from a replaced Worker session', async (t) => {
const context = await fixture(t);
const worker = await registerWorker(context, 'worker-replaced', SESSION_A);
const reference = await createRun(context);
const claimed = await context.leaseService.claim(
{ workerId: worker.id },
claimRequest(reference, worker, TOKEN_A),
);
await moveToRunning(context, reference, claimed);
context.setNow(START + 5);
await registerWorker(context, worker.id, SESSION_B);
const service = new RemoteRunCompletionService(context.runs, context.leases, {
clock: { now: () => START + 10 },
createEventId: nextId,
});
await assert.rejects(
service.complete(
{ workerId: worker.id },
{
runId: reference.run.id,
attemptId: reference.attempt.id,
callbackSequence: 1,
result: {
outcome: 'succeeded',
startedAtMs: START + 2,
finishedAtMs: START + 9,
exitCode: 0,
},
executorType: 'remote_worker',
workerId: worker.id,
workerSessionId: worker.sessionId,
workerGeneration: worker.generation,
leaseGeneration: claimed.lease.leaseGeneration,
leaseToken: claimed.lease.leaseToken,
expectedLeaseVersion: claimed.lease.version,
},
),
(error) =>
error instanceof RunDispatchLeaseFenceRejectedError &&
error.reason === 'worker_unavailable',
);
assert.equal(
(await context.runs.findRunById(reference.run.id)).status,
'running',
);
});
test('serializes two control-plane claims for one Attempt', async (t) => {
const root = await fs.mkdtemp(path.join(os.tmpdir(), 'ql3-run-leases-'));
const storage = path.join(root, 'leases.sqlite');
const databases = [
new Sequelize({ dialect: 'sqlite', storage, logging: false }),
new Sequelize({ dialect: 'sqlite', storage, logging: false }),
];
t.after(async () => {
await Promise.all(databases.map((database) => database.close()));
await fs.rm(root, { recursive: true, force: true });
});
await setupDatabase(databases[0]);
for (const database of databases) {
await database.query('PRAGMA journal_mode=WAL');
await database.query('PRAGMA busy_timeout=1000');
}
let nowMs = START;
const workerService = new WorkerControlService(
new LegacySequelizeWorkerRegistryRepository(databases[0]),
{ leaseDurationMs: 60_000, clock: { now: () => nowMs } },
);
const context = {
workers: workerService,
runs: new LegacySequelizeRunRepository(databases[0]),
getNow: () => nowMs,
};
const workerA = await registerWorker(context, 'worker-race-a', SESSION_A);
const workerB = await registerWorker(context, 'worker-race-b', SESSION_B);
const reference = await createRun(context);
const services = databases.map(
(database) =>
new RunDispatchLeaseService(
new LegacySequelizeRunDispatchLeaseRepository(database),
{
leaseDurationMs: 5_000,
clock: { now: () => nowMs },
createEventId: nextId,
},
),
);
const results = await Promise.all([
services[0].claim(
{ workerId: workerA.id },
claimRequest(reference, workerA, TOKEN_A),
),
services[1].claim(
{ workerId: workerB.id },
claimRequest(reference, workerB, TOKEN_B),
),
]);
assert.deepEqual(results.map((result) => result.status).sort(), [
'claimed',
'leased',
]);
});