john
581 lines
18 KiB
JavaScript
581 lines
18 KiB
JavaScript
import pg from 'pg';
|
|
import { createInMemoryExecutorJobStore } from './executor-gateway.mjs';
|
|
|
|
const { Pool } = pg;
|
|
const DEFAULT_EXECUTOR_SCHEMA = 'memind_orchestrator';
|
|
const EXECUTOR_JOB_TABLE = 'executor_jobs';
|
|
const EXECUTOR_JOB_EVENT_TABLE = 'executor_job_events';
|
|
const EXECUTOR_WORKER_TABLE = 'executor_workers';
|
|
|
|
function normalizeMode(value) {
|
|
const mode = String(value ?? 'postgres').trim().toLowerCase();
|
|
if (mode === 'memory' || mode === 'postgres') return mode;
|
|
throw new Error(`Unsupported executor job store mode: ${value}`);
|
|
}
|
|
|
|
function normalizeSchema(value) {
|
|
const schema = String(value ?? DEFAULT_EXECUTOR_SCHEMA).trim();
|
|
if (!/^[a-z_][a-z0-9_]{0,62}$/i.test(schema)) {
|
|
throw new Error('Invalid Executor Gateway PostgreSQL schema');
|
|
}
|
|
return schema;
|
|
}
|
|
|
|
function parseState(value) {
|
|
if (value == null) return null;
|
|
if (typeof value === 'object') return structuredClone(value);
|
|
try {
|
|
return JSON.parse(String(value));
|
|
} catch {
|
|
return null;
|
|
}
|
|
}
|
|
|
|
function projectRow(row) {
|
|
return parseState(row?.state_json);
|
|
}
|
|
|
|
function projectEventRow(row) {
|
|
return parseState(row?.event_json);
|
|
}
|
|
|
|
function eventWithSequence(event, jobId, sequence) {
|
|
return {
|
|
...structuredClone(event),
|
|
jobId,
|
|
sequence,
|
|
};
|
|
}
|
|
|
|
async function appendEvent(client, eventTable, jobId, event) {
|
|
if (!event) return null;
|
|
const sequenceResult = await client.query(
|
|
`SELECT COALESCE(MAX(sequence), 0) + 1 AS next_sequence
|
|
FROM ${eventTable}
|
|
WHERE job_id = $1`,
|
|
[jobId],
|
|
);
|
|
const sequence = Number(sequenceResult.rows?.[0]?.next_sequence ?? 1);
|
|
const persistedEvent = eventWithSequence(event, jobId, sequence);
|
|
await client.query(
|
|
`INSERT INTO ${eventTable}
|
|
(event_id, job_id, sequence, event_json, created_at)
|
|
VALUES ($1, $2, $3, $4::jsonb, $5)`,
|
|
[
|
|
persistedEvent.eventId,
|
|
jobId,
|
|
sequence,
|
|
JSON.stringify(persistedEvent),
|
|
persistedEvent.timestamp,
|
|
],
|
|
);
|
|
return persistedEvent;
|
|
}
|
|
|
|
async function withTransaction(pool, operation) {
|
|
if (typeof pool?.connect !== 'function') {
|
|
throw new Error('PostgreSQL Executor Job Store requires transactional pool.connect');
|
|
}
|
|
const client = await pool.connect();
|
|
try {
|
|
await client.query('BEGIN');
|
|
const result = await operation(client);
|
|
await client.query('COMMIT');
|
|
return result;
|
|
} catch (error) {
|
|
try {
|
|
await client.query('ROLLBACK');
|
|
} catch {
|
|
// Preserve the transactional operation failure.
|
|
}
|
|
throw error;
|
|
} finally {
|
|
client.release?.();
|
|
}
|
|
}
|
|
|
|
export function createPostgresExecutorJobStore({
|
|
pool,
|
|
schema = DEFAULT_EXECUTOR_SCHEMA,
|
|
} = {}) {
|
|
if (!pool?.query) throw new Error('PostgreSQL Executor Job Store requires a pool');
|
|
const normalizedSchema = normalizeSchema(schema);
|
|
const table = `"${normalizedSchema}"."${EXECUTOR_JOB_TABLE}"`;
|
|
const eventTable = `"${normalizedSchema}"."${EXECUTOR_JOB_EVENT_TABLE}"`;
|
|
const workerTable = `"${normalizedSchema}"."${EXECUTOR_WORKER_TABLE}"`;
|
|
|
|
return {
|
|
kind: 'postgres',
|
|
durable: true,
|
|
|
|
async ensureSchema() {
|
|
await pool.query(`CREATE SCHEMA IF NOT EXISTS "${normalizedSchema}"`);
|
|
await pool.query(`
|
|
CREATE TABLE IF NOT EXISTS ${table} (
|
|
job_id VARCHAR(128) PRIMARY KEY,
|
|
idempotency_key VARCHAR(200) NOT NULL UNIQUE,
|
|
request_fingerprint CHAR(64) NOT NULL,
|
|
state_json JSONB NOT NULL,
|
|
created_at BIGINT NOT NULL,
|
|
updated_at BIGINT NOT NULL
|
|
)
|
|
`);
|
|
await pool.query(`
|
|
CREATE INDEX IF NOT EXISTS executor_jobs_updated_at_idx
|
|
ON ${table} (updated_at DESC)
|
|
`);
|
|
await pool.query(`
|
|
CREATE TABLE IF NOT EXISTS ${eventTable} (
|
|
event_id VARCHAR(128) PRIMARY KEY,
|
|
job_id VARCHAR(128) NOT NULL REFERENCES ${table} (job_id) ON DELETE CASCADE,
|
|
sequence INTEGER NOT NULL,
|
|
event_json JSONB NOT NULL,
|
|
created_at BIGINT NOT NULL,
|
|
UNIQUE (job_id, sequence)
|
|
)
|
|
`);
|
|
await pool.query(`
|
|
CREATE INDEX IF NOT EXISTS executor_job_events_job_sequence_idx
|
|
ON ${eventTable} (job_id, sequence)
|
|
`);
|
|
await pool.query(`
|
|
INSERT INTO ${eventTable} (event_id, job_id, sequence, event_json, created_at)
|
|
SELECT
|
|
'evt_' || md5(job_id || ':snapshot-import'),
|
|
job_id,
|
|
1,
|
|
jsonb_build_object(
|
|
'version', 'executor-job-event-v1',
|
|
'eventId', 'evt_' || md5(job_id || ':snapshot-import'),
|
|
'jobId', job_id,
|
|
'sequence', 1,
|
|
'type', 'executor_job_snapshot_imported',
|
|
'timestamp', created_at,
|
|
'data', jsonb_build_object(
|
|
'status', state_json->>'status',
|
|
'reason', state_json->>'reason'
|
|
)
|
|
),
|
|
created_at
|
|
FROM ${table} jobs
|
|
WHERE NOT EXISTS (
|
|
SELECT 1 FROM ${eventTable} events WHERE events.job_id = jobs.job_id
|
|
)
|
|
ON CONFLICT DO NOTHING
|
|
`);
|
|
await pool.query(`
|
|
CREATE TABLE IF NOT EXISTS ${workerTable} (
|
|
worker_id VARCHAR(128) PRIMARY KEY,
|
|
state_json JSONB NOT NULL,
|
|
last_seen_at BIGINT NOT NULL
|
|
)
|
|
`);
|
|
await pool.query(`
|
|
CREATE INDEX IF NOT EXISTS executor_workers_last_seen_idx
|
|
ON ${workerTable} (last_seen_at DESC)
|
|
`);
|
|
},
|
|
|
|
async getById(jobId) {
|
|
const result = await pool.query(
|
|
`SELECT state_json FROM ${table} WHERE job_id = $1 LIMIT 1`,
|
|
[String(jobId ?? '').trim()],
|
|
);
|
|
return projectRow(result.rows?.[0]);
|
|
},
|
|
|
|
async getByIdempotencyKey(key) {
|
|
const result = await pool.query(
|
|
`SELECT state_json FROM ${table} WHERE idempotency_key = $1 LIMIT 1`,
|
|
[String(key ?? '').trim()],
|
|
);
|
|
return projectRow(result.rows?.[0]);
|
|
},
|
|
|
|
async createIfAbsent(record, { initialEvent = null } = {}) {
|
|
const created = await withTransaction(pool, async (client) => {
|
|
const inserted = await client.query(
|
|
`INSERT INTO ${table}
|
|
(job_id, idempotency_key, request_fingerprint, state_json, created_at, updated_at)
|
|
VALUES ($1, $2, $3, $4::jsonb, $5, $6)
|
|
ON CONFLICT DO NOTHING
|
|
RETURNING state_json`,
|
|
[
|
|
record.jobId,
|
|
record.idempotencyKey,
|
|
record.requestFingerprint,
|
|
JSON.stringify(record),
|
|
record.createdAt,
|
|
record.updatedAt,
|
|
],
|
|
);
|
|
const insertedRecord = projectRow(inserted.rows?.[0]);
|
|
if (!insertedRecord || !initialEvent) return insertedRecord;
|
|
const event = eventWithSequence(initialEvent, record.jobId, 1);
|
|
await client.query(
|
|
`INSERT INTO ${eventTable}
|
|
(event_id, job_id, sequence, event_json, created_at)
|
|
VALUES ($1, $2, $3, $4::jsonb, $5)`,
|
|
[
|
|
event.eventId,
|
|
record.jobId,
|
|
event.sequence,
|
|
JSON.stringify(event),
|
|
event.timestamp,
|
|
],
|
|
);
|
|
return insertedRecord;
|
|
});
|
|
if (created) return { created: true, record: created };
|
|
|
|
const existingResult = await pool.query(
|
|
`SELECT state_json FROM ${table} WHERE idempotency_key = $1 LIMIT 1`,
|
|
[record.idempotencyKey],
|
|
);
|
|
const existing = projectRow(existingResult.rows?.[0]);
|
|
if (existing) return { created: false, record: existing };
|
|
|
|
const error = new Error(`Executor job id already exists: ${record.jobId}`);
|
|
error.code = 'EXECUTOR_JOB_ID_CONFLICT';
|
|
error.status = 409;
|
|
throw error;
|
|
},
|
|
|
|
async update(record, { event = null } = {}) {
|
|
return withTransaction(pool, async (client) => {
|
|
const locked = await client.query(
|
|
`SELECT job_id FROM ${table} WHERE job_id = $1 FOR UPDATE`,
|
|
[record.jobId],
|
|
);
|
|
if (!locked.rows?.[0]) return null;
|
|
const result = await client.query(
|
|
`UPDATE ${table}
|
|
SET state_json = $2::jsonb,
|
|
updated_at = $3
|
|
WHERE job_id = $1
|
|
RETURNING state_json`,
|
|
[record.jobId, JSON.stringify(record), record.updatedAt],
|
|
);
|
|
await appendEvent(client, eventTable, record.jobId, event);
|
|
return projectRow(result.rows?.[0]);
|
|
});
|
|
},
|
|
|
|
async claimNext({
|
|
workerId,
|
|
executors,
|
|
now,
|
|
leaseToken,
|
|
leaseDurationMs,
|
|
eventFactory,
|
|
}) {
|
|
return withTransaction(pool, async (client) => {
|
|
const selected = await client.query(
|
|
`SELECT state_json
|
|
FROM ${table}
|
|
WHERE (
|
|
state_json->>'status' = 'queued'
|
|
OR (
|
|
state_json->>'status' = 'retryable'
|
|
AND COALESCE((state_json->>'nextAttemptAt')::BIGINT, 0) <= $1
|
|
)
|
|
)
|
|
AND state_json->>'executor' = ANY($2::text[])
|
|
ORDER BY created_at ASC, job_id ASC
|
|
FOR UPDATE SKIP LOCKED
|
|
LIMIT 1`,
|
|
[now, executors],
|
|
);
|
|
const current = projectRow(selected.rows?.[0]);
|
|
if (!current) return null;
|
|
const updated = {
|
|
...current,
|
|
status: 'leased',
|
|
reason: null,
|
|
attempts: Number(current.attempts ?? 0) + 1,
|
|
nextAttemptAt: null,
|
|
lease: {
|
|
token: leaseToken,
|
|
workerId,
|
|
acquiredAt: now,
|
|
heartbeatAt: now,
|
|
expiresAt: now + leaseDurationMs,
|
|
},
|
|
updatedAt: now,
|
|
completedAt: null,
|
|
};
|
|
const result = await client.query(
|
|
`UPDATE ${table}
|
|
SET state_json = $2::jsonb,
|
|
updated_at = $3
|
|
WHERE job_id = $1
|
|
RETURNING state_json`,
|
|
[updated.jobId, JSON.stringify(updated), updated.updatedAt],
|
|
);
|
|
await appendEvent(
|
|
client,
|
|
eventTable,
|
|
updated.jobId,
|
|
typeof eventFactory === 'function' ? eventFactory(updated, current) : null,
|
|
);
|
|
return projectRow(result.rows?.[0]);
|
|
});
|
|
},
|
|
|
|
async transition(jobId, {
|
|
expectedStatuses = [],
|
|
leaseToken = null,
|
|
updater,
|
|
eventFactory = null,
|
|
} = {}) {
|
|
return withTransaction(pool, async (client) => {
|
|
const locked = await client.query(
|
|
`SELECT state_json FROM ${table} WHERE job_id = $1 FOR UPDATE`,
|
|
[String(jobId ?? '').trim()],
|
|
);
|
|
const current = projectRow(locked.rows?.[0]);
|
|
if (!current) return { applied: false, record: null, reason: 'not_found' };
|
|
if (expectedStatuses.length && !expectedStatuses.includes(current.status)) {
|
|
return { applied: false, record: current, reason: 'status_mismatch' };
|
|
}
|
|
if (leaseToken != null && current.lease?.token !== leaseToken) {
|
|
return { applied: false, record: current, reason: 'lease_mismatch' };
|
|
}
|
|
const updated = updater(structuredClone(current));
|
|
if (!updated || updated.jobId !== current.jobId) {
|
|
throw new Error('Executor job transition must preserve job identity');
|
|
}
|
|
const result = await client.query(
|
|
`UPDATE ${table}
|
|
SET state_json = $2::jsonb,
|
|
updated_at = $3
|
|
WHERE job_id = $1
|
|
RETURNING state_json`,
|
|
[current.jobId, JSON.stringify(updated), updated.updatedAt],
|
|
);
|
|
await appendEvent(
|
|
client,
|
|
eventTable,
|
|
current.jobId,
|
|
typeof eventFactory === 'function' ? eventFactory(updated, current) : null,
|
|
);
|
|
return {
|
|
applied: true,
|
|
record: projectRow(result.rows?.[0]),
|
|
reason: null,
|
|
};
|
|
});
|
|
},
|
|
|
|
async listExpiredLeases({ now, limit = 100 } = {}) {
|
|
const pageSize = Math.min(500, Math.max(1, Number(limit) || 100));
|
|
const result = await pool.query(
|
|
`SELECT state_json
|
|
FROM ${table}
|
|
WHERE state_json->>'status' = ANY($1::text[])
|
|
AND COALESCE((state_json->'lease'->>'expiresAt')::BIGINT, 0) <= $2
|
|
ORDER BY COALESCE((state_json->'lease'->>'expiresAt')::BIGINT, 0) ASC
|
|
LIMIT $3`,
|
|
[['leased', 'running'], now, pageSize],
|
|
);
|
|
return (result.rows ?? []).map(projectRow).filter(Boolean);
|
|
},
|
|
|
|
async getQueueStats({ now = Date.now() } = {}) {
|
|
const result = await pool.query(
|
|
`SELECT
|
|
COUNT(*)::BIGINT AS total,
|
|
COUNT(*) FILTER (
|
|
WHERE state_json->>'status' = 'queued'
|
|
OR (
|
|
state_json->>'status' = 'retryable'
|
|
AND COALESCE((state_json->>'nextAttemptAt')::BIGINT, 0) <= $1
|
|
)
|
|
)::BIGINT AS claimable,
|
|
COUNT(*) FILTER (
|
|
WHERE state_json->>'status' = ANY($2::text[])
|
|
AND COALESCE((state_json->'lease'->>'expiresAt')::BIGINT, 0) <= $1
|
|
)::BIGINT AS expired_leases,
|
|
state_json->>'status' AS status,
|
|
COUNT(*)::BIGINT AS status_count
|
|
FROM ${table}
|
|
GROUP BY state_json->>'status'`,
|
|
[now, ['leased', 'running']],
|
|
);
|
|
const counts = {};
|
|
let total = 0;
|
|
let claimable = 0;
|
|
let expiredLeases = 0;
|
|
for (const row of result.rows ?? []) {
|
|
if (row.status) counts[row.status] = Number(row.status_count ?? 0);
|
|
total += Number(row.status_count ?? 0);
|
|
claimable += Number(row.claimable ?? 0);
|
|
expiredLeases += Number(row.expired_leases ?? 0);
|
|
}
|
|
return { total, counts, claimable, expiredLeases };
|
|
},
|
|
|
|
async getAdmissionStats({
|
|
tenantId = null,
|
|
userId = null,
|
|
now = Date.now(),
|
|
windowMs = 60_000,
|
|
} = {}) {
|
|
const result = await pool.query(
|
|
`SELECT
|
|
COUNT(*) FILTER (
|
|
WHERE state_json->>'status' = ANY($1::text[])
|
|
)::BIGINT AS global_active,
|
|
COUNT(*) FILTER (
|
|
WHERE state_json->>'status' = ANY($1::text[])
|
|
AND $2::text IS NOT NULL
|
|
AND state_json->'request'->'subject'->>'tenantId' = $2
|
|
)::BIGINT AS tenant_active,
|
|
COUNT(*) FILTER (
|
|
WHERE state_json->>'status' = ANY($1::text[])
|
|
AND $3::text IS NOT NULL
|
|
AND state_json->'request'->'subject'->>'userId' = $3
|
|
)::BIGINT AS user_active,
|
|
COUNT(*) FILTER (
|
|
WHERE $2::text IS NOT NULL
|
|
AND state_json->'request'->'subject'->>'tenantId' = $2
|
|
AND created_at >= $4
|
|
)::BIGINT AS tenant_recent,
|
|
COUNT(*) FILTER (
|
|
WHERE $3::text IS NOT NULL
|
|
AND state_json->'request'->'subject'->>'userId' = $3
|
|
AND created_at >= $4
|
|
)::BIGINT AS user_recent
|
|
FROM ${table}`,
|
|
[
|
|
['queued', 'retryable', 'leased', 'running'],
|
|
tenantId,
|
|
userId,
|
|
now - windowMs,
|
|
],
|
|
);
|
|
const row = result.rows?.[0] ?? {};
|
|
return {
|
|
globalActive: Number(row.global_active ?? 0),
|
|
tenantActive: Number(row.tenant_active ?? 0),
|
|
userActive: Number(row.user_active ?? 0),
|
|
tenantRecent: Number(row.tenant_recent ?? 0),
|
|
userRecent: Number(row.user_recent ?? 0),
|
|
};
|
|
},
|
|
|
|
async recordWorkerHeartbeat(record) {
|
|
await pool.query(
|
|
`INSERT INTO ${workerTable} (worker_id, state_json, last_seen_at)
|
|
VALUES ($1, $2::jsonb, $3)
|
|
ON CONFLICT (worker_id) DO UPDATE
|
|
SET state_json = EXCLUDED.state_json,
|
|
last_seen_at = EXCLUDED.last_seen_at`,
|
|
[record.workerId, JSON.stringify(record), record.lastSeenAt],
|
|
);
|
|
return structuredClone(record);
|
|
},
|
|
|
|
async listWorkers({ now = Date.now(), staleAfterMs = 60_000 } = {}) {
|
|
const result = await pool.query(
|
|
`SELECT state_json
|
|
FROM ${workerTable}
|
|
ORDER BY last_seen_at DESC
|
|
LIMIT 500`,
|
|
);
|
|
return (result.rows ?? [])
|
|
.map(projectRow)
|
|
.filter(Boolean)
|
|
.map((worker) => ({
|
|
...worker,
|
|
stale: Number(worker.lastSeenAt ?? 0) < now - staleAfterMs,
|
|
}));
|
|
},
|
|
|
|
async listEvents(jobId, { after = 0, limit = 500 } = {}) {
|
|
const cursor = Math.max(0, Number(after) || 0);
|
|
const pageSize = Math.min(500, Math.max(1, Number(limit) || 500));
|
|
const result = await pool.query(
|
|
`SELECT event_json
|
|
FROM ${eventTable}
|
|
WHERE job_id = $1 AND sequence > $2
|
|
ORDER BY sequence ASC
|
|
LIMIT $3`,
|
|
[String(jobId ?? '').trim(), cursor, pageSize],
|
|
);
|
|
return (result.rows ?? []).map(projectEventRow).filter(Boolean);
|
|
},
|
|
};
|
|
}
|
|
|
|
export async function createExecutorJobPersistence({
|
|
mode = process.env.MEMIND_ORCHESTRATOR_CHECKPOINT_MODE,
|
|
connectionString = process.env.MEMIND_ORCHESTRATOR_DATABASE_URL,
|
|
schema = process.env.MEMIND_ORCHESTRATOR_DATABASE_SCHEMA,
|
|
poolFactory = (config) => new Pool(config),
|
|
} = {}) {
|
|
const normalizedMode = normalizeMode(mode);
|
|
if (normalizedMode === 'memory') {
|
|
const store = createInMemoryExecutorJobStore();
|
|
return {
|
|
kind: 'memory',
|
|
durable: false,
|
|
store,
|
|
async probe() {
|
|
return true;
|
|
},
|
|
async close() {},
|
|
};
|
|
}
|
|
|
|
const normalizedConnectionString = String(connectionString ?? '').trim();
|
|
if (!normalizedConnectionString) {
|
|
const error = new Error(
|
|
'MEMIND_ORCHESTRATOR_DATABASE_URL is required for PostgreSQL Executor Job Store',
|
|
);
|
|
error.code = 'ORCHESTRATOR_DATABASE_URL_REQUIRED';
|
|
throw error;
|
|
}
|
|
|
|
const normalizedSchema = normalizeSchema(schema);
|
|
const pool = poolFactory({ connectionString: normalizedConnectionString });
|
|
const store = createPostgresExecutorJobStore({ pool, schema: normalizedSchema });
|
|
try {
|
|
await store.ensureSchema();
|
|
} catch (error) {
|
|
try {
|
|
await pool.end?.();
|
|
} catch {
|
|
// Preserve the schema initialization failure.
|
|
}
|
|
throw error;
|
|
}
|
|
return {
|
|
kind: 'postgres',
|
|
durable: true,
|
|
store,
|
|
async probe() {
|
|
await pool.query('SELECT 1 AS ok');
|
|
return true;
|
|
},
|
|
async close() {
|
|
await pool.end?.();
|
|
},
|
|
};
|
|
}
|
|
|
|
export const executorJobStoreInternals = {
|
|
DEFAULT_EXECUTOR_SCHEMA,
|
|
EXECUTOR_JOB_TABLE,
|
|
EXECUTOR_JOB_EVENT_TABLE,
|
|
EXECUTOR_WORKER_TABLE,
|
|
appendEvent,
|
|
eventWithSequence,
|
|
normalizeMode,
|
|
normalizeSchema,
|
|
parseState,
|
|
projectEventRow,
|
|
projectRow,
|
|
withTransaction,
|
|
};
|