mirror of
https://github.com/whyour/qinglong.git
synced 2026-09-21 18:08:20 +08:00
252 lines
11 KiB
TypeScript
252 lines
11 KiB
TypeScript
// 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();
|
|
}
|
|
}
|
|
}
|