mirror of
https://github.com/whyour/qinglong.git
synced 2026-09-22 19:29:13 +08:00
feat(ql3): establish 3.0 incubation baseline
This commit is contained in:
whyour
@@ -0,0 +1,307 @@
|
||||
// PostgreSQL source for remote-execution dispatch candidates and lease recovery.
|
||||
import type {
|
||||
ClusterDispatchCandidate,
|
||||
ClusterDispatchCandidatePage,
|
||||
ClusterDispatchCandidateCursor,
|
||||
ClusterDispatchRecovery,
|
||||
ClusterDispatchSource,
|
||||
} from '@qinglong/runtime-core/remote-dispatch';
|
||||
import {
|
||||
assertRemoteDispatchPageSize,
|
||||
normalizeClusterDispatchCandidate,
|
||||
normalizeClusterDispatchCursor,
|
||||
} from '@qinglong/runtime-core/remote-dispatch';
|
||||
import type {
|
||||
PostgresPool,
|
||||
RunDispatchLeaseRecord,
|
||||
RunDispatchLeaseStatus,
|
||||
} from '@qinglong/runtime-core';
|
||||
import {
|
||||
RUN_DISPATCH_LEASE_STATUSES,
|
||||
assertRunDispatchId,
|
||||
assertRunDispatchLeaseRecord,
|
||||
} from '@qinglong/runtime-core';
|
||||
|
||||
type Row = Record<string, unknown>;
|
||||
|
||||
const CANDIDATE_COLUMNS = `
|
||||
run.id AS "runId",
|
||||
attempt.id AS "attemptId",
|
||||
COALESCE(workflow_task.project_id, run.project_id) AS "projectId",
|
||||
COALESCE(workflow_task.task_id, run.task_id) AS "taskId",
|
||||
COALESCE(workflow_task.task_revision, run.task_revision) AS "taskRevision",
|
||||
run.priority,
|
||||
COALESCE(workflow_task.admitted_at_ms, run.queued_at_ms) AS "queuedAtMs",
|
||||
attempt.created_at_ms AS "attemptCreatedAtMs",
|
||||
attempt.attempt AS "attemptNumber",
|
||||
attempt.executor_type AS "executorType"`.trim();
|
||||
|
||||
const LEASE_COLUMNS = `
|
||||
lease.status AS "leaseStatus",
|
||||
lease.version AS "leaseVersion",
|
||||
lease.lease_generation AS "leaseGeneration",
|
||||
lease.worker_id AS "workerId",
|
||||
lease.worker_session_id AS "workerSessionId",
|
||||
lease.worker_generation AS "workerGeneration",
|
||||
lease.lease_token_digest AS "leaseTokenDigest",
|
||||
lease.acquired_at_ms AS "acquiredAtMs",
|
||||
lease.renewed_at_ms AS "renewedAtMs",
|
||||
lease.expires_at_ms AS "expiresAtMs",
|
||||
lease.released_at_ms AS "releasedAtMs",
|
||||
lease.release_reason AS "releaseReason",
|
||||
lease.completed_at_ms AS "completedAtMs",
|
||||
lease.updated_at_ms AS "leaseUpdatedAtMs"`.trim();
|
||||
|
||||
function string(row: Row, key: string): string {
|
||||
const value = row[key];
|
||||
if (typeof value !== 'string' || value.length === 0) {
|
||||
throw new TypeError(`PostgreSQL cluster dispatch ${key} is invalid`);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
function integer(row: Row, key: string): number {
|
||||
const value = row[key];
|
||||
const normalized = typeof value === 'string' && /^(0|[1-9]\d*)$/.test(value)
|
||||
? Number(value)
|
||||
: value;
|
||||
if (typeof normalized !== 'number' || !Number.isSafeInteger(normalized)) {
|
||||
throw new TypeError(`PostgreSQL cluster dispatch ${key} is invalid`);
|
||||
}
|
||||
return normalized;
|
||||
}
|
||||
|
||||
function optionalInteger(row: Row, key: string): number | undefined {
|
||||
return row[key] === null || row[key] === undefined ? undefined : integer(row, key);
|
||||
}
|
||||
|
||||
function optionalString(row: Row, key: string): string | undefined {
|
||||
return row[key] === null || row[key] === undefined ? undefined : string(row, key);
|
||||
}
|
||||
|
||||
function candidate(row: Row): ClusterDispatchCandidate {
|
||||
return normalizeClusterDispatchCandidate({
|
||||
runId: string(row, 'runId'),
|
||||
attemptId: string(row, 'attemptId'),
|
||||
projectId: string(row, 'projectId'),
|
||||
taskId: string(row, 'taskId'),
|
||||
taskRevision: string(row, 'taskRevision'),
|
||||
priority: integer(row, 'priority'),
|
||||
queuedAtMs: integer(row, 'queuedAtMs'),
|
||||
attemptCreatedAtMs: integer(row, 'attemptCreatedAtMs'),
|
||||
attemptNumber: integer(row, 'attemptNumber'),
|
||||
executorType: string(row, 'executorType') as ClusterDispatchCandidate['executorType'],
|
||||
});
|
||||
}
|
||||
|
||||
function lease(row: Row): RunDispatchLeaseRecord {
|
||||
const status = string(row, 'leaseStatus') as RunDispatchLeaseStatus;
|
||||
if (!RUN_DISPATCH_LEASE_STATUSES.includes(status)) {
|
||||
throw new TypeError('PostgreSQL cluster dispatch lease status is invalid');
|
||||
}
|
||||
const releasedAtMs = optionalInteger(row, 'releasedAtMs');
|
||||
const releaseReason = optionalString(row, 'releaseReason');
|
||||
const completedAtMs = optionalInteger(row, 'completedAtMs');
|
||||
const value: RunDispatchLeaseRecord = Object.freeze({
|
||||
attemptId: string(row, 'attemptId'),
|
||||
runId: string(row, 'runId'),
|
||||
status,
|
||||
version: integer(row, 'leaseVersion'),
|
||||
leaseGeneration: integer(row, 'leaseGeneration'),
|
||||
workerId: string(row, 'workerId'),
|
||||
workerSessionId: string(row, 'workerSessionId'),
|
||||
workerGeneration: integer(row, 'workerGeneration'),
|
||||
leaseTokenDigest: string(row, 'leaseTokenDigest'),
|
||||
acquiredAtMs: integer(row, 'acquiredAtMs'),
|
||||
renewedAtMs: integer(row, 'renewedAtMs'),
|
||||
expiresAtMs: integer(row, 'expiresAtMs'),
|
||||
updatedAtMs: integer(row, 'leaseUpdatedAtMs'),
|
||||
...(releasedAtMs === undefined ? {} : { releasedAtMs }),
|
||||
...(releaseReason === undefined ? {} : {
|
||||
releaseReason: releaseReason as NonNullable<RunDispatchLeaseRecord['releaseReason']>,
|
||||
}),
|
||||
...(completedAtMs === undefined ? {} : { completedAtMs }),
|
||||
});
|
||||
assertRunDispatchLeaseRecord(value);
|
||||
return value;
|
||||
}
|
||||
|
||||
function cursor(value: ClusterDispatchCandidate): ClusterDispatchCandidateCursor {
|
||||
return Object.freeze({
|
||||
priority: value.priority,
|
||||
queuedAtMs: value.queuedAtMs,
|
||||
attemptCreatedAtMs: value.attemptCreatedAtMs,
|
||||
attemptId: value.attemptId,
|
||||
});
|
||||
}
|
||||
|
||||
export class PostgresClusterDispatchSource implements ClusterDispatchSource {
|
||||
constructor(private readonly pool: PostgresPool) {
|
||||
if (!pool || typeof pool.query !== 'function') {
|
||||
throw new TypeError('PostgreSQL cluster dispatch pool is invalid');
|
||||
}
|
||||
}
|
||||
|
||||
async listClusterDispatchCandidates(options: Readonly<{
|
||||
limit: number;
|
||||
after?: ClusterDispatchCandidateCursor;
|
||||
}>): Promise<ClusterDispatchCandidatePage> {
|
||||
if (!options || typeof options !== 'object' || Array.isArray(options)) {
|
||||
throw new TypeError('PostgreSQL cluster dispatch options are invalid');
|
||||
}
|
||||
const keys = Object.keys(options);
|
||||
if (!keys.includes('limit') || keys.some((key) => !['after', 'limit'].includes(key))) {
|
||||
throw new TypeError('PostgreSQL cluster dispatch options shape is invalid');
|
||||
}
|
||||
assertRemoteDispatchPageSize(options.limit);
|
||||
const after = options.after === undefined
|
||||
? undefined
|
||||
: normalizeClusterDispatchCursor(options.after);
|
||||
const result = await this.pool.query<Row>(
|
||||
`WITH observation AS MATERIALIZED (
|
||||
SELECT floor(extract(epoch FROM clock_timestamp()) * 1000)::bigint
|
||||
AS observed_at_ms
|
||||
), candidates AS MATERIALIZED (
|
||||
SELECT ${CANDIDATE_COLUMNS}
|
||||
FROM observation
|
||||
JOIN "ql3"."runs" AS run ON true
|
||||
JOIN "ql3"."run_attempts" AS attempt ON attempt.run_id = run.id
|
||||
LEFT JOIN
|
||||
"ql3"."plugin_package_workflow_task_attempt_admissions"
|
||||
AS workflow_task
|
||||
ON workflow_task.attempt_id = attempt.id
|
||||
LEFT JOIN "ql3"."step_runs" AS workflow_step
|
||||
ON workflow_step.run_id = workflow_task.run_id
|
||||
AND workflow_step.id = workflow_task.step_run_id
|
||||
LEFT JOIN "ql3"."run_dispatch_leases" AS lease
|
||||
ON lease.attempt_id = attempt.id
|
||||
WHERE run.execution_owner = 'runtime'
|
||||
AND run.cancel_requested_at_ms IS NULL
|
||||
AND attempt.status = 'claimed'
|
||||
AND attempt.executor_type = 'remote_worker'
|
||||
AND (
|
||||
(
|
||||
workflow_task.attempt_id IS NULL
|
||||
AND run.status IN ('queued', 'dispatching')
|
||||
AND run.queued_at_ms IS NOT NULL
|
||||
) OR (
|
||||
workflow_task.attempt_id IS NOT NULL
|
||||
AND run.status = 'running'
|
||||
AND attempt.step_run_id = workflow_task.step_run_id
|
||||
AND workflow_step.status = 'ready'
|
||||
AND workflow_step.version = workflow_task.step_run_version
|
||||
AND workflow_step.step_run_digest =
|
||||
workflow_task.step_run_digest
|
||||
)
|
||||
)
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM "ql3"."run_attempts" AS newer
|
||||
WHERE newer.run_id = run.id
|
||||
AND newer.attempt > attempt.attempt
|
||||
AND (
|
||||
workflow_task.attempt_id IS NULL
|
||||
OR newer.step_run_id = attempt.step_run_id
|
||||
)
|
||||
)
|
||||
AND (
|
||||
lease.attempt_id IS NULL OR lease.status = 'released'
|
||||
OR (lease.status = 'leased' AND lease.expires_at_ms <= observation.observed_at_ms)
|
||||
)
|
||||
AND (
|
||||
$1::integer IS NULL OR run.priority < $1
|
||||
OR (
|
||||
run.priority = $1 AND
|
||||
COALESCE(workflow_task.admitted_at_ms, run.queued_at_ms) > $2
|
||||
)
|
||||
OR (run.priority = $1
|
||||
AND COALESCE(
|
||||
workflow_task.admitted_at_ms, run.queued_at_ms
|
||||
) = $2
|
||||
AND attempt.created_at_ms > $3)
|
||||
OR (run.priority = $1
|
||||
AND COALESCE(
|
||||
workflow_task.admitted_at_ms, run.queued_at_ms
|
||||
) = $2
|
||||
AND attempt.created_at_ms = $3 AND attempt.id > $4)
|
||||
)
|
||||
ORDER BY run.priority DESC,
|
||||
COALESCE(
|
||||
workflow_task.admitted_at_ms, run.queued_at_ms
|
||||
),
|
||||
attempt.created_at_ms, attempt.id
|
||||
LIMIT $5
|
||||
)
|
||||
SELECT observation.observed_at_ms AS "observedAtMs", candidates.*
|
||||
FROM observation LEFT JOIN candidates ON true
|
||||
ORDER BY candidates.priority DESC, candidates."queuedAtMs",
|
||||
candidates."attemptCreatedAtMs", candidates."attemptId"`,
|
||||
[
|
||||
after?.priority ?? null,
|
||||
after?.queuedAtMs ?? null,
|
||||
after?.attemptCreatedAtMs ?? null,
|
||||
after?.attemptId ?? null,
|
||||
options.limit + 1,
|
||||
],
|
||||
);
|
||||
if (result.rows.length < 1) {
|
||||
throw new TypeError('PostgreSQL cluster dispatch observation is missing');
|
||||
}
|
||||
const observedAtMs = integer(result.rows[0]!, 'observedAtMs');
|
||||
const mapped = result.rows
|
||||
.filter((row) => row.attemptId !== null && row.attemptId !== undefined)
|
||||
.map(candidate);
|
||||
const truncated = mapped.length > options.limit;
|
||||
const candidates = Object.freeze(mapped.slice(0, options.limit));
|
||||
const last = candidates.at(-1);
|
||||
return Object.freeze({
|
||||
observedAtMs,
|
||||
candidates,
|
||||
truncated,
|
||||
...(truncated && last ? { next: cursor(last) } : {}),
|
||||
});
|
||||
}
|
||||
|
||||
async findClusterDispatchRecovery(offerId: string): Promise<ClusterDispatchRecovery | null> {
|
||||
assertRunDispatchId('offerId', offerId);
|
||||
const result = await this.pool.query<Row>(
|
||||
`WITH observation AS MATERIALIZED (
|
||||
SELECT floor(extract(epoch FROM clock_timestamp()) * 1000)::bigint
|
||||
AS observed_at_ms
|
||||
)
|
||||
SELECT observation.observed_at_ms AS "observedAtMs",
|
||||
${CANDIDATE_COLUMNS}, ${LEASE_COLUMNS},
|
||||
(
|
||||
worker.worker_id = lease.worker_id
|
||||
AND worker.session_id = lease.worker_session_id
|
||||
AND worker.generation = lease.worker_generation
|
||||
AND worker.status IN ('online', 'draining')
|
||||
AND worker.lease_expires_at_ms > observation.observed_at_ms
|
||||
) AS "workerCurrent"
|
||||
FROM observation
|
||||
JOIN "ql3"."run_dispatch_leases" AS lease ON lease.offer_id = $1
|
||||
JOIN "ql3"."runs" AS run ON run.id = lease.run_id
|
||||
JOIN "ql3"."run_attempts" AS attempt ON attempt.id = lease.attempt_id
|
||||
LEFT JOIN
|
||||
"ql3"."plugin_package_workflow_task_attempt_admissions"
|
||||
AS workflow_task
|
||||
ON workflow_task.attempt_id = attempt.id
|
||||
LEFT JOIN "ql3"."step_runs" AS workflow_step
|
||||
ON workflow_step.run_id = workflow_task.run_id
|
||||
AND workflow_step.id = workflow_task.step_run_id
|
||||
LEFT JOIN "ql3"."worker_sessions" AS worker
|
||||
ON worker.worker_id = lease.worker_id
|
||||
LIMIT 2`,
|
||||
[offerId],
|
||||
);
|
||||
if (result.rows.length === 0) return null;
|
||||
if (result.rows.length !== 1 || typeof result.rows[0]!.workerCurrent !== 'boolean') {
|
||||
throw new TypeError('PostgreSQL cluster dispatch recovery is invalid');
|
||||
}
|
||||
return Object.freeze({
|
||||
observedAtMs: integer(result.rows[0]!, 'observedAtMs'),
|
||||
candidate: candidate(result.rows[0]!),
|
||||
lease: lease(result.rows[0]!),
|
||||
workerCurrent: result.rows[0]!.workerCurrent as boolean,
|
||||
});
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
+96
@@ -0,0 +1,96 @@
|
||||
// PostgreSQL Remote Worker recovery evidence is owned by this domain.
|
||||
import type {
|
||||
ClusterControlRecoveryEvidence,
|
||||
ClusterControlRecoveryEvidenceInspectionContext,
|
||||
ClusterControlRecoveryExecutorEvidenceProvider,
|
||||
ClusterControlRecoveryProbeTarget,
|
||||
} from '@qinglong/runtime-core';
|
||||
import type { PostgresPool } from '@qinglong/runtime-core';
|
||||
import type { WorkerExecutionAttestationRepository } from '@qinglong/runtime-core/worker-attestation';
|
||||
|
||||
export const REMOTE_WORKER_ATTESTATION_LIMITS = Object.freeze({
|
||||
defaultRunningFreshnessMs: 30_000,
|
||||
maxRunningFreshnessMs: 300_000,
|
||||
});
|
||||
|
||||
type Row = Record<string, unknown>;
|
||||
|
||||
export class PostgresRemoteWorkerAttestationEvidenceProvider
|
||||
implements ClusterControlRecoveryExecutorEvidenceProvider
|
||||
{
|
||||
readonly executorType = 'remote-worker';
|
||||
readonly requiredIdentity = Object.freeze([
|
||||
'workerId',
|
||||
'workerSessionId',
|
||||
'workerGeneration',
|
||||
'leaseTokenDigest',
|
||||
'leaseGeneration',
|
||||
'leaseVersion',
|
||||
'offerId',
|
||||
'executorHandle',
|
||||
] as const);
|
||||
|
||||
private readonly runningFreshnessMs: number;
|
||||
|
||||
constructor(
|
||||
private readonly pool: PostgresPool,
|
||||
private readonly attestations: WorkerExecutionAttestationRepository,
|
||||
options: Readonly<{ runningFreshnessMs?: number }> = {},
|
||||
) {
|
||||
const freshness =
|
||||
options.runningFreshnessMs ??
|
||||
REMOTE_WORKER_ATTESTATION_LIMITS.defaultRunningFreshnessMs;
|
||||
if (
|
||||
!Number.isSafeInteger(freshness) ||
|
||||
freshness < 1_000 ||
|
||||
freshness > REMOTE_WORKER_ATTESTATION_LIMITS.maxRunningFreshnessMs
|
||||
) {
|
||||
throw new RangeError('Remote Worker attestation freshness is invalid');
|
||||
}
|
||||
this.runningFreshnessMs = freshness;
|
||||
}
|
||||
|
||||
async inspect(
|
||||
target: ClusterControlRecoveryProbeTarget,
|
||||
context: ClusterControlRecoveryEvidenceInspectionContext,
|
||||
): Promise<ClusterControlRecoveryEvidence> {
|
||||
if (context.signal.aborted) {
|
||||
return Object.freeze({ status: 'unknown', reason: 'provider_unavailable' });
|
||||
}
|
||||
const attestation = await this.attestations.findLatestExact({
|
||||
runId: target.runId,
|
||||
attemptId: target.attemptId,
|
||||
workerId: target.workerId!,
|
||||
workerSessionId: target.workerSessionId!,
|
||||
workerGeneration: target.workerGeneration!,
|
||||
leaseTokenDigest: target.leaseTokenDigest!,
|
||||
leaseGeneration: target.leaseGeneration!,
|
||||
leaseVersion: target.leaseVersion!,
|
||||
offerId: target.offerId!,
|
||||
callbackSequence: target.callbackSequence,
|
||||
executorHandle: target.executorHandle!,
|
||||
});
|
||||
if (!attestation) {
|
||||
return Object.freeze({ status: 'unknown', reason: 'provider_unavailable' });
|
||||
}
|
||||
if (attestation.state === 'stopped') {
|
||||
return Object.freeze({ status: 'not_running' });
|
||||
}
|
||||
const observed = await this.pool.query<Row>(
|
||||
`SELECT floor(extract(epoch FROM statement_timestamp()) * 1000)::bigint AS "observedAtMs"`,
|
||||
);
|
||||
const raw = observed.rows[0]?.observedAtMs;
|
||||
const observedAtMs =
|
||||
typeof raw === 'string' && /^\d+$/.test(raw) ? Number(raw) : raw;
|
||||
if (
|
||||
observed.rows.length !== 1 ||
|
||||
typeof observedAtMs !== 'number' ||
|
||||
!Number.isSafeInteger(observedAtMs) ||
|
||||
observedAtMs < attestation.receivedAtMs ||
|
||||
observedAtMs - attestation.receivedAtMs > this.runningFreshnessMs
|
||||
) {
|
||||
return Object.freeze({ status: 'unknown', reason: 'provider_unavailable' });
|
||||
}
|
||||
return Object.freeze({ status: 'running' });
|
||||
}
|
||||
}
|
||||
+1117
File diff suppressed because it is too large
Load Diff
+584
@@ -0,0 +1,584 @@
|
||||
// PostgreSQL Remote Worker lease control is owned by this domain.
|
||||
import type { PostgresClient, PostgresPool } from '@qinglong/runtime-core';
|
||||
import { digestRunDispatchLeaseToken } from '@qinglong/runtime-core';
|
||||
import {
|
||||
InvalidRemoteWorkerLeaseControlError,
|
||||
RemoteWorkerLeaseControlFenceRejectedError,
|
||||
RemoteWorkerLeaseControlUnavailableError,
|
||||
assertRemoteWorkerLeaseControlDuration,
|
||||
normalizeRemoteWorkerLeaseControlCommand,
|
||||
normalizeRemoteWorkerLeaseControlResult,
|
||||
type RemoteWorkerLeaseControlCommand,
|
||||
type RemoteWorkerLeaseControlRepository,
|
||||
type RemoteWorkerLeaseControlResult,
|
||||
type RemoteWorkerStopReason,
|
||||
type RemoteWorkerTerminalStatus,
|
||||
} from '@qinglong/runtime-core/remote-worker-lease-control';
|
||||
import {
|
||||
normalizeStepRunRecord,
|
||||
type StepRunRecord,
|
||||
} from '@qinglong/runtime-core/step-run';
|
||||
import { lockAttemptAuthority } from '../run/attemptAuthorityLock';
|
||||
|
||||
type Row = Record<string, unknown>;
|
||||
type MutationCommand = RemoteWorkerLeaseControlCommand & Readonly<{
|
||||
leaseDurationMs: number;
|
||||
timeoutEventId: string;
|
||||
}>;
|
||||
|
||||
const TERMINAL = new Set<RemoteWorkerTerminalStatus>([
|
||||
'succeeded', 'failed', 'cancelled', 'timed_out', 'lost',
|
||||
]);
|
||||
const STOP_REASONS = new Set<RemoteWorkerStopReason>([
|
||||
'user', 'policy', 'shutdown', 'reconcile', 'timeout',
|
||||
]);
|
||||
|
||||
function text(row: Row, key: string): string {
|
||||
const value = row[key];
|
||||
if (typeof value !== 'string' || value.length < 1) {
|
||||
throw new TypeError(`PostgreSQL Worker lease control ${key} is invalid`);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
function integer(row: Row, key: string): number {
|
||||
const raw = row[key];
|
||||
const value = typeof raw === 'string' && /^(0|[1-9]\d*)$/.test(raw)
|
||||
? Number(raw)
|
||||
: raw;
|
||||
if (typeof value !== 'number' || !Number.isSafeInteger(value) || value < 0) {
|
||||
throw new TypeError(`PostgreSQL Worker lease control ${key} is invalid`);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
function optionalInteger(row: Row, key: string): number | undefined {
|
||||
return row[key] === null || row[key] === undefined
|
||||
? undefined
|
||||
: integer(row, key);
|
||||
}
|
||||
|
||||
function optionalText(row: Row, key: string): string | undefined {
|
||||
return row[key] === null || row[key] === undefined
|
||||
? undefined
|
||||
: text(row, key);
|
||||
}
|
||||
|
||||
function isWorkflowTaskAttempt(aggregate: Row): boolean {
|
||||
return (
|
||||
aggregate.workflowAttemptId !== null &&
|
||||
aggregate.workflowAttemptId !== undefined
|
||||
);
|
||||
}
|
||||
|
||||
function workflowStepRun(aggregate: Row): Readonly<StepRunRecord> {
|
||||
try {
|
||||
const value = normalizeStepRunRecord(
|
||||
aggregate.workflowStepJson as StepRunRecord,
|
||||
);
|
||||
if (
|
||||
value.runId !== aggregate.runId ||
|
||||
value.id !== aggregate.workflowStepRunId ||
|
||||
value.id !== aggregate.attemptStepRunId ||
|
||||
value.version !== integer(aggregate, 'workflowStepVersion') ||
|
||||
value.stepRunDigest !== aggregate.workflowStepDigest
|
||||
) {
|
||||
throw new TypeError();
|
||||
}
|
||||
return value;
|
||||
} catch {
|
||||
throw new TypeError(
|
||||
'PostgreSQL Worker lease control Workflow StepRun is invalid',
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
function workflowStepAtAdmissionEpoch(aggregate: Row): boolean {
|
||||
const stepRun = workflowStepRun(aggregate);
|
||||
return (
|
||||
stepRun.version === integer(aggregate, 'admittedWorkflowStepVersion') &&
|
||||
stepRun.stepRunDigest === aggregate.admittedWorkflowStepDigest
|
||||
);
|
||||
}
|
||||
|
||||
function reject(
|
||||
command: Pick<RemoteWorkerLeaseControlCommand, 'attemptId'>,
|
||||
reason: ConstructorParameters<
|
||||
typeof RemoteWorkerLeaseControlFenceRejectedError
|
||||
>[1],
|
||||
): never {
|
||||
throw new RemoteWorkerLeaseControlFenceRejectedError(
|
||||
command.attemptId,
|
||||
reason,
|
||||
);
|
||||
}
|
||||
|
||||
function normalizeMutationCommand(value: MutationCommand): MutationCommand {
|
||||
if (!value || typeof value !== 'object' || Array.isArray(value)) {
|
||||
throw new InvalidRemoteWorkerLeaseControlError(
|
||||
'lease control mutation is invalid',
|
||||
);
|
||||
}
|
||||
const { leaseDurationMs, timeoutEventId, ...wire } = value;
|
||||
const command = normalizeRemoteWorkerLeaseControlCommand(wire);
|
||||
assertRemoteWorkerLeaseControlDuration(leaseDurationMs);
|
||||
if (
|
||||
typeof timeoutEventId !== 'string' || timeoutEventId.length < 1 ||
|
||||
timeoutEventId.length > 36 || /[\u0000-\u001f\u007f]/.test(timeoutEventId)
|
||||
) {
|
||||
throw new InvalidRemoteWorkerLeaseControlError(
|
||||
'timeout event identity is invalid',
|
||||
);
|
||||
}
|
||||
return Object.freeze({ ...command, leaseDurationMs, timeoutEventId });
|
||||
}
|
||||
|
||||
async function begin(client: PostgresClient): Promise<void> {
|
||||
await client.query('BEGIN');
|
||||
await client.query("SET LOCAL statement_timeout = '5s'");
|
||||
await client.query("SET LOCAL lock_timeout = '1s'");
|
||||
await client.query("SET LOCAL idle_in_transaction_session_timeout = '10s'");
|
||||
}
|
||||
|
||||
async function databaseNow(client: PostgresClient): Promise<number> {
|
||||
const result = await client.query<Row>(`
|
||||
SELECT floor(extract(epoch FROM statement_timestamp()) * 1000)::bigint AS "nowMs"
|
||||
`);
|
||||
if (result.rows.length !== 1) {
|
||||
throw new TypeError('PostgreSQL Worker lease control clock is invalid');
|
||||
}
|
||||
return integer(result.rows[0]!, 'nowMs');
|
||||
}
|
||||
|
||||
function common(
|
||||
command: RemoteWorkerLeaseControlCommand,
|
||||
): Pick<
|
||||
RemoteWorkerLeaseControlResult,
|
||||
'projectId' | 'runId' | 'attemptId' | 'offerId' | 'leaseGeneration'
|
||||
> {
|
||||
return {
|
||||
projectId: command.projectId,
|
||||
runId: command.runId,
|
||||
attemptId: command.attemptId,
|
||||
offerId: command.offerId,
|
||||
leaseGeneration: command.leaseGeneration,
|
||||
};
|
||||
}
|
||||
|
||||
function assertIdentity(
|
||||
command: RemoteWorkerLeaseControlCommand,
|
||||
worker: Row | undefined,
|
||||
aggregate: Row | undefined,
|
||||
lease: Row | undefined,
|
||||
observedAtMs: number,
|
||||
): asserts aggregate is Row {
|
||||
if (!aggregate || !lease) reject(command, 'missing');
|
||||
if (!worker) reject(command, 'worker_unavailable');
|
||||
if (worker.workerId !== command.workerId) reject(command, 'worker_mismatch');
|
||||
if (worker.sessionId !== command.workerSessionId) {
|
||||
reject(command, 'worker_session_mismatch');
|
||||
}
|
||||
if (integer(worker, 'workerGeneration') !== command.workerGeneration) {
|
||||
reject(command, 'worker_generation_mismatch');
|
||||
}
|
||||
if (
|
||||
!['online', 'draining'].includes(String(worker.workerStatus)) ||
|
||||
integer(worker, 'workerLeaseExpiresAtMs') <= observedAtMs
|
||||
) reject(command, 'worker_unavailable');
|
||||
const digest = digestRunDispatchLeaseToken(command.leaseToken);
|
||||
if (
|
||||
aggregate.runId !== command.runId ||
|
||||
aggregate.attemptRunId !== command.runId ||
|
||||
lease.runId !== command.runId
|
||||
) reject(command, 'run_mismatch');
|
||||
if (aggregate.projectId !== command.projectId) {
|
||||
reject(command, 'project_mismatch');
|
||||
}
|
||||
if (aggregate.executionOwner !== 'runtime') {
|
||||
reject(command, 'execution_owner_mismatch');
|
||||
}
|
||||
if (aggregate.executorType !== 'remote_worker') {
|
||||
reject(command, 'executor_mismatch');
|
||||
}
|
||||
if (
|
||||
lease.workerId !== command.workerId ||
|
||||
aggregate.attemptWorkerId !== command.workerId
|
||||
) reject(command, 'worker_mismatch');
|
||||
if (
|
||||
lease.workerSessionId !== command.workerSessionId ||
|
||||
aggregate.attemptWorkerSessionId !== command.workerSessionId
|
||||
) reject(command, 'worker_session_mismatch');
|
||||
if (
|
||||
integer(lease, 'workerGeneration') !== command.workerGeneration ||
|
||||
integer(aggregate, 'attemptWorkerGeneration') !== command.workerGeneration
|
||||
) reject(command, 'worker_generation_mismatch');
|
||||
if (
|
||||
integer(lease, 'leaseGeneration') !== command.leaseGeneration ||
|
||||
integer(aggregate, 'attemptLeaseGeneration') !== command.leaseGeneration
|
||||
) reject(command, 'lease_generation_mismatch');
|
||||
if (
|
||||
lease.leaseTokenDigest !== digest ||
|
||||
aggregate.attemptLeaseTokenDigest !== digest
|
||||
) reject(command, 'lease_token_mismatch');
|
||||
if (
|
||||
lease.offerId !== command.offerId ||
|
||||
aggregate.attemptOfferId !== command.offerId
|
||||
) reject(command, 'offer_mismatch');
|
||||
if (isWorkflowTaskAttempt(aggregate)) {
|
||||
workflowStepRun(aggregate);
|
||||
} else if (
|
||||
aggregate.attemptStepRunId !== null &&
|
||||
aggregate.attemptStepRunId !== undefined
|
||||
) {
|
||||
reject(command, 'run_mismatch');
|
||||
}
|
||||
}
|
||||
|
||||
export class PostgresRemoteWorkerLeaseControlRepository
|
||||
implements RemoteWorkerLeaseControlRepository {
|
||||
constructor(private readonly pool: PostgresPool) {
|
||||
if (!pool || typeof pool.connect !== 'function') {
|
||||
throw new TypeError('PostgreSQL Worker lease control pool is invalid');
|
||||
}
|
||||
}
|
||||
|
||||
async control(
|
||||
value: MutationCommand,
|
||||
): Promise<Readonly<RemoteWorkerLeaseControlResult>> {
|
||||
const command = normalizeMutationCommand(value);
|
||||
return this.transaction(async (client) => {
|
||||
await lockAttemptAuthority(client, command.attemptId);
|
||||
const worker = await client.query<Row>(`
|
||||
SELECT worker_id AS "workerId", session_id AS "sessionId",
|
||||
generation AS "workerGeneration", status AS "workerStatus",
|
||||
lease_expires_at_ms AS "workerLeaseExpiresAtMs"
|
||||
FROM "ql3"."worker_sessions" WHERE worker_id = $1 FOR UPDATE
|
||||
`, [command.workerId]);
|
||||
const aggregate = await client.query<Row>(`
|
||||
SELECT run.id AS "runId",
|
||||
COALESCE(workflow_task.project_id, run.project_id)
|
||||
AS "projectId",
|
||||
run.status AS "runStatus", run.execution_owner AS "executionOwner",
|
||||
run.cancel_requested_at_ms AS "cancelRequestedAtMs",
|
||||
run.cancel_reason AS "cancelReason", run.version AS "runVersion",
|
||||
run.event_sequence AS "eventSequence",
|
||||
attempt.id AS "attemptId", attempt.run_id AS "attemptRunId",
|
||||
attempt.step_run_id AS "attemptStepRunId",
|
||||
attempt.status AS "attemptStatus",
|
||||
attempt.executor_type AS "executorType",
|
||||
attempt.worker_id AS "attemptWorkerId",
|
||||
attempt.worker_session_id AS "attemptWorkerSessionId",
|
||||
attempt.worker_generation AS "attemptWorkerGeneration",
|
||||
attempt.lease_generation AS "attemptLeaseGeneration",
|
||||
attempt.lease_version AS "attemptLeaseVersion",
|
||||
attempt.lease_token_digest AS "attemptLeaseTokenDigest",
|
||||
attempt.offer_id AS "attemptOfferId",
|
||||
attempt.deadline_at_ms AS "deadlineAtMs",
|
||||
workflow_task.attempt_id AS "workflowAttemptId",
|
||||
workflow_task.step_run_id AS "workflowStepRunId",
|
||||
workflow_task.step_run_version AS
|
||||
"admittedWorkflowStepVersion",
|
||||
workflow_task.step_run_digest AS
|
||||
"admittedWorkflowStepDigest"
|
||||
FROM "ql3"."runs" AS run
|
||||
INNER JOIN "ql3"."run_attempts" AS attempt ON attempt.id = $2
|
||||
LEFT JOIN
|
||||
"ql3"."plugin_package_workflow_task_attempt_admissions"
|
||||
AS workflow_task
|
||||
ON workflow_task.attempt_id = attempt.id
|
||||
WHERE run.id = $1 FOR UPDATE OF run, attempt
|
||||
`, [command.runId, command.attemptId]);
|
||||
const aggregateRow = aggregate.rows[0];
|
||||
let lockedAggregate = aggregateRow;
|
||||
if (aggregateRow && isWorkflowTaskAttempt(aggregateRow)) {
|
||||
const step = await client.query<Row>(
|
||||
`SELECT version AS "workflowStepVersion",
|
||||
step_run_digest AS "workflowStepDigest",
|
||||
step_run_json AS "workflowStepJson"
|
||||
FROM "ql3"."step_runs"
|
||||
WHERE run_id = $1 AND id = $2
|
||||
FOR UPDATE`,
|
||||
[command.runId, aggregateRow.workflowStepRunId],
|
||||
);
|
||||
lockedAggregate = Object.freeze({
|
||||
...aggregateRow,
|
||||
...(step.rows[0] ?? {}),
|
||||
});
|
||||
}
|
||||
const lease = await client.query<Row>(`
|
||||
SELECT run_id AS "runId", status AS "leaseStatus",
|
||||
version AS "leaseVersion", lease_generation AS "leaseGeneration",
|
||||
worker_id AS "workerId", worker_session_id AS "workerSessionId",
|
||||
worker_generation AS "workerGeneration",
|
||||
lease_token_digest AS "leaseTokenDigest", offer_id AS "offerId",
|
||||
expires_at_ms AS "leaseExpiresAtMs"
|
||||
FROM "ql3"."run_dispatch_leases"
|
||||
WHERE attempt_id = $1 FOR UPDATE
|
||||
`, [command.attemptId]);
|
||||
const observedAtMs = await databaseNow(client);
|
||||
assertIdentity(
|
||||
command,
|
||||
worker.rows[0],
|
||||
lockedAggregate,
|
||||
lease.rows[0],
|
||||
observedAtMs,
|
||||
);
|
||||
const state = lockedAggregate!;
|
||||
const currentLease = lease.rows[0]!;
|
||||
const runStatus = String(state.runStatus);
|
||||
const attemptStatus = String(state.attemptStatus);
|
||||
const workflowTask = isWorkflowTaskAttempt(state);
|
||||
const stepStatus = workflowTask
|
||||
? workflowStepRun(state).status
|
||||
: undefined;
|
||||
const attemptTerminal = TERMINAL.has(
|
||||
attemptStatus as RemoteWorkerTerminalStatus,
|
||||
);
|
||||
const scopeTerminal = workflowTask
|
||||
? TERMINAL.has(stepStatus as RemoteWorkerTerminalStatus)
|
||||
: TERMINAL.has(runStatus as RemoteWorkerTerminalStatus);
|
||||
if (
|
||||
scopeTerminal ||
|
||||
attemptTerminal
|
||||
) {
|
||||
if (
|
||||
(workflowTask ? stepStatus !== attemptStatus : runStatus !== attemptStatus) ||
|
||||
!scopeTerminal ||
|
||||
!attemptTerminal ||
|
||||
currentLease.leaseStatus !== 'completed'
|
||||
) reject(command, 'state_mismatch');
|
||||
return normalizeRemoteWorkerLeaseControlResult({
|
||||
status: 'terminal',
|
||||
...common(command),
|
||||
terminalStatus: attemptStatus as RemoteWorkerTerminalStatus,
|
||||
});
|
||||
}
|
||||
if (
|
||||
currentLease.leaseStatus !== 'leased' ||
|
||||
integer(currentLease, 'leaseVersion') !== command.expectedLeaseVersion ||
|
||||
integer(state, 'attemptLeaseVersion') !== command.expectedLeaseVersion
|
||||
) reject(command, 'version_mismatch');
|
||||
if (integer(currentLease, 'leaseExpiresAtMs') <= observedAtMs) {
|
||||
reject(command, 'lease_expired');
|
||||
}
|
||||
if (
|
||||
!['claimed', 'starting', 'running'].includes(attemptStatus) ||
|
||||
(workflowTask
|
||||
? runStatus !== 'running' ||
|
||||
(attemptStatus === 'running'
|
||||
? stepStatus !== 'running'
|
||||
: stepStatus !== 'ready' ||
|
||||
!workflowStepAtAdmissionEpoch(state))
|
||||
: !['dispatching', 'running'].includes(runStatus) ||
|
||||
(attemptStatus === 'running') !== (runStatus === 'running'))
|
||||
) reject(command, 'state_mismatch');
|
||||
|
||||
let cancelRequestedAtMs = optionalInteger(state, 'cancelRequestedAtMs');
|
||||
let cancelReason = optionalText(state, 'cancelReason');
|
||||
let workflowTimeoutRequestedAtMs: number | undefined;
|
||||
const deadlineAtMs = optionalInteger(state, 'deadlineAtMs');
|
||||
if (
|
||||
cancelRequestedAtMs === undefined &&
|
||||
deadlineAtMs !== undefined &&
|
||||
deadlineAtMs <= observedAtMs
|
||||
) {
|
||||
const runVersion = integer(state, 'runVersion');
|
||||
const eventSequence = integer(state, 'eventSequence') + 1;
|
||||
if (runVersion >= 2_147_483_647 || eventSequence > 2_147_483_647) {
|
||||
throw new RangeError('Worker lease control Run counter overflowed');
|
||||
}
|
||||
const timeoutDedupeKey =
|
||||
`remote-timeout:${command.attemptId}:${command.leaseGeneration}`;
|
||||
if (workflowTask) {
|
||||
const existing = await client.query<Row>(
|
||||
`SELECT created_at_ms AS "createdAtMs"
|
||||
FROM "ql3"."run_events"
|
||||
WHERE run_id = $1 AND dedupe_key = $2`,
|
||||
[command.runId, timeoutDedupeKey],
|
||||
);
|
||||
if (existing.rows.length > 1) {
|
||||
throw new TypeError(
|
||||
'PostgreSQL Worker Workflow timeout event is invalid',
|
||||
);
|
||||
}
|
||||
if (existing.rows.length === 1) {
|
||||
workflowTimeoutRequestedAtMs = integer(
|
||||
existing.rows[0]!,
|
||||
'createdAtMs',
|
||||
);
|
||||
} else {
|
||||
const updated = await client.query(`
|
||||
UPDATE "ql3"."runs"
|
||||
SET version = $2, event_sequence = $3
|
||||
WHERE id = $1 AND version = $4
|
||||
AND cancel_requested_at_ms IS NULL
|
||||
AND status = 'running'
|
||||
`, [
|
||||
command.runId,
|
||||
runVersion + 1,
|
||||
eventSequence,
|
||||
runVersion,
|
||||
]);
|
||||
if (updated.rowCount !== 1) reject(command, 'version_mismatch');
|
||||
await client.query(`
|
||||
INSERT INTO "ql3"."run_events" (
|
||||
id, run_id, sequence, type, dedupe_key, actor_type, actor_id,
|
||||
attempt_id, step_run_id, payload, created_at_ms
|
||||
) VALUES (
|
||||
$1, $2, $3, 'workflow.task_timeout_requested', $4, 'system',
|
||||
'runtime:timeout', $5, $6, $7::jsonb, $8
|
||||
)
|
||||
`, [
|
||||
command.timeoutEventId,
|
||||
command.runId,
|
||||
eventSequence,
|
||||
timeoutDedupeKey,
|
||||
command.attemptId,
|
||||
text(state, 'workflowStepRunId'),
|
||||
JSON.stringify({
|
||||
attempt_id: command.attemptId,
|
||||
step_run_id: text(state, 'workflowStepRunId'),
|
||||
lease_generation: command.leaseGeneration,
|
||||
execution_scope: 'workflow_task',
|
||||
reason: 'timeout',
|
||||
deadline_at_ms: deadlineAtMs,
|
||||
}),
|
||||
observedAtMs,
|
||||
]);
|
||||
workflowTimeoutRequestedAtMs = observedAtMs;
|
||||
}
|
||||
} else {
|
||||
const updated = await client.query(`
|
||||
UPDATE "ql3"."runs"
|
||||
SET cancel_requested_at_ms = $2, cancel_reason = 'timeout',
|
||||
version = $3, event_sequence = $4
|
||||
WHERE id = $1 AND version = $5 AND cancel_requested_at_ms IS NULL
|
||||
AND status IN ('dispatching', 'running')
|
||||
`, [
|
||||
command.runId,
|
||||
observedAtMs,
|
||||
runVersion + 1,
|
||||
eventSequence,
|
||||
runVersion,
|
||||
]);
|
||||
if (updated.rowCount !== 1) reject(command, 'version_mismatch');
|
||||
await client.query(`
|
||||
INSERT INTO "ql3"."run_events" (
|
||||
id, run_id, sequence, type, dedupe_key, actor_type, actor_id,
|
||||
attempt_id, step_run_id, payload, created_at_ms
|
||||
) VALUES ($1, $2, $3, 'run.cancel_requested', $4, 'system',
|
||||
'runtime:timeout', $5, NULL, $6::jsonb, $7)
|
||||
`, [
|
||||
command.timeoutEventId,
|
||||
command.runId,
|
||||
eventSequence,
|
||||
timeoutDedupeKey,
|
||||
command.attemptId,
|
||||
JSON.stringify({
|
||||
attempt_id: command.attemptId,
|
||||
lease_generation: command.leaseGeneration,
|
||||
reason: 'timeout',
|
||||
deadline_at_ms: deadlineAtMs,
|
||||
}),
|
||||
observedAtMs,
|
||||
]);
|
||||
cancelRequestedAtMs = observedAtMs;
|
||||
cancelReason = 'timeout';
|
||||
}
|
||||
}
|
||||
if (
|
||||
(cancelRequestedAtMs === undefined) !== (cancelReason === undefined) ||
|
||||
(cancelReason !== undefined &&
|
||||
!STOP_REASONS.has(cancelReason as RemoteWorkerStopReason))
|
||||
) {
|
||||
throw new TypeError('PostgreSQL Worker cancellation intent is invalid');
|
||||
}
|
||||
const stopReason =
|
||||
cancelReason ??
|
||||
(workflowTimeoutRequestedAtMs === undefined ? undefined : 'timeout');
|
||||
const stopRequestedAtMs =
|
||||
cancelRequestedAtMs ?? workflowTimeoutRequestedAtMs;
|
||||
|
||||
const nextVersion = command.expectedLeaseVersion + 1;
|
||||
if (nextVersion > 2_147_483_647) {
|
||||
throw new RangeError('Worker lease control version overflowed');
|
||||
}
|
||||
const expiresAtMs = observedAtMs + command.leaseDurationMs;
|
||||
const renewed = await client.query<Row>(`
|
||||
UPDATE "ql3"."run_dispatch_leases"
|
||||
SET version = $2, renewed_at_ms = $3, expires_at_ms = $4,
|
||||
updated_at_ms = $3
|
||||
WHERE attempt_id = $1 AND status = 'leased' AND version = $5
|
||||
RETURNING renewed_at_ms AS "renewedAtMs", expires_at_ms AS "expiresAtMs"
|
||||
`, [
|
||||
command.attemptId,
|
||||
nextVersion,
|
||||
observedAtMs,
|
||||
expiresAtMs,
|
||||
command.expectedLeaseVersion,
|
||||
]);
|
||||
const attempt = await client.query(`
|
||||
UPDATE "ql3"."run_attempts"
|
||||
SET lease_version = $2, lease_expires_at_ms = $3
|
||||
WHERE id = $1 AND lease_version = $4
|
||||
AND status IN ('claimed', 'starting', 'running')
|
||||
`, [
|
||||
command.attemptId,
|
||||
nextVersion,
|
||||
expiresAtMs,
|
||||
command.expectedLeaseVersion,
|
||||
]);
|
||||
if (renewed.rows.length !== 1 || attempt.rowCount !== 1) {
|
||||
reject(command, 'version_mismatch');
|
||||
}
|
||||
const result = {
|
||||
status: stopReason === undefined
|
||||
? 'renewed' as const
|
||||
: 'stop_requested' as const,
|
||||
...common(command),
|
||||
leaseVersion: nextVersion,
|
||||
renewedAtMs: integer(renewed.rows[0]!, 'renewedAtMs'),
|
||||
expiresAtMs: integer(renewed.rows[0]!, 'expiresAtMs'),
|
||||
...(stopReason === undefined
|
||||
? {}
|
||||
: {
|
||||
stop: Object.freeze({
|
||||
reason: stopReason as RemoteWorkerStopReason,
|
||||
requestedAtMs: stopRequestedAtMs!,
|
||||
}),
|
||||
}),
|
||||
};
|
||||
return normalizeRemoteWorkerLeaseControlResult(result);
|
||||
});
|
||||
}
|
||||
|
||||
private async transaction<T>(
|
||||
work: (client: PostgresClient) => Promise<T>,
|
||||
): Promise<T> {
|
||||
let client: PostgresClient;
|
||||
try {
|
||||
client = await this.pool.connect();
|
||||
} catch (error) {
|
||||
throw new RemoteWorkerLeaseControlUnavailableError({ cause: error });
|
||||
}
|
||||
try {
|
||||
await begin(client);
|
||||
const value = await work(client);
|
||||
await client.query('COMMIT');
|
||||
return value;
|
||||
} catch (error) {
|
||||
try {
|
||||
await client.query('ROLLBACK');
|
||||
} catch {
|
||||
// Preserve the originating failure.
|
||||
}
|
||||
if (
|
||||
error instanceof RemoteWorkerLeaseControlFenceRejectedError ||
|
||||
error instanceof InvalidRemoteWorkerLeaseControlError
|
||||
) throw error;
|
||||
throw new RemoteWorkerLeaseControlUnavailableError({ cause: error });
|
||||
} finally {
|
||||
client.release();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,739 @@
|
||||
// PostgreSQL Remote Execution dispatch leases are owned by this domain.
|
||||
import type {
|
||||
ClaimRunDispatchLeaseCommand,
|
||||
ClaimRunDispatchLeaseResult,
|
||||
PostgresClient,
|
||||
PostgresPool,
|
||||
ReleaseRunDispatchLeaseCommand,
|
||||
RenewRunDispatchLeaseCommand,
|
||||
RunDispatchLeaseRecord,
|
||||
RunDispatchLeaseRepository,
|
||||
RunDispatchLeaseStatus,
|
||||
} from '@qinglong/runtime-core';
|
||||
import {
|
||||
RUN_DISPATCH_LEASE_STATUSES,
|
||||
RunDispatchLeaseFenceRejectedError,
|
||||
assertRunDispatchId,
|
||||
assertRunDispatchLeaseDuration,
|
||||
assertRunDispatchLeaseFence,
|
||||
assertRunDispatchLeaseRecord,
|
||||
digestRunDispatchLeaseToken,
|
||||
} from '@qinglong/runtime-core';
|
||||
import { lockAttemptAuthority } from '../run/attemptAuthorityLock';
|
||||
|
||||
type Row = Record<string, unknown>;
|
||||
|
||||
const LEASE_COLUMNS = `
|
||||
attempt_id AS "attemptId",
|
||||
run_id AS "runId",
|
||||
status AS "status",
|
||||
version AS "version",
|
||||
lease_generation AS "leaseGeneration",
|
||||
worker_id AS "workerId",
|
||||
worker_session_id AS "workerSessionId",
|
||||
worker_generation AS "workerGeneration",
|
||||
lease_token_digest AS "leaseTokenDigest",
|
||||
acquired_at_ms AS "acquiredAtMs",
|
||||
renewed_at_ms AS "renewedAtMs",
|
||||
expires_at_ms AS "expiresAtMs",
|
||||
released_at_ms AS "releasedAtMs",
|
||||
release_reason AS "releaseReason",
|
||||
completed_at_ms AS "completedAtMs",
|
||||
updated_at_ms AS "updatedAtMs"
|
||||
`.trim();
|
||||
|
||||
function string(row: Row, key: string): string {
|
||||
const value = row[key];
|
||||
if (typeof value !== 'string' || value.length === 0) {
|
||||
throw new TypeError(`PostgreSQL Run dispatch lease ${key} is invalid`);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
function integer(row: Row, key: string): number {
|
||||
const value = row[key];
|
||||
const normalized =
|
||||
typeof value === 'string' && /^(0|[1-9]\d*)$/.test(value)
|
||||
? Number(value)
|
||||
: value;
|
||||
if (typeof normalized !== 'number' || !Number.isSafeInteger(normalized)) {
|
||||
throw new TypeError(`PostgreSQL Run dispatch lease ${key} is invalid`);
|
||||
}
|
||||
return normalized;
|
||||
}
|
||||
|
||||
function optionalInteger(row: Row, key: string): number | undefined {
|
||||
return row[key] === null || row[key] === undefined ? undefined : integer(row, key);
|
||||
}
|
||||
|
||||
function optionalString(row: Row, key: string): string | undefined {
|
||||
return row[key] === null || row[key] === undefined ? undefined : string(row, key);
|
||||
}
|
||||
|
||||
function lease(row: Row): RunDispatchLeaseRecord {
|
||||
const status = string(row, 'status');
|
||||
if (!RUN_DISPATCH_LEASE_STATUSES.includes(status as RunDispatchLeaseStatus)) {
|
||||
throw new TypeError('PostgreSQL Run dispatch lease status is invalid');
|
||||
}
|
||||
const releasedAtMs = optionalInteger(row, 'releasedAtMs');
|
||||
const releaseReason = optionalString(row, 'releaseReason');
|
||||
const completedAtMs = optionalInteger(row, 'completedAtMs');
|
||||
const value: RunDispatchLeaseRecord = Object.freeze({
|
||||
attemptId: string(row, 'attemptId'),
|
||||
runId: string(row, 'runId'),
|
||||
status: status as RunDispatchLeaseStatus,
|
||||
version: integer(row, 'version'),
|
||||
leaseGeneration: integer(row, 'leaseGeneration'),
|
||||
workerId: string(row, 'workerId'),
|
||||
workerSessionId: string(row, 'workerSessionId'),
|
||||
workerGeneration: integer(row, 'workerGeneration'),
|
||||
leaseTokenDigest: string(row, 'leaseTokenDigest'),
|
||||
acquiredAtMs: integer(row, 'acquiredAtMs'),
|
||||
renewedAtMs: integer(row, 'renewedAtMs'),
|
||||
expiresAtMs: integer(row, 'expiresAtMs'),
|
||||
updatedAtMs: integer(row, 'updatedAtMs'),
|
||||
...(releasedAtMs === undefined ? {} : { releasedAtMs }),
|
||||
...(releaseReason === undefined
|
||||
? {}
|
||||
: {
|
||||
releaseReason: releaseReason as NonNullable<
|
||||
RunDispatchLeaseRecord['releaseReason']
|
||||
>,
|
||||
}),
|
||||
...(completedAtMs === undefined ? {} : { completedAtMs }),
|
||||
});
|
||||
assertRunDispatchLeaseRecord(value);
|
||||
return value;
|
||||
}
|
||||
|
||||
function bounded(name: string, value: string): void {
|
||||
assertRunDispatchId(name, value);
|
||||
}
|
||||
|
||||
function assertClaim(command: ClaimRunDispatchLeaseCommand): void {
|
||||
bounded('runId', command.runId);
|
||||
bounded('attemptId', command.attemptId);
|
||||
bounded('eventId', command.eventId);
|
||||
bounded('offerId', command.offerId);
|
||||
assertRunDispatchLeaseFence({
|
||||
...command,
|
||||
leaseGeneration: 1,
|
||||
expectedVersion: 0,
|
||||
});
|
||||
assertRunDispatchLeaseDuration(command.leaseDurationMs);
|
||||
}
|
||||
|
||||
function assertRenew(command: RenewRunDispatchLeaseCommand): void {
|
||||
bounded('attemptId', command.attemptId);
|
||||
assertRunDispatchLeaseFence(command);
|
||||
assertRunDispatchLeaseDuration(command.leaseDurationMs);
|
||||
}
|
||||
|
||||
function assertRelease(command: ReleaseRunDispatchLeaseCommand): void {
|
||||
bounded('runId', command.runId);
|
||||
bounded('attemptId', command.attemptId);
|
||||
bounded('eventId', command.eventId);
|
||||
assertRunDispatchLeaseFence(command);
|
||||
if (
|
||||
!['declined', 'shutdown', 'start_failed', 'capacity_changed'].includes(
|
||||
command.reason,
|
||||
)
|
||||
) {
|
||||
throw new TypeError('Run dispatch release reason is invalid');
|
||||
}
|
||||
}
|
||||
|
||||
async function begin(client: PostgresClient): Promise<void> {
|
||||
await client.query('BEGIN');
|
||||
await client.query("SET LOCAL statement_timeout = '5s'");
|
||||
await client.query("SET LOCAL lock_timeout = '1s'");
|
||||
await client.query("SET LOCAL idle_in_transaction_session_timeout = '10s'");
|
||||
}
|
||||
|
||||
async function now(client: PostgresClient): Promise<number> {
|
||||
const result = await client.query<Row>(`
|
||||
SELECT floor(extract(epoch FROM statement_timestamp()) * 1000)::bigint AS "nowMs"
|
||||
`);
|
||||
if (result.rows.length !== 1) {
|
||||
throw new TypeError('PostgreSQL Run lease observation is invalid');
|
||||
}
|
||||
return integer(result.rows[0]!, 'nowMs');
|
||||
}
|
||||
|
||||
function sameFence(
|
||||
current: RunDispatchLeaseRecord,
|
||||
command: {
|
||||
workerId: string;
|
||||
workerSessionId: string;
|
||||
workerGeneration: number;
|
||||
leaseGeneration: number;
|
||||
leaseToken: string;
|
||||
},
|
||||
): boolean {
|
||||
return (
|
||||
current.workerId === command.workerId &&
|
||||
current.workerSessionId === command.workerSessionId &&
|
||||
current.workerGeneration === command.workerGeneration &&
|
||||
current.leaseGeneration === command.leaseGeneration &&
|
||||
current.leaseTokenDigest === digestRunDispatchLeaseToken(command.leaseToken)
|
||||
);
|
||||
}
|
||||
|
||||
function assertLeaseFence(
|
||||
current: RunDispatchLeaseRecord | null,
|
||||
command: RenewRunDispatchLeaseCommand | ReleaseRunDispatchLeaseCommand,
|
||||
observedAtMs: number,
|
||||
): asserts current is RunDispatchLeaseRecord {
|
||||
if (!current) throw new RunDispatchLeaseFenceRejectedError(command.attemptId, 'missing');
|
||||
if ('runId' in command && current.runId !== command.runId) {
|
||||
throw new RunDispatchLeaseFenceRejectedError(command.attemptId, 'run_mismatch');
|
||||
}
|
||||
if (current.workerId !== command.workerId) {
|
||||
throw new RunDispatchLeaseFenceRejectedError(command.attemptId, 'worker_mismatch');
|
||||
}
|
||||
if (current.workerSessionId !== command.workerSessionId) {
|
||||
throw new RunDispatchLeaseFenceRejectedError(
|
||||
command.attemptId,
|
||||
'worker_session_mismatch',
|
||||
);
|
||||
}
|
||||
if (current.workerGeneration !== command.workerGeneration) {
|
||||
throw new RunDispatchLeaseFenceRejectedError(
|
||||
command.attemptId,
|
||||
'worker_generation_mismatch',
|
||||
);
|
||||
}
|
||||
if (current.leaseGeneration !== command.leaseGeneration) {
|
||||
throw new RunDispatchLeaseFenceRejectedError(
|
||||
command.attemptId,
|
||||
'lease_generation_mismatch',
|
||||
);
|
||||
}
|
||||
if (current.leaseTokenDigest !== digestRunDispatchLeaseToken(command.leaseToken)) {
|
||||
throw new RunDispatchLeaseFenceRejectedError(
|
||||
command.attemptId,
|
||||
'lease_token_mismatch',
|
||||
);
|
||||
}
|
||||
if (current.version !== command.expectedVersion) {
|
||||
throw new RunDispatchLeaseFenceRejectedError(command.attemptId, 'version_mismatch');
|
||||
}
|
||||
if (current.status !== 'leased') {
|
||||
throw new RunDispatchLeaseFenceRejectedError(command.attemptId, 'not_leased');
|
||||
}
|
||||
if (current.expiresAtMs <= observedAtMs) {
|
||||
throw new RunDispatchLeaseFenceRejectedError(command.attemptId, 'lease_expired');
|
||||
}
|
||||
}
|
||||
|
||||
function workerCurrent(
|
||||
row: Row | undefined,
|
||||
command: {
|
||||
workerId: string;
|
||||
workerSessionId: string;
|
||||
workerGeneration: number;
|
||||
},
|
||||
observedAtMs: number,
|
||||
allowDraining: boolean,
|
||||
): boolean {
|
||||
return Boolean(
|
||||
row &&
|
||||
row.workerId === command.workerId &&
|
||||
row.sessionId === command.workerSessionId &&
|
||||
integer(row, 'generation') === command.workerGeneration &&
|
||||
(row.status === 'online' || (allowDraining && row.status === 'draining')) &&
|
||||
integer(row, 'leaseExpiresAtMs') > observedAtMs,
|
||||
);
|
||||
}
|
||||
|
||||
export class PostgresRunDispatchLeaseRepository
|
||||
implements RunDispatchLeaseRepository
|
||||
{
|
||||
constructor(private readonly pool: PostgresPool) {}
|
||||
|
||||
async findByAttemptId(
|
||||
attemptId: string,
|
||||
): Promise<RunDispatchLeaseRecord | null> {
|
||||
bounded('attemptId', attemptId);
|
||||
const result = await this.pool.query<Row>(
|
||||
`SELECT ${LEASE_COLUMNS} FROM "ql3"."run_dispatch_leases" WHERE attempt_id = $1`,
|
||||
[attemptId],
|
||||
);
|
||||
if (result.rows.length > 1) {
|
||||
throw new TypeError('PostgreSQL Run lease lookup returned multiple rows');
|
||||
}
|
||||
return result.rows[0] ? lease(result.rows[0]) : null;
|
||||
}
|
||||
|
||||
async claim(
|
||||
command: ClaimRunDispatchLeaseCommand,
|
||||
): Promise<ClaimRunDispatchLeaseResult> {
|
||||
assertClaim(command);
|
||||
return this.transaction(async (client) => {
|
||||
await lockAttemptAuthority(client, command.attemptId);
|
||||
const workerResult = await client.query<Row>(
|
||||
`
|
||||
SELECT worker_id AS "workerId", session_id AS "sessionId",
|
||||
generation, status, max_concurrent_runs AS "maxConcurrentRuns",
|
||||
available_slots AS "availableSlots",
|
||||
lease_expires_at_ms AS "leaseExpiresAtMs"
|
||||
FROM "ql3"."worker_sessions" WHERE worker_id = $1 FOR UPDATE
|
||||
`,
|
||||
[command.workerId],
|
||||
);
|
||||
const worker = workerResult.rows[0];
|
||||
const aggregateResult = await client.query<Row>(
|
||||
`
|
||||
SELECT run.id AS "runId", run.status AS "runStatus",
|
||||
run.execution_owner AS "executionOwner",
|
||||
run.cancel_requested_at_ms AS "cancelRequestedAtMs",
|
||||
run.version AS "runVersion", run.event_sequence AS "eventSequence",
|
||||
attempt.id AS "attemptId", attempt.status AS "attemptStatus",
|
||||
attempt.run_id AS "attemptRunId",
|
||||
attempt.step_run_id AS "attemptStepRunId",
|
||||
workflow_task.attempt_id AS "workflowAttemptId",
|
||||
workflow_task.step_run_id AS "workflowStepRunId",
|
||||
workflow_task.step_run_version AS
|
||||
"admittedWorkflowStepVersion",
|
||||
workflow_task.step_run_digest AS
|
||||
"admittedWorkflowStepDigest"
|
||||
FROM "ql3"."runs" AS run
|
||||
INNER JOIN "ql3"."run_attempts" AS attempt ON attempt.id = $2
|
||||
LEFT JOIN
|
||||
"ql3"."plugin_package_workflow_task_attempt_admissions"
|
||||
AS workflow_task
|
||||
ON workflow_task.attempt_id = attempt.id
|
||||
WHERE run.id = $1
|
||||
FOR UPDATE OF run, attempt
|
||||
`,
|
||||
[command.runId, command.attemptId],
|
||||
);
|
||||
const aggregate = aggregateResult.rows[0];
|
||||
const workflowAttempt =
|
||||
aggregate?.workflowAttemptId !== null &&
|
||||
aggregate?.workflowAttemptId !== undefined;
|
||||
const workflowStep = workflowAttempt
|
||||
? (
|
||||
await client.query<Row>(
|
||||
`SELECT status AS "workflowStepStatus",
|
||||
version AS "workflowStepVersion",
|
||||
step_run_digest AS "workflowStepDigest"
|
||||
FROM "ql3"."step_runs"
|
||||
WHERE run_id = $1 AND id = $2
|
||||
FOR UPDATE`,
|
||||
[command.runId, aggregate!.workflowStepRunId],
|
||||
)
|
||||
).rows[0]
|
||||
: undefined;
|
||||
const workflowEligible = Boolean(
|
||||
workflowAttempt &&
|
||||
aggregate!.runStatus === 'running' &&
|
||||
aggregate!.attemptStepRunId === aggregate!.workflowStepRunId &&
|
||||
workflowStep?.workflowStepStatus === 'ready' &&
|
||||
integer(workflowStep!, 'workflowStepVersion') ===
|
||||
integer(aggregate!, 'admittedWorkflowStepVersion') &&
|
||||
workflowStep!.workflowStepDigest ===
|
||||
aggregate!.admittedWorkflowStepDigest,
|
||||
);
|
||||
const runEligible = Boolean(
|
||||
!workflowAttempt &&
|
||||
['queued', 'dispatching'].includes(String(aggregate?.runStatus)),
|
||||
);
|
||||
if (
|
||||
!aggregate ||
|
||||
aggregate.executionOwner !== 'runtime' ||
|
||||
aggregate.cancelRequestedAtMs !== null ||
|
||||
aggregate.attemptRunId !== command.runId ||
|
||||
aggregate.attemptStatus !== 'claimed' ||
|
||||
(!runEligible && !workflowEligible)
|
||||
) {
|
||||
return Object.freeze({ status: 'not_eligible' as const });
|
||||
}
|
||||
const currentResult = await client.query<Row>(
|
||||
`SELECT ${LEASE_COLUMNS}, offer_id AS "offerId" FROM "ql3"."run_dispatch_leases" WHERE attempt_id = $1 FOR UPDATE`,
|
||||
[command.attemptId],
|
||||
);
|
||||
const current = currentResult.rows[0] ? lease(currentResult.rows[0]) : null;
|
||||
const observedAtMs = await now(client);
|
||||
if (!workerCurrent(worker, command, observedAtMs, false)) {
|
||||
return Object.freeze({ status: 'worker_unavailable' as const });
|
||||
}
|
||||
if (current && current.runId !== command.runId) {
|
||||
throw new RunDispatchLeaseFenceRejectedError(command.attemptId, 'run_mismatch');
|
||||
}
|
||||
if (current?.status === 'leased' && current.expiresAtMs > observedAtMs) {
|
||||
if (
|
||||
sameFence(current, {
|
||||
...command,
|
||||
leaseGeneration: current.leaseGeneration,
|
||||
}) &&
|
||||
currentResult.rows[0]!.offerId === command.offerId
|
||||
) {
|
||||
return Object.freeze({ status: 'idempotent' as const, lease: current });
|
||||
}
|
||||
return Object.freeze({ status: 'leased' as const, lease: current });
|
||||
}
|
||||
if (current?.status === 'completed') {
|
||||
return Object.freeze({ status: 'not_eligible' as const });
|
||||
}
|
||||
const active = await client.query<Row>(
|
||||
`
|
||||
SELECT count(*)::integer AS "activeCount"
|
||||
FROM "ql3"."run_dispatch_leases"
|
||||
WHERE worker_id = $1 AND worker_session_id = $2
|
||||
AND worker_generation = $3 AND status = 'leased'
|
||||
AND expires_at_ms > $4 AND attempt_id <> $5
|
||||
`,
|
||||
[
|
||||
command.workerId,
|
||||
command.workerSessionId,
|
||||
command.workerGeneration,
|
||||
observedAtMs,
|
||||
command.attemptId,
|
||||
],
|
||||
);
|
||||
const activeCount = integer(active.rows[0]!, 'activeCount');
|
||||
if (
|
||||
activeCount >= integer(worker!, 'maxConcurrentRuns') ||
|
||||
activeCount >= integer(worker!, 'availableSlots')
|
||||
) {
|
||||
return Object.freeze({ status: 'capacity_exhausted' as const });
|
||||
}
|
||||
const leaseGeneration = (current?.leaseGeneration ?? 0) + 1;
|
||||
const version = current ? current.version + 1 : 0;
|
||||
if (leaseGeneration > 2_147_483_647 || version > 2_147_483_647) {
|
||||
throw new RangeError('Run dispatch lease generation or version overflowed');
|
||||
}
|
||||
const digest = digestRunDispatchLeaseToken(command.leaseToken);
|
||||
const expiresAtMs = observedAtMs + command.leaseDurationMs;
|
||||
const persisted = await client.query<Row>(
|
||||
`
|
||||
INSERT INTO "ql3"."run_dispatch_leases" (
|
||||
attempt_id, run_id, status, version, lease_generation,
|
||||
worker_id, worker_session_id, worker_generation,
|
||||
lease_token_digest, offer_id, acquired_at_ms, renewed_at_ms,
|
||||
expires_at_ms, released_at_ms, release_reason, completed_at_ms,
|
||||
updated_at_ms
|
||||
) VALUES ($1, $2, 'leased', $3, $4, $5, $6, $7, $8, $9, $10, $10, $11, NULL, NULL, NULL, $10)
|
||||
ON CONFLICT (attempt_id) DO UPDATE SET
|
||||
run_id = EXCLUDED.run_id,
|
||||
status = EXCLUDED.status,
|
||||
version = EXCLUDED.version,
|
||||
lease_generation = EXCLUDED.lease_generation,
|
||||
worker_id = EXCLUDED.worker_id,
|
||||
worker_session_id = EXCLUDED.worker_session_id,
|
||||
worker_generation = EXCLUDED.worker_generation,
|
||||
lease_token_digest = EXCLUDED.lease_token_digest,
|
||||
offer_id = EXCLUDED.offer_id,
|
||||
acquired_at_ms = EXCLUDED.acquired_at_ms,
|
||||
renewed_at_ms = EXCLUDED.renewed_at_ms,
|
||||
expires_at_ms = EXCLUDED.expires_at_ms,
|
||||
released_at_ms = NULL,
|
||||
release_reason = NULL,
|
||||
completed_at_ms = NULL,
|
||||
updated_at_ms = EXCLUDED.updated_at_ms
|
||||
RETURNING ${LEASE_COLUMNS}
|
||||
`,
|
||||
[
|
||||
command.attemptId,
|
||||
command.runId,
|
||||
version,
|
||||
leaseGeneration,
|
||||
command.workerId,
|
||||
command.workerSessionId,
|
||||
command.workerGeneration,
|
||||
digest,
|
||||
command.offerId,
|
||||
observedAtMs,
|
||||
expiresAtMs,
|
||||
],
|
||||
);
|
||||
const attemptUpdate = await client.query(
|
||||
`
|
||||
UPDATE "ql3"."run_attempts"
|
||||
SET worker_id = $3, worker_session_id = $4, worker_generation = $5,
|
||||
lease_token = NULL, lease_token_digest = $6,
|
||||
lease_generation = $7, lease_version = $8,
|
||||
lease_expires_at_ms = $9, offer_id = $10
|
||||
WHERE id = $1 AND run_id = $2 AND status = 'claimed'
|
||||
`,
|
||||
[
|
||||
command.attemptId,
|
||||
command.runId,
|
||||
command.workerId,
|
||||
command.workerSessionId,
|
||||
command.workerGeneration,
|
||||
digest,
|
||||
leaseGeneration,
|
||||
version,
|
||||
expiresAtMs,
|
||||
command.offerId,
|
||||
],
|
||||
);
|
||||
if (attemptUpdate.rowCount !== 1) {
|
||||
throw new RunDispatchLeaseFenceRejectedError(command.attemptId, 'version_mismatch');
|
||||
}
|
||||
const nextSequence = integer(aggregate, 'eventSequence') + 1;
|
||||
const nextVersion = integer(aggregate, 'runVersion') + 1;
|
||||
const eventType = workflowAttempt
|
||||
? 'workflow.task_dispatch_leased'
|
||||
: aggregate.runStatus === 'queued'
|
||||
? 'run.dispatching'
|
||||
: 'run.dispatch_reclaimed';
|
||||
const runUpdate = await client.query(
|
||||
`
|
||||
UPDATE "ql3"."runs"
|
||||
SET status = $2, version = $3, event_sequence = $4
|
||||
WHERE id = $1 AND version = $5
|
||||
`,
|
||||
[
|
||||
command.runId,
|
||||
workflowAttempt ? 'running' : 'dispatching',
|
||||
nextVersion,
|
||||
nextSequence,
|
||||
integer(aggregate, 'runVersion'),
|
||||
],
|
||||
);
|
||||
if (runUpdate.rowCount !== 1) {
|
||||
throw new RunDispatchLeaseFenceRejectedError(command.attemptId, 'version_mismatch');
|
||||
}
|
||||
await client.query(
|
||||
`
|
||||
INSERT INTO "ql3"."run_events" (
|
||||
id, run_id, sequence, type, dedupe_key, actor_type, actor_id,
|
||||
attempt_id, step_run_id, payload, created_at_ms
|
||||
) VALUES ($1, $2, $3, $4, $5, 'worker', $6, $7, $8, $9::jsonb, $10)
|
||||
`,
|
||||
[
|
||||
command.eventId,
|
||||
command.runId,
|
||||
nextSequence,
|
||||
eventType,
|
||||
`run-dispatch:${command.attemptId}:${leaseGeneration}:claimed`,
|
||||
command.workerId,
|
||||
command.attemptId,
|
||||
workflowAttempt ? aggregate.workflowStepRunId : null,
|
||||
JSON.stringify({
|
||||
attempt_id: command.attemptId,
|
||||
lease_generation: leaseGeneration,
|
||||
execution_scope: workflowAttempt ? 'workflow_task' : 'run',
|
||||
from_status: workflowAttempt ? 'ready' : aggregate.runStatus,
|
||||
to_status: workflowAttempt ? 'ready' : 'dispatching',
|
||||
}),
|
||||
observedAtMs,
|
||||
],
|
||||
);
|
||||
if (persisted.rows.length !== 1) {
|
||||
throw new TypeError('PostgreSQL Run lease claim returned no row');
|
||||
}
|
||||
return Object.freeze({ status: 'claimed' as const, lease: lease(persisted.rows[0]!) });
|
||||
});
|
||||
}
|
||||
|
||||
async renew(
|
||||
command: RenewRunDispatchLeaseCommand,
|
||||
): Promise<RunDispatchLeaseRecord> {
|
||||
assertRenew(command);
|
||||
return this.transaction(async (client) => {
|
||||
await lockAttemptAuthority(client, command.attemptId);
|
||||
const workerResult = await this.lockWorker(client, command.workerId);
|
||||
const current = await this.lockLease(client, command.attemptId);
|
||||
const observedAtMs = await now(client);
|
||||
if (!workerCurrent(workerResult, command, observedAtMs, true)) {
|
||||
throw new RunDispatchLeaseFenceRejectedError(
|
||||
command.attemptId,
|
||||
'worker_unavailable',
|
||||
);
|
||||
}
|
||||
assertLeaseFence(current, command, observedAtMs);
|
||||
const nextVersion = current.version + 1;
|
||||
if (nextVersion > 2_147_483_647) {
|
||||
throw new RangeError('Run dispatch lease version overflowed');
|
||||
}
|
||||
const expiresAtMs = observedAtMs + command.leaseDurationMs;
|
||||
const result = await client.query<Row>(
|
||||
`
|
||||
UPDATE "ql3"."run_dispatch_leases"
|
||||
SET version = $2, renewed_at_ms = $3, expires_at_ms = $4, updated_at_ms = $3
|
||||
WHERE attempt_id = $1 AND version = $5
|
||||
RETURNING ${LEASE_COLUMNS}
|
||||
`,
|
||||
[command.attemptId, nextVersion, observedAtMs, expiresAtMs, current.version],
|
||||
);
|
||||
const attempt = await client.query(
|
||||
`
|
||||
UPDATE "ql3"."run_attempts"
|
||||
SET lease_version = $2, lease_expires_at_ms = $3
|
||||
WHERE id = $1 AND worker_id = $4 AND worker_session_id = $5
|
||||
AND worker_generation = $6 AND lease_generation = $7
|
||||
AND lease_token_digest = $8 AND lease_version = $9
|
||||
AND status IN ('claimed', 'starting', 'running')
|
||||
`,
|
||||
[
|
||||
command.attemptId,
|
||||
nextVersion,
|
||||
expiresAtMs,
|
||||
command.workerId,
|
||||
command.workerSessionId,
|
||||
command.workerGeneration,
|
||||
command.leaseGeneration,
|
||||
digestRunDispatchLeaseToken(command.leaseToken),
|
||||
command.expectedVersion,
|
||||
],
|
||||
);
|
||||
if (result.rows.length !== 1 || attempt.rowCount !== 1) {
|
||||
throw new RunDispatchLeaseFenceRejectedError(command.attemptId, 'version_mismatch');
|
||||
}
|
||||
return lease(result.rows[0]!);
|
||||
});
|
||||
}
|
||||
|
||||
async release(
|
||||
command: ReleaseRunDispatchLeaseCommand,
|
||||
): Promise<RunDispatchLeaseRecord> {
|
||||
assertRelease(command);
|
||||
return this.transaction(async (client) => {
|
||||
await lockAttemptAuthority(client, command.attemptId);
|
||||
const worker = await this.lockWorker(client, command.workerId);
|
||||
const current = await this.lockLease(client, command.attemptId);
|
||||
const aggregate = await client.query<Row>(
|
||||
`
|
||||
SELECT run.version AS "runVersion", run.event_sequence AS "eventSequence",
|
||||
attempt.status AS "attemptStatus",
|
||||
workflow_task.step_run_id AS "workflowStepRunId"
|
||||
FROM "ql3"."runs" AS run
|
||||
INNER JOIN "ql3"."run_attempts" AS attempt ON attempt.id = $2
|
||||
LEFT JOIN
|
||||
"ql3"."plugin_package_workflow_task_attempt_admissions"
|
||||
AS workflow_task
|
||||
ON workflow_task.attempt_id = attempt.id
|
||||
WHERE run.id = $1 AND attempt.run_id = run.id
|
||||
FOR UPDATE OF run, attempt
|
||||
`,
|
||||
[command.runId, command.attemptId],
|
||||
);
|
||||
const observedAtMs = await now(client);
|
||||
if (!workerCurrent(worker, command, observedAtMs, true)) {
|
||||
throw new RunDispatchLeaseFenceRejectedError(
|
||||
command.attemptId,
|
||||
'worker_unavailable',
|
||||
);
|
||||
}
|
||||
assertLeaseFence(current, command, observedAtMs);
|
||||
if (aggregate.rows[0]?.attemptStatus !== 'claimed') {
|
||||
throw new RunDispatchLeaseFenceRejectedError(command.attemptId, 'not_leased');
|
||||
}
|
||||
const nextVersion = current.version + 1;
|
||||
const result = await client.query<Row>(
|
||||
`
|
||||
UPDATE "ql3"."run_dispatch_leases"
|
||||
SET status = 'released', version = $2, released_at_ms = $3,
|
||||
release_reason = $4, completed_at_ms = NULL, updated_at_ms = $3
|
||||
WHERE attempt_id = $1 AND version = $5
|
||||
RETURNING ${LEASE_COLUMNS}
|
||||
`,
|
||||
[command.attemptId, nextVersion, observedAtMs, command.reason, current.version],
|
||||
);
|
||||
const attempt = await client.query(
|
||||
`
|
||||
UPDATE "ql3"."run_attempts"
|
||||
SET worker_id = NULL, worker_session_id = NULL, worker_generation = NULL,
|
||||
lease_token = NULL, lease_token_digest = NULL,
|
||||
lease_generation = NULL, lease_version = NULL,
|
||||
lease_expires_at_ms = NULL, offer_id = NULL
|
||||
WHERE id = $1 AND lease_version = $2
|
||||
`,
|
||||
[command.attemptId, command.expectedVersion],
|
||||
);
|
||||
const aggregateRow = aggregate.rows[0]!;
|
||||
const workflowStepRunId = optionalString(
|
||||
aggregateRow,
|
||||
'workflowStepRunId',
|
||||
);
|
||||
const runVersion = integer(aggregateRow, 'runVersion');
|
||||
const eventSequence = integer(aggregateRow, 'eventSequence') + 1;
|
||||
const run = await client.query(
|
||||
`UPDATE "ql3"."runs" SET version = $2, event_sequence = $3 WHERE id = $1 AND version = $4`,
|
||||
[command.runId, runVersion + 1, eventSequence, runVersion],
|
||||
);
|
||||
if (result.rows.length !== 1 || attempt.rowCount !== 1 || run.rowCount !== 1) {
|
||||
throw new RunDispatchLeaseFenceRejectedError(command.attemptId, 'version_mismatch');
|
||||
}
|
||||
await client.query(
|
||||
`
|
||||
INSERT INTO "ql3"."run_events" (
|
||||
id, run_id, sequence, type, dedupe_key, actor_type, actor_id,
|
||||
attempt_id, step_run_id, payload, created_at_ms
|
||||
) VALUES ($1, $2, $3, $4, $5, 'worker', $6, $7, $8, $9::jsonb, $10)
|
||||
`,
|
||||
[
|
||||
command.eventId,
|
||||
command.runId,
|
||||
eventSequence,
|
||||
workflowStepRunId
|
||||
? 'workflow.task_dispatch_released'
|
||||
: 'run.dispatch_released',
|
||||
`run-dispatch:${command.attemptId}:${command.leaseGeneration}:released`,
|
||||
command.workerId,
|
||||
command.attemptId,
|
||||
workflowStepRunId ?? null,
|
||||
JSON.stringify({
|
||||
attempt_id: command.attemptId,
|
||||
lease_generation: command.leaseGeneration,
|
||||
execution_scope: workflowStepRunId ? 'workflow_task' : 'run',
|
||||
reason: command.reason,
|
||||
}),
|
||||
observedAtMs,
|
||||
],
|
||||
);
|
||||
return lease(result.rows[0]!);
|
||||
});
|
||||
}
|
||||
|
||||
private async lockWorker(
|
||||
client: PostgresClient,
|
||||
workerId: string,
|
||||
): Promise<Row | undefined> {
|
||||
const result = await client.query<Row>(
|
||||
`
|
||||
SELECT worker_id AS "workerId", session_id AS "sessionId",
|
||||
generation, status, lease_expires_at_ms AS "leaseExpiresAtMs"
|
||||
FROM "ql3"."worker_sessions" WHERE worker_id = $1 FOR UPDATE
|
||||
`,
|
||||
[workerId],
|
||||
);
|
||||
return result.rows[0];
|
||||
}
|
||||
|
||||
private async lockLease(
|
||||
client: PostgresClient,
|
||||
attemptId: string,
|
||||
): Promise<RunDispatchLeaseRecord | null> {
|
||||
const result = await client.query<Row>(
|
||||
`SELECT ${LEASE_COLUMNS} FROM "ql3"."run_dispatch_leases" WHERE attempt_id = $1 FOR UPDATE`,
|
||||
[attemptId],
|
||||
);
|
||||
return result.rows[0] ? lease(result.rows[0]) : null;
|
||||
}
|
||||
|
||||
private async transaction<T>(
|
||||
work: (client: PostgresClient) => Promise<T>,
|
||||
): Promise<T> {
|
||||
const client = await this.pool.connect();
|
||||
try {
|
||||
await begin(client);
|
||||
const result = await work(client);
|
||||
await client.query('COMMIT');
|
||||
return result;
|
||||
} catch (error) {
|
||||
try {
|
||||
await client.query('ROLLBACK');
|
||||
} catch {
|
||||
// Preserve the originating failure.
|
||||
}
|
||||
throw error;
|
||||
} finally {
|
||||
client.release();
|
||||
}
|
||||
}
|
||||
}
|
||||
+251
@@ -0,0 +1,251 @@
|
||||
// PostgreSQL Worker execution attestations are owned by this domain.
|
||||
import {
|
||||
WorkerExecutionAttestationFenceRejectedError,
|
||||
WorkerExecutionAttestationUnavailableError,
|
||||
normalizeSubmitWorkerExecutionAttestationCommand,
|
||||
normalizeWorkerExecutionAttestation,
|
||||
type SubmitWorkerExecutionAttestationCommand,
|
||||
type SubmitWorkerExecutionAttestationResult,
|
||||
type WorkerExecutionAttestationRecord,
|
||||
type WorkerExecutionAttestationRepository,
|
||||
} from '@qinglong/runtime-core/worker-attestation';
|
||||
import type { PostgresClient, PostgresPool } from '@qinglong/runtime-core';
|
||||
import { lockAttemptAuthority } from '../run/attemptAuthorityLock';
|
||||
|
||||
type Row = Record<string, unknown>;
|
||||
|
||||
const COLUMNS = `
|
||||
attestation_id AS "attestationId", run_id AS "runId",
|
||||
attempt_id AS "attemptId", sequence, state, worker_id AS "workerId",
|
||||
worker_session_id AS "workerSessionId",
|
||||
worker_generation AS "workerGeneration",
|
||||
lease_token_digest AS "leaseTokenDigest",
|
||||
lease_generation AS "leaseGeneration", lease_version AS "leaseVersion",
|
||||
offer_id AS "offerId", callback_sequence AS "callbackSequence",
|
||||
executor_handle AS "executorHandle", journal_revision AS "journalRevision",
|
||||
received_at_ms AS "receivedAtMs"
|
||||
`.trim();
|
||||
|
||||
function text(row: Row, key: string): string {
|
||||
const value = row[key];
|
||||
if (typeof value !== 'string' || value.length < 1) {
|
||||
throw new TypeError(`PostgreSQL Worker attestation ${key} is invalid`);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
function integer(row: Row, key: string): number {
|
||||
const raw = row[key];
|
||||
const value = typeof raw === 'string' && /^\d+$/.test(raw) ? Number(raw) : raw;
|
||||
if (typeof value !== 'number' || !Number.isSafeInteger(value)) {
|
||||
throw new TypeError(`PostgreSQL Worker attestation ${key} is invalid`);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
function record(row: Row): Readonly<WorkerExecutionAttestationRecord> {
|
||||
return normalizeWorkerExecutionAttestation({
|
||||
attestationId: text(row, 'attestationId'),
|
||||
runId: text(row, 'runId'),
|
||||
attemptId: text(row, 'attemptId'),
|
||||
sequence: integer(row, 'sequence'),
|
||||
state: text(row, 'state') as WorkerExecutionAttestationRecord['state'],
|
||||
workerId: text(row, 'workerId'),
|
||||
workerSessionId: text(row, 'workerSessionId'),
|
||||
workerGeneration: integer(row, 'workerGeneration'),
|
||||
leaseTokenDigest: text(row, 'leaseTokenDigest'),
|
||||
leaseGeneration: integer(row, 'leaseGeneration'),
|
||||
leaseVersion: integer(row, 'leaseVersion'),
|
||||
offerId: text(row, 'offerId'),
|
||||
callbackSequence: integer(row, 'callbackSequence'),
|
||||
executorHandle: text(row, 'executorHandle'),
|
||||
journalRevision: integer(row, 'journalRevision'),
|
||||
receivedAtMs: integer(row, 'receivedAtMs'),
|
||||
});
|
||||
}
|
||||
|
||||
function sameCommand(
|
||||
current: Readonly<WorkerExecutionAttestationRecord>,
|
||||
command: Readonly<SubmitWorkerExecutionAttestationCommand>,
|
||||
): boolean {
|
||||
const { receivedAtMs: _receivedAtMs, ...semantic } = current;
|
||||
return JSON.stringify(semantic) === JSON.stringify(command);
|
||||
}
|
||||
|
||||
async function begin(client: PostgresClient): Promise<void> {
|
||||
await client.query('BEGIN');
|
||||
await client.query("SET LOCAL statement_timeout = '5s'");
|
||||
await client.query("SET LOCAL lock_timeout = '1s'");
|
||||
await client.query("SET LOCAL idle_in_transaction_session_timeout = '10s'");
|
||||
}
|
||||
|
||||
export class PostgresWorkerExecutionAttestationRepository
|
||||
implements WorkerExecutionAttestationRepository
|
||||
{
|
||||
constructor(private readonly pool: PostgresPool) {}
|
||||
|
||||
async submit(
|
||||
input: SubmitWorkerExecutionAttestationCommand,
|
||||
): Promise<SubmitWorkerExecutionAttestationResult> {
|
||||
const command = normalizeSubmitWorkerExecutionAttestationCommand(input);
|
||||
const client = await this.pool.connect();
|
||||
try {
|
||||
await begin(client);
|
||||
await lockAttemptAuthority(client, command.attemptId);
|
||||
const replayResult = await client.query<Row>(
|
||||
`SELECT ${COLUMNS} FROM "ql3"."worker_execution_attestations" WHERE attestation_id = $1`,
|
||||
[command.attestationId],
|
||||
);
|
||||
if (replayResult.rows.length > 1) {
|
||||
throw new WorkerExecutionAttestationUnavailableError();
|
||||
}
|
||||
if (replayResult.rows[0]) {
|
||||
const current = record(replayResult.rows[0]);
|
||||
if (!sameCommand(current, command)) {
|
||||
throw new WorkerExecutionAttestationFenceRejectedError('attestation_id_conflict');
|
||||
}
|
||||
await client.query('COMMIT');
|
||||
return Object.freeze({ status: 'existing', attestation: current });
|
||||
}
|
||||
|
||||
const authority = await client.query<Row>(
|
||||
`
|
||||
SELECT attempt.run_id AS "attemptRunId", attempt.status AS "attemptStatus",
|
||||
attempt.worker_id AS "attemptWorkerId",
|
||||
attempt.worker_session_id AS "attemptWorkerSessionId",
|
||||
attempt.worker_generation AS "attemptWorkerGeneration",
|
||||
attempt.lease_token_digest AS "attemptLeaseTokenDigest",
|
||||
attempt.lease_generation AS "attemptLeaseGeneration",
|
||||
attempt.lease_version AS "attemptLeaseVersion",
|
||||
attempt.offer_id AS "attemptOfferId",
|
||||
attempt.callback_sequence AS "attemptCallbackSequence",
|
||||
attempt.executor_handle AS "attemptExecutorHandle",
|
||||
session.session_id AS "sessionId",
|
||||
session.generation AS "sessionGeneration",
|
||||
lease.run_id AS "leaseRunId", lease.status AS "leaseStatus",
|
||||
lease.worker_id AS "leaseWorkerId",
|
||||
lease.worker_session_id AS "leaseWorkerSessionId",
|
||||
lease.worker_generation AS "leaseWorkerGeneration",
|
||||
lease.lease_token_digest AS "leaseTokenDigest",
|
||||
lease.lease_generation AS "leaseGeneration",
|
||||
lease.version AS "leaseVersion", lease.offer_id AS "leaseOfferId"
|
||||
FROM "ql3"."run_attempts" AS attempt
|
||||
JOIN "ql3"."run_dispatch_leases" AS lease ON lease.attempt_id = attempt.id
|
||||
JOIN "ql3"."worker_sessions" AS session ON session.worker_id = lease.worker_id
|
||||
WHERE attempt.id = $1
|
||||
`,
|
||||
[command.attemptId],
|
||||
);
|
||||
const row = authority.rows[0];
|
||||
const matches =
|
||||
row &&
|
||||
row.attemptRunId === command.runId &&
|
||||
(row.attemptStatus === 'starting' || row.attemptStatus === 'running') &&
|
||||
row.attemptWorkerId === command.workerId &&
|
||||
row.attemptWorkerSessionId === command.workerSessionId &&
|
||||
integer(row, 'attemptWorkerGeneration') === command.workerGeneration &&
|
||||
row.attemptLeaseTokenDigest === command.leaseTokenDigest &&
|
||||
integer(row, 'attemptLeaseGeneration') === command.leaseGeneration &&
|
||||
integer(row, 'attemptLeaseVersion') === command.leaseVersion &&
|
||||
row.attemptOfferId === command.offerId &&
|
||||
integer(row, 'attemptCallbackSequence') === command.callbackSequence &&
|
||||
row.attemptExecutorHandle === command.executorHandle &&
|
||||
row.sessionId === command.workerSessionId &&
|
||||
integer(row, 'sessionGeneration') === command.workerGeneration &&
|
||||
row.leaseRunId === command.runId &&
|
||||
row.leaseStatus === 'leased' &&
|
||||
row.leaseWorkerId === command.workerId &&
|
||||
row.leaseWorkerSessionId === command.workerSessionId &&
|
||||
integer(row, 'leaseWorkerGeneration') === command.workerGeneration &&
|
||||
row.leaseTokenDigest === command.leaseTokenDigest &&
|
||||
integer(row, 'leaseGeneration') === command.leaseGeneration &&
|
||||
integer(row, 'leaseVersion') === command.leaseVersion &&
|
||||
row.leaseOfferId === command.offerId;
|
||||
if (!matches) {
|
||||
throw new WorkerExecutionAttestationFenceRejectedError('authority_mismatch');
|
||||
}
|
||||
|
||||
const previous = await client.query<Row>(
|
||||
`
|
||||
SELECT sequence, state, journal_revision AS "journalRevision"
|
||||
FROM "ql3"."worker_execution_attestations"
|
||||
WHERE attempt_id = $1 AND lease_generation = $2
|
||||
ORDER BY sequence DESC LIMIT 1
|
||||
`,
|
||||
[command.attemptId, command.leaseGeneration],
|
||||
);
|
||||
const last = previous.rows[0];
|
||||
if (
|
||||
command.sequence !== (last ? integer(last, 'sequence') + 1 : 1) ||
|
||||
(last && integer(last, 'journalRevision') >= command.journalRevision) ||
|
||||
(last?.state === 'stopped' && command.state !== 'stopped')
|
||||
) {
|
||||
throw new WorkerExecutionAttestationFenceRejectedError('sequence_mismatch');
|
||||
}
|
||||
const nowResult = await client.query<Row>(
|
||||
`SELECT floor(extract(epoch FROM statement_timestamp()) * 1000)::bigint AS "receivedAtMs"`,
|
||||
);
|
||||
const receivedAtMs = integer(nowResult.rows[0]!, 'receivedAtMs');
|
||||
const inserted = await client.query<Row>(
|
||||
`
|
||||
INSERT INTO "ql3"."worker_execution_attestations" (
|
||||
attestation_id, run_id, attempt_id, sequence, state, worker_id,
|
||||
worker_session_id, worker_generation, lease_token_digest,
|
||||
lease_generation, lease_version, offer_id, callback_sequence,
|
||||
executor_handle, journal_revision, received_at_ms
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16
|
||||
) RETURNING ${COLUMNS}
|
||||
`,
|
||||
[
|
||||
command.attestationId, command.runId, command.attemptId,
|
||||
command.sequence, command.state, command.workerId,
|
||||
command.workerSessionId, command.workerGeneration,
|
||||
command.leaseTokenDigest, command.leaseGeneration,
|
||||
command.leaseVersion, command.offerId, command.callbackSequence,
|
||||
command.executorHandle, command.journalRevision, receivedAtMs,
|
||||
],
|
||||
);
|
||||
const attestation = record(inserted.rows[0]!);
|
||||
await client.query('COMMIT');
|
||||
return Object.freeze({ status: 'created', attestation });
|
||||
} catch (error) {
|
||||
try { await client.query('ROLLBACK'); } catch { /* preserve root */ }
|
||||
if (
|
||||
error instanceof WorkerExecutionAttestationFenceRejectedError ||
|
||||
error instanceof WorkerExecutionAttestationUnavailableError
|
||||
) throw error;
|
||||
throw new WorkerExecutionAttestationUnavailableError();
|
||||
} finally {
|
||||
client.release();
|
||||
}
|
||||
}
|
||||
|
||||
async findLatestExact(
|
||||
target: Parameters<WorkerExecutionAttestationRepository['findLatestExact']>[0],
|
||||
): Promise<Readonly<WorkerExecutionAttestationRecord> | null> {
|
||||
const result = await this.pool.query<Row>(
|
||||
`
|
||||
SELECT ${COLUMNS}
|
||||
FROM "ql3"."worker_execution_attestations"
|
||||
WHERE run_id = $1 AND attempt_id = $2 AND worker_id = $3
|
||||
AND worker_session_id = $4 AND worker_generation = $5
|
||||
AND lease_token_digest = $6 AND lease_generation = $7
|
||||
AND lease_version = $8 AND offer_id = $9
|
||||
AND callback_sequence = $10 AND executor_handle = $11
|
||||
ORDER BY sequence DESC LIMIT 1
|
||||
`,
|
||||
[
|
||||
target.runId, target.attemptId, target.workerId,
|
||||
target.workerSessionId, target.workerGeneration,
|
||||
target.leaseTokenDigest, target.leaseGeneration, target.leaseVersion,
|
||||
target.offerId, target.callbackSequence, target.executorHandle,
|
||||
],
|
||||
);
|
||||
try {
|
||||
return result.rows[0] ? record(result.rows[0]) : null;
|
||||
} catch {
|
||||
throw new WorkerExecutionAttestationUnavailableError();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,675 @@
|
||||
// PostgreSQL Remote Worker sessions are owned by this domain.
|
||||
import type {
|
||||
AvailableWorkerSessionPage,
|
||||
HeartbeatWorkerSessionCommand,
|
||||
PostgresClient,
|
||||
PostgresPool,
|
||||
RegisterWorkerSessionCommand,
|
||||
RegisterWorkerSessionResult,
|
||||
TransitionWorkerSessionCommand,
|
||||
WorkerSessionRecord,
|
||||
WorkerSessionRepository,
|
||||
WorkerSessionStatus,
|
||||
} from '@qinglong/runtime-core';
|
||||
import {
|
||||
MAX_AVAILABLE_WORKER_PAGE_SIZE,
|
||||
MAX_WORKER_CONCURRENT_RUNS,
|
||||
WORKER_SESSION_STATUSES,
|
||||
WorkerSessionConflictError,
|
||||
WorkerSessionFenceRejectedError,
|
||||
assertWorkerCapabilitiesSnapshot,
|
||||
assertWorkerConcurrency,
|
||||
assertWorkerId,
|
||||
assertWorkerSessionId,
|
||||
assertWorkerSessionLeaseDuration,
|
||||
assertWorkerSessionRecord,
|
||||
} from '@qinglong/runtime-core';
|
||||
import {
|
||||
WorkerCredentialDeliveryConflictError,
|
||||
WorkerCredentialDeliveryUnavailableError,
|
||||
normalizeAuthenticatedWorkerCredentialIdentity,
|
||||
normalizeWorkerCredentialDeliveryRecord,
|
||||
type AuthenticatedWorkerCredentialIdentity,
|
||||
type AuthenticatedWorkerSessionRepository,
|
||||
type WorkerCredentialDeliveryRecord,
|
||||
} from '@qinglong/runtime-core/worker-credential-delivery';
|
||||
|
||||
type Row = Record<string, unknown>;
|
||||
|
||||
interface DeliveryRow extends Row {
|
||||
deliveryId: unknown;
|
||||
version: unknown;
|
||||
state: unknown;
|
||||
workerId: unknown;
|
||||
credentialId: unknown;
|
||||
credentialVersion: unknown;
|
||||
previousCredentialId: unknown;
|
||||
secretDigest: unknown;
|
||||
tokenDigest: unknown;
|
||||
deploymentTargetDigest: unknown;
|
||||
deploymentGeneration: unknown;
|
||||
stagedAtMs: unknown;
|
||||
credentialCommittedAtMs: unknown;
|
||||
publishedAtMs: unknown;
|
||||
publicationDigest: unknown;
|
||||
observedAtMs: unknown;
|
||||
observedSessionId: unknown;
|
||||
observedSessionVersion: unknown;
|
||||
previousRevokedAtMs: unknown;
|
||||
}
|
||||
|
||||
const SELECT_COLUMNS = `
|
||||
worker_id AS "workerId",
|
||||
session_id AS "sessionId",
|
||||
generation AS "generation",
|
||||
status AS "status",
|
||||
version AS "version",
|
||||
capabilities_json AS "capabilitiesJson",
|
||||
capabilities_hash AS "capabilitiesHash",
|
||||
max_concurrent_runs AS "maxConcurrentRuns",
|
||||
available_slots AS "availableSlots",
|
||||
registered_at_ms AS "registeredAtMs",
|
||||
last_heartbeat_at_ms AS "lastHeartbeatAtMs",
|
||||
lease_expires_at_ms AS "leaseExpiresAtMs",
|
||||
updated_at_ms AS "updatedAtMs"
|
||||
`.trim();
|
||||
|
||||
const DELIVERY_COLUMNS = `
|
||||
delivery_id AS "deliveryId", version, state,
|
||||
worker_id AS "workerId", credential_id AS "credentialId",
|
||||
credential_version AS "credentialVersion",
|
||||
previous_credential_id AS "previousCredentialId",
|
||||
secret_digest AS "secretDigest", token_digest AS "tokenDigest",
|
||||
deployment_target_digest AS "deploymentTargetDigest",
|
||||
deployment_generation AS "deploymentGeneration",
|
||||
staged_at_ms AS "stagedAtMs",
|
||||
credential_committed_at_ms AS "credentialCommittedAtMs",
|
||||
published_at_ms AS "publishedAtMs",
|
||||
publication_digest AS "publicationDigest",
|
||||
observed_at_ms AS "observedAtMs",
|
||||
observed_session_id AS "observedSessionId",
|
||||
observed_session_version AS "observedSessionVersion",
|
||||
previous_revoked_at_ms AS "previousRevokedAtMs"
|
||||
`.trim();
|
||||
|
||||
function string(row: Row, key: string): string {
|
||||
const value = row[key];
|
||||
if (typeof value !== 'string' || value.length === 0) {
|
||||
throw new TypeError(`PostgreSQL Worker session ${key} is invalid`);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
function integer(row: Row, key: string): number {
|
||||
const value = row[key];
|
||||
const normalized =
|
||||
typeof value === 'string' && /^(0|[1-9]\d*)$/.test(value)
|
||||
? Number(value)
|
||||
: value;
|
||||
if (typeof normalized !== 'number' || !Number.isSafeInteger(normalized)) {
|
||||
throw new TypeError(`PostgreSQL Worker session ${key} is invalid`);
|
||||
}
|
||||
return normalized;
|
||||
}
|
||||
|
||||
function nullableString(row: Row, key: string): string | null {
|
||||
return row[key] === null ? null : string(row, key);
|
||||
}
|
||||
|
||||
function nullableInteger(row: Row, key: string): number | null {
|
||||
return row[key] === null ? null : integer(row, key);
|
||||
}
|
||||
|
||||
function deliveryRecord(
|
||||
row: DeliveryRow,
|
||||
): Readonly<WorkerCredentialDeliveryRecord> {
|
||||
return normalizeWorkerCredentialDeliveryRecord({
|
||||
deliveryId: string(row, 'deliveryId'),
|
||||
version: integer(row, 'version'),
|
||||
state: string(row, 'state') as WorkerCredentialDeliveryRecord['state'],
|
||||
workerId: string(row, 'workerId'),
|
||||
credentialId: string(row, 'credentialId'),
|
||||
credentialVersion: integer(row, 'credentialVersion'),
|
||||
previousCredentialId: nullableString(row, 'previousCredentialId'),
|
||||
secretDigest: string(row, 'secretDigest'),
|
||||
tokenDigest: string(row, 'tokenDigest'),
|
||||
deploymentTargetDigest: string(row, 'deploymentTargetDigest'),
|
||||
deploymentGeneration: string(row, 'deploymentGeneration'),
|
||||
stagedAtMs: integer(row, 'stagedAtMs'),
|
||||
credentialCommittedAtMs: integer(row, 'credentialCommittedAtMs'),
|
||||
publishedAtMs: nullableInteger(row, 'publishedAtMs'),
|
||||
publicationDigest: nullableString(row, 'publicationDigest'),
|
||||
observedAtMs: nullableInteger(row, 'observedAtMs'),
|
||||
observedSessionId: nullableString(row, 'observedSessionId'),
|
||||
observedSessionVersion: nullableInteger(row, 'observedSessionVersion'),
|
||||
previousRevokedAtMs: nullableInteger(row, 'previousRevokedAtMs'),
|
||||
});
|
||||
}
|
||||
|
||||
function sameDeliveryIdentity(
|
||||
current: Readonly<WorkerCredentialDeliveryRecord>,
|
||||
previous: Readonly<WorkerCredentialDeliveryRecord>,
|
||||
): boolean {
|
||||
return (
|
||||
current.deliveryId === previous.deliveryId &&
|
||||
current.workerId === previous.workerId &&
|
||||
current.credentialId === previous.credentialId &&
|
||||
current.credentialVersion === previous.credentialVersion &&
|
||||
current.previousCredentialId === previous.previousCredentialId &&
|
||||
current.secretDigest === previous.secretDigest &&
|
||||
current.tokenDigest === previous.tokenDigest &&
|
||||
current.deploymentTargetDigest === previous.deploymentTargetDigest &&
|
||||
current.deploymentGeneration === previous.deploymentGeneration &&
|
||||
current.stagedAtMs === previous.stagedAtMs &&
|
||||
current.credentialCommittedAtMs === previous.credentialCommittedAtMs
|
||||
);
|
||||
}
|
||||
|
||||
async function observeCredentialDelivery(
|
||||
client: PostgresClient,
|
||||
requestedCredential: AuthenticatedWorkerCredentialIdentity,
|
||||
worker: Readonly<WorkerSessionRecord>,
|
||||
nowMs: number,
|
||||
): Promise<void> {
|
||||
const credential = normalizeAuthenticatedWorkerCredentialIdentity(
|
||||
requestedCredential,
|
||||
);
|
||||
if (credential.workerId !== worker.workerId || worker.version < 1) {
|
||||
throw new WorkerCredentialDeliveryConflictError();
|
||||
}
|
||||
const result = await client.query<DeliveryRow>(
|
||||
`SELECT ${DELIVERY_COLUMNS}
|
||||
FROM "ql3"."worker_credential_deliveries"
|
||||
WHERE worker_id = $1
|
||||
AND credential_id = $2
|
||||
AND credential_version = $3
|
||||
ORDER BY delivery_id ASC, version ASC
|
||||
LIMIT 5`,
|
||||
[credential.workerId, credential.credentialId, credential.credentialVersion],
|
||||
);
|
||||
if (result.rows.length === 0) return;
|
||||
if (result.rows.length > 4) {
|
||||
throw new WorkerCredentialDeliveryConflictError();
|
||||
}
|
||||
const records = result.rows.map(deliveryRecord);
|
||||
for (let index = 0; index < records.length; index += 1) {
|
||||
const current = records[index]!;
|
||||
const previous = records[index - 1];
|
||||
if (current.version !== index + 1) {
|
||||
throw new WorkerCredentialDeliveryConflictError();
|
||||
}
|
||||
if (!previous) continue;
|
||||
if (
|
||||
!sameDeliveryIdentity(current, previous) ||
|
||||
(current.version >= 3 &&
|
||||
(current.publishedAtMs !== previous.publishedAtMs ||
|
||||
current.publicationDigest !== previous.publicationDigest)) ||
|
||||
(current.version >= 4 &&
|
||||
(current.observedAtMs !== previous.observedAtMs ||
|
||||
current.observedSessionId !== previous.observedSessionId ||
|
||||
current.observedSessionVersion !== previous.observedSessionVersion))
|
||||
) {
|
||||
throw new WorkerCredentialDeliveryConflictError();
|
||||
}
|
||||
}
|
||||
const latest = records.at(-1)!;
|
||||
if (latest.state === 'credential_committed') {
|
||||
throw new WorkerCredentialDeliveryUnavailableError();
|
||||
}
|
||||
if (latest.state !== 'published') return;
|
||||
const observed = normalizeWorkerCredentialDeliveryRecord({
|
||||
...latest,
|
||||
version: 3,
|
||||
state: 'observed',
|
||||
observedAtMs: nowMs,
|
||||
observedSessionId: worker.sessionId,
|
||||
observedSessionVersion: worker.version,
|
||||
});
|
||||
await client.query(
|
||||
`INSERT INTO "ql3"."worker_credential_deliveries" (
|
||||
delivery_id, version, state, worker_id, credential_id,
|
||||
credential_version, previous_credential_id, secret_digest,
|
||||
token_digest, deployment_target_digest, deployment_generation,
|
||||
staged_at_ms, credential_committed_at_ms, published_at_ms,
|
||||
publication_digest, observed_at_ms, observed_session_id,
|
||||
observed_session_version, previous_revoked_at_ms
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10,
|
||||
$11, $12, $13, $14, $15, $16, $17, $18, $19
|
||||
)`,
|
||||
[
|
||||
observed.deliveryId, observed.version, observed.state,
|
||||
observed.workerId, observed.credentialId, observed.credentialVersion,
|
||||
observed.previousCredentialId, observed.secretDigest,
|
||||
observed.tokenDigest, observed.deploymentTargetDigest,
|
||||
observed.deploymentGeneration, observed.stagedAtMs,
|
||||
observed.credentialCommittedAtMs, observed.publishedAtMs,
|
||||
observed.publicationDigest, observed.observedAtMs,
|
||||
observed.observedSessionId, observed.observedSessionVersion,
|
||||
observed.previousRevokedAtMs,
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
function record(row: Row): WorkerSessionRecord {
|
||||
const status = string(row, 'status');
|
||||
if (!WORKER_SESSION_STATUSES.includes(status as WorkerSessionStatus)) {
|
||||
throw new TypeError('PostgreSQL Worker session status is invalid');
|
||||
}
|
||||
const result: WorkerSessionRecord = Object.freeze({
|
||||
workerId: string(row, 'workerId'),
|
||||
sessionId: string(row, 'sessionId'),
|
||||
generation: integer(row, 'generation'),
|
||||
status: status as WorkerSessionStatus,
|
||||
version: integer(row, 'version'),
|
||||
capabilitiesJson: string(row, 'capabilitiesJson'),
|
||||
capabilitiesHash: string(row, 'capabilitiesHash'),
|
||||
maxConcurrentRuns: integer(row, 'maxConcurrentRuns'),
|
||||
availableSlots: integer(row, 'availableSlots'),
|
||||
registeredAtMs: integer(row, 'registeredAtMs'),
|
||||
lastHeartbeatAtMs: integer(row, 'lastHeartbeatAtMs'),
|
||||
leaseExpiresAtMs: integer(row, 'leaseExpiresAtMs'),
|
||||
updatedAtMs: integer(row, 'updatedAtMs'),
|
||||
});
|
||||
assertWorkerSessionRecord(result);
|
||||
return result;
|
||||
}
|
||||
|
||||
function assertRegister(command: RegisterWorkerSessionCommand): void {
|
||||
assertWorkerId(command.workerId);
|
||||
assertWorkerSessionId(command.sessionId);
|
||||
assertWorkerCapabilitiesSnapshot(
|
||||
command.capabilitiesJson,
|
||||
command.capabilitiesHash,
|
||||
);
|
||||
assertWorkerConcurrency(command.maxConcurrentRuns, command.availableSlots);
|
||||
assertWorkerSessionLeaseDuration(command.leaseDurationMs);
|
||||
}
|
||||
|
||||
function assertHeartbeat(command: HeartbeatWorkerSessionCommand): void {
|
||||
assertWorkerId(command.workerId);
|
||||
assertWorkerSessionId(command.sessionId);
|
||||
if (
|
||||
!Number.isSafeInteger(command.availableSlots) ||
|
||||
command.availableSlots < 0 ||
|
||||
command.availableSlots > MAX_WORKER_CONCURRENT_RUNS
|
||||
) {
|
||||
throw new RangeError('Worker heartbeat availableSlots is invalid');
|
||||
}
|
||||
for (const [name, value, minimum] of [
|
||||
['generation', command.generation, 1],
|
||||
['expectedVersion', command.expectedVersion, 0],
|
||||
] as const) {
|
||||
if (!Number.isSafeInteger(value) || value < minimum) {
|
||||
throw new RangeError(`Worker heartbeat ${name} is invalid`);
|
||||
}
|
||||
}
|
||||
assertWorkerSessionLeaseDuration(command.leaseDurationMs);
|
||||
}
|
||||
|
||||
function assertTransition(command: TransitionWorkerSessionCommand): void {
|
||||
assertWorkerId(command.workerId);
|
||||
assertWorkerSessionId(command.sessionId);
|
||||
if (command.status !== 'draining' && command.status !== 'offline') {
|
||||
throw new TypeError('Worker transition status is invalid');
|
||||
}
|
||||
if (
|
||||
!Number.isSafeInteger(command.generation) ||
|
||||
command.generation < 1 ||
|
||||
!Number.isSafeInteger(command.expectedVersion) ||
|
||||
command.expectedVersion < 0
|
||||
) {
|
||||
throw new RangeError('Worker transition fence is invalid');
|
||||
}
|
||||
}
|
||||
|
||||
async function begin(client: PostgresClient): Promise<void> {
|
||||
await client.query('BEGIN');
|
||||
await client.query("SET LOCAL statement_timeout = '5s'");
|
||||
await client.query("SET LOCAL lock_timeout = '1s'");
|
||||
await client.query("SET LOCAL idle_in_transaction_session_timeout = '10s'");
|
||||
}
|
||||
|
||||
async function observedAtMs(client: PostgresClient): Promise<number> {
|
||||
const result = await client.query<Row>(`
|
||||
SELECT floor(extract(epoch FROM statement_timestamp()) * 1000)::bigint
|
||||
AS "observedAtMs"
|
||||
`);
|
||||
if (result.rows.length !== 1) {
|
||||
throw new TypeError('PostgreSQL Worker observation is invalid');
|
||||
}
|
||||
return integer(result.rows[0]!, 'observedAtMs');
|
||||
}
|
||||
|
||||
function fence(
|
||||
current: WorkerSessionRecord | null,
|
||||
command: {
|
||||
workerId: string;
|
||||
sessionId: string;
|
||||
generation: number;
|
||||
expectedVersion: number;
|
||||
},
|
||||
nowMs: number,
|
||||
): void {
|
||||
if (!current) throw new WorkerSessionFenceRejectedError(command.workerId, 'missing');
|
||||
if (current.sessionId !== command.sessionId) {
|
||||
throw new WorkerSessionFenceRejectedError(command.workerId, 'session_mismatch');
|
||||
}
|
||||
if (current.generation !== command.generation) {
|
||||
throw new WorkerSessionFenceRejectedError(command.workerId, 'generation_mismatch');
|
||||
}
|
||||
if (current.version !== command.expectedVersion) {
|
||||
throw new WorkerSessionFenceRejectedError(command.workerId, 'version_mismatch');
|
||||
}
|
||||
if (current.status === 'offline') {
|
||||
throw new WorkerSessionFenceRejectedError(command.workerId, 'offline');
|
||||
}
|
||||
if (current.leaseExpiresAtMs <= nowMs) {
|
||||
throw new WorkerSessionFenceRejectedError(command.workerId, 'lease_expired');
|
||||
}
|
||||
}
|
||||
|
||||
export class PostgresWorkerSessionRepository
|
||||
implements WorkerSessionRepository, AuthenticatedWorkerSessionRepository
|
||||
{
|
||||
constructor(private readonly pool: PostgresPool) {}
|
||||
|
||||
async findById(workerId: string): Promise<WorkerSessionRecord | null> {
|
||||
assertWorkerId(workerId);
|
||||
const result = await this.pool.query<Row>(
|
||||
`SELECT ${SELECT_COLUMNS} FROM "ql3"."worker_sessions" WHERE worker_id = $1`,
|
||||
[workerId],
|
||||
);
|
||||
if (result.rows.length > 1) {
|
||||
throw new TypeError('PostgreSQL Worker lookup returned multiple rows');
|
||||
}
|
||||
return result.rows[0] ? record(result.rows[0]) : null;
|
||||
}
|
||||
|
||||
async register(
|
||||
command: RegisterWorkerSessionCommand,
|
||||
): Promise<RegisterWorkerSessionResult> {
|
||||
assertRegister(command);
|
||||
const client = await this.pool.connect();
|
||||
try {
|
||||
await begin(client);
|
||||
await client.query(
|
||||
`SELECT pg_advisory_xact_lock(hashtextextended($1, 684022785147727641::bigint))`,
|
||||
[command.workerId],
|
||||
);
|
||||
const currentResult = await client.query<Row>(
|
||||
`SELECT ${SELECT_COLUMNS} FROM "ql3"."worker_sessions" WHERE worker_id = $1 FOR UPDATE`,
|
||||
[command.workerId],
|
||||
);
|
||||
const current = currentResult.rows[0] ? record(currentResult.rows[0]) : null;
|
||||
const nowMs = await observedAtMs(client);
|
||||
if (current?.sessionId === command.sessionId) {
|
||||
if (
|
||||
current.capabilitiesJson !== command.capabilitiesJson ||
|
||||
current.capabilitiesHash !== command.capabilitiesHash ||
|
||||
current.maxConcurrentRuns !== command.maxConcurrentRuns ||
|
||||
current.availableSlots !== command.availableSlots
|
||||
) {
|
||||
throw new WorkerSessionConflictError(command.workerId);
|
||||
}
|
||||
if (current.status === 'offline') {
|
||||
throw new WorkerSessionFenceRejectedError(command.workerId, 'offline');
|
||||
}
|
||||
if (current.leaseExpiresAtMs <= nowMs) {
|
||||
throw new WorkerSessionFenceRejectedError(command.workerId, 'lease_expired');
|
||||
}
|
||||
await client.query('COMMIT');
|
||||
return Object.freeze({ worker: current, replacedSession: false });
|
||||
}
|
||||
const generation = current ? current.generation + 1 : 1;
|
||||
const version = current ? current.version + 1 : 0;
|
||||
if (generation > 2_147_483_647 || version > 2_147_483_647) {
|
||||
throw new RangeError('Worker session generation or version overflowed');
|
||||
}
|
||||
const expiresAtMs = nowMs + command.leaseDurationMs;
|
||||
const result = await client.query<Row>(
|
||||
`
|
||||
INSERT INTO "ql3"."worker_sessions" (
|
||||
worker_id, session_id, generation, status, version,
|
||||
capabilities_json, capabilities_hash, max_concurrent_runs,
|
||||
available_slots, registered_at_ms, last_heartbeat_at_ms,
|
||||
lease_expires_at_ms, updated_at_ms
|
||||
) VALUES ($1, $2, $3, 'online', $4, $5, $6, $7, $8, $9, $9, $10, $9)
|
||||
ON CONFLICT (worker_id) DO UPDATE SET
|
||||
session_id = EXCLUDED.session_id,
|
||||
generation = EXCLUDED.generation,
|
||||
status = EXCLUDED.status,
|
||||
version = EXCLUDED.version,
|
||||
capabilities_json = EXCLUDED.capabilities_json,
|
||||
capabilities_hash = EXCLUDED.capabilities_hash,
|
||||
max_concurrent_runs = EXCLUDED.max_concurrent_runs,
|
||||
available_slots = EXCLUDED.available_slots,
|
||||
registered_at_ms = EXCLUDED.registered_at_ms,
|
||||
last_heartbeat_at_ms = EXCLUDED.last_heartbeat_at_ms,
|
||||
lease_expires_at_ms = EXCLUDED.lease_expires_at_ms,
|
||||
updated_at_ms = EXCLUDED.updated_at_ms
|
||||
RETURNING ${SELECT_COLUMNS}
|
||||
`,
|
||||
[
|
||||
command.workerId,
|
||||
command.sessionId,
|
||||
generation,
|
||||
version,
|
||||
command.capabilitiesJson,
|
||||
command.capabilitiesHash,
|
||||
command.maxConcurrentRuns,
|
||||
command.availableSlots,
|
||||
nowMs,
|
||||
expiresAtMs,
|
||||
],
|
||||
);
|
||||
if (result.rows.length !== 1) {
|
||||
throw new TypeError('PostgreSQL Worker registration returned no row');
|
||||
}
|
||||
await client.query('COMMIT');
|
||||
return Object.freeze({
|
||||
worker: record(result.rows[0]!),
|
||||
replacedSession: current !== null,
|
||||
});
|
||||
} catch (error) {
|
||||
try {
|
||||
await client.query('ROLLBACK');
|
||||
} catch {
|
||||
// Preserve the originating failure.
|
||||
}
|
||||
throw error;
|
||||
} finally {
|
||||
client.release();
|
||||
}
|
||||
}
|
||||
|
||||
async heartbeat(
|
||||
command: HeartbeatWorkerSessionCommand,
|
||||
): Promise<WorkerSessionRecord> {
|
||||
return this.heartbeatInternal(command);
|
||||
}
|
||||
|
||||
async heartbeatAuthenticated(
|
||||
command: HeartbeatWorkerSessionCommand,
|
||||
credential: AuthenticatedWorkerCredentialIdentity,
|
||||
): Promise<WorkerSessionRecord> {
|
||||
return this.heartbeatInternal(command, credential);
|
||||
}
|
||||
|
||||
private async heartbeatInternal(
|
||||
command: HeartbeatWorkerSessionCommand,
|
||||
credential?: AuthenticatedWorkerCredentialIdentity,
|
||||
): Promise<WorkerSessionRecord> {
|
||||
assertHeartbeat(command);
|
||||
return this.mutate(command.workerId, async (client, nowMs, current) => {
|
||||
fence(current, command, nowMs);
|
||||
if (!current) throw new WorkerSessionFenceRejectedError(command.workerId, 'missing');
|
||||
assertWorkerConcurrency(current.maxConcurrentRuns, command.availableSlots);
|
||||
const result = await client.query<Row>(
|
||||
`
|
||||
UPDATE "ql3"."worker_sessions"
|
||||
SET version = version + 1,
|
||||
available_slots = CASE WHEN status = 'draining' THEN 0 ELSE $5 END,
|
||||
last_heartbeat_at_ms = $6,
|
||||
lease_expires_at_ms = $7,
|
||||
updated_at_ms = $6
|
||||
WHERE worker_id = $1 AND session_id = $2 AND generation = $3 AND version = $4
|
||||
RETURNING ${SELECT_COLUMNS}
|
||||
`,
|
||||
[
|
||||
command.workerId,
|
||||
command.sessionId,
|
||||
command.generation,
|
||||
command.expectedVersion,
|
||||
command.availableSlots,
|
||||
nowMs,
|
||||
nowMs + command.leaseDurationMs,
|
||||
],
|
||||
);
|
||||
if (result.rows.length !== 1) {
|
||||
throw new WorkerSessionFenceRejectedError(command.workerId, 'version_mismatch');
|
||||
}
|
||||
const worker = record(result.rows[0]!);
|
||||
if (credential) {
|
||||
await observeCredentialDelivery(client, credential, worker, nowMs);
|
||||
}
|
||||
return worker;
|
||||
});
|
||||
}
|
||||
|
||||
async transition(
|
||||
command: TransitionWorkerSessionCommand,
|
||||
): Promise<WorkerSessionRecord> {
|
||||
return this.transitionInternal(command);
|
||||
}
|
||||
|
||||
async transitionAuthenticated(
|
||||
command: TransitionWorkerSessionCommand,
|
||||
credential: AuthenticatedWorkerCredentialIdentity,
|
||||
): Promise<WorkerSessionRecord> {
|
||||
return this.transitionInternal(command, credential);
|
||||
}
|
||||
|
||||
private async transitionInternal(
|
||||
command: TransitionWorkerSessionCommand,
|
||||
credential?: AuthenticatedWorkerCredentialIdentity,
|
||||
): Promise<WorkerSessionRecord> {
|
||||
assertTransition(command);
|
||||
return this.mutate(command.workerId, async (client, nowMs, current) => {
|
||||
fence(current, command, nowMs);
|
||||
const result = await client.query<Row>(
|
||||
`
|
||||
UPDATE "ql3"."worker_sessions"
|
||||
SET version = version + 1,
|
||||
status = $5::varchar,
|
||||
available_slots = 0,
|
||||
lease_expires_at_ms = CASE
|
||||
WHEN $5::varchar = 'offline' THEN $6
|
||||
ELSE lease_expires_at_ms
|
||||
END,
|
||||
updated_at_ms = $6
|
||||
WHERE worker_id = $1 AND session_id = $2 AND generation = $3 AND version = $4
|
||||
RETURNING ${SELECT_COLUMNS}
|
||||
`,
|
||||
[
|
||||
command.workerId,
|
||||
command.sessionId,
|
||||
command.generation,
|
||||
command.expectedVersion,
|
||||
command.status,
|
||||
nowMs,
|
||||
],
|
||||
);
|
||||
if (result.rows.length !== 1) {
|
||||
throw new WorkerSessionFenceRejectedError(command.workerId, 'version_mismatch');
|
||||
}
|
||||
const worker = record(result.rows[0]!);
|
||||
if (credential) {
|
||||
await observeCredentialDelivery(client, credential, worker, nowMs);
|
||||
}
|
||||
return worker;
|
||||
});
|
||||
}
|
||||
|
||||
async listAvailable(
|
||||
options: Readonly<{ afterWorkerId?: string; limit?: number }> = {},
|
||||
): Promise<AvailableWorkerSessionPage> {
|
||||
const limit = options.limit ?? 16;
|
||||
if (!Number.isInteger(limit) || limit < 1 || limit > MAX_AVAILABLE_WORKER_PAGE_SIZE) {
|
||||
throw new RangeError(
|
||||
`Worker page limit must be between 1 and ${MAX_AVAILABLE_WORKER_PAGE_SIZE}`,
|
||||
);
|
||||
}
|
||||
if (options.afterWorkerId !== undefined) assertWorkerId(options.afterWorkerId);
|
||||
const result = await this.pool.query<Row>(
|
||||
`
|
||||
WITH observation AS (
|
||||
SELECT floor(extract(epoch FROM statement_timestamp()) * 1000)::bigint AS observed_at_ms
|
||||
), available AS (
|
||||
SELECT ${SELECT_COLUMNS}
|
||||
FROM "ql3"."worker_sessions", observation
|
||||
WHERE status = 'online'
|
||||
AND available_slots > 0
|
||||
AND lease_expires_at_ms > observation.observed_at_ms
|
||||
AND ($1::varchar IS NULL OR worker_id > $1)
|
||||
ORDER BY worker_id
|
||||
LIMIT $2
|
||||
)
|
||||
SELECT observation.observed_at_ms AS "observedAtMs", available.*
|
||||
FROM observation LEFT JOIN available ON TRUE
|
||||
ORDER BY available."workerId"
|
||||
`,
|
||||
[options.afterWorkerId ?? null, limit + 1],
|
||||
);
|
||||
if (result.rows.length < 1 || result.rows.length > limit + 1) {
|
||||
throw new TypeError('PostgreSQL Worker page violated its bound');
|
||||
}
|
||||
const observed = integer(result.rows[0]!, 'observedAtMs');
|
||||
const workers = result.rows[0]!.workerId === null
|
||||
? []
|
||||
: result.rows.map((row) => record(row));
|
||||
const page = workers.slice(0, limit);
|
||||
const last = page.at(-1);
|
||||
return Object.freeze({
|
||||
observedAtMs: observed,
|
||||
workers: Object.freeze(page),
|
||||
truncated: workers.length > limit,
|
||||
...(workers.length > limit && last ? { nextCursor: last.workerId } : {}),
|
||||
});
|
||||
}
|
||||
|
||||
private async mutate<T>(
|
||||
workerId: string,
|
||||
work: (
|
||||
client: PostgresClient,
|
||||
nowMs: number,
|
||||
current: WorkerSessionRecord | null,
|
||||
) => Promise<T>,
|
||||
): Promise<T> {
|
||||
const client = await this.pool.connect();
|
||||
try {
|
||||
await begin(client);
|
||||
const currentResult = await client.query<Row>(
|
||||
`SELECT ${SELECT_COLUMNS} FROM "ql3"."worker_sessions" WHERE worker_id = $1 FOR UPDATE`,
|
||||
[workerId],
|
||||
);
|
||||
const nowMs = await observedAtMs(client);
|
||||
const result = await work(
|
||||
client,
|
||||
nowMs,
|
||||
currentResult.rows[0] ? record(currentResult.rows[0]) : null,
|
||||
);
|
||||
await client.query('COMMIT');
|
||||
return result;
|
||||
} catch (error) {
|
||||
try {
|
||||
await client.query('ROLLBACK');
|
||||
} catch {
|
||||
// Preserve the originating failure.
|
||||
}
|
||||
throw error;
|
||||
} finally {
|
||||
client.release();
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user