fix(mcp): устранить конкуренцию фоновых записей при опросе
This commit is contained in:
1 parent
2ac36a6ca5
commit
cc190a381c
9 files changed
+167
-54
No files matched your search
@@ -2,7 +2,7 @@ import { randomUUID } from "node:crypto"
|
||||
import { Client, StreamableHTTPClientTransport } from "@modelcontextprotocol/client"
|
||||
import { mcpTemporaryGrantSchema } from "@mmapp/contracts/mcp"
|
||||
import type { AuthUser } from "../../lib/permissions.js"
|
||||
import { audit, assertConnectionVisible, getConnection, transaction, validatePrincipal } from "./store.js"
|
||||
import { audit, assertConnectionVisible, connectionWrite, getConnection, validatePrincipal } from "./store.js"
|
||||
import { canonicalMcpUrl, McpError, randomToken, tokenDigest } from "./security.js"
|
||||
|
||||
export interface ProtocolProbe {
|
||||
@@ -21,7 +21,7 @@ const capabilitiesTool = "mm_capabilities_get"
|
||||
* that operation. The ordinary connection credentials are never recovered. */
|
||||
export async function withMcpProbeCredential<T>(connectionId: string, actor: AuthUser, work: (token: string, readCapabilities: boolean) => Promise<T>): Promise<T> {
|
||||
const credentialId = randomUUID(), grantId = randomUUID(), token = randomToken("access")
|
||||
const readCapabilities = await transaction(async client => {
|
||||
const readCapabilities = await connectionWrite(async client => {
|
||||
await client.query("SELECT id FROM mcp_connections WHERE id=$1 FOR UPDATE", [connectionId])
|
||||
const connection = await getConnection(connectionId, client)
|
||||
assertConnectionVisible(connection, actor)
|
||||
@@ -42,9 +42,11 @@ export async function withMcpProbeCredential<T>(connectionId: string, actor: Aut
|
||||
let completed = false
|
||||
try { const result = await work(token, readCapabilities); completed = true; return result }
|
||||
finally {
|
||||
await transaction(async client => {
|
||||
await client.query("DELETE FROM mcp_grants WHERE id=$1 AND connection_id=$2", [grantId, connectionId])
|
||||
await connectionWrite(async client => {
|
||||
// Use the same connection → credential → grant lock order as revocation.
|
||||
await client.query("SELECT id FROM mcp_connections WHERE id=$1 FOR UPDATE", [connectionId])
|
||||
await client.query("DELETE FROM mcp_credentials WHERE id=$1 AND connection_id=$2", [credentialId, connectionId])
|
||||
await client.query("DELETE FROM mcp_grants WHERE id=$1 AND connection_id=$2", [grantId, connectionId])
|
||||
await audit(client, actor.id, "connection.handshake.end", connectionId, { credentialId, completed })
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import { withPartitionWrite } from "../../services/partition-write-transaction.js"
|
||||
import { withBackgroundPartitionWrite } from "../../services/partition-write-transaction.js"
|
||||
import { randomUUID } from "node:crypto"
|
||||
import { z } from "zod"
|
||||
import type { PoolClient } from "pg"
|
||||
@@ -440,7 +440,7 @@ export function registerScopedRecoveryService() {
|
||||
|
||||
/** Restart maintenance is read-only with respect to the network. */
|
||||
export async function recoverScopedRecoveryIntents() {
|
||||
await withPartitionWrite(({ client }) => client.query(`WITH interrupted AS (
|
||||
await withBackgroundPartitionWrite(({ client }) => client.query(`WITH interrupted AS (
|
||||
SELECT intent.source_operation_id,intent.source_step_index
|
||||
FROM mcp_scoped_recovery_intents intent
|
||||
WHERE intent.state='intent' AND (intent.deadline<=now() OR EXISTS (
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import { withPartitionWrite } from "../../services/partition-write-transaction.js"
|
||||
import { partitionWriteFailureDiagnostic, withBackgroundPartitionWrite } from "../../services/partition-write-transaction.js"
|
||||
import { pollErrorLog } from "../../services/server-poll-error-provenance.js"
|
||||
import { vrfPreparationDiagnostic } from "./vrf-preparation-diagnostic.js"
|
||||
import { createVrfRuntimeDiagnosticBoundary } from "../../services/vrf-execution-diagnostic.js"
|
||||
@@ -782,7 +782,7 @@ export async function recoverMcpOperations() {
|
||||
// A second backend must not interrupt a runner still holding its session lease.
|
||||
const { rows } = await pool.query("SELECT id FROM mcp_operations WHERE state IN ('pending','running') AND created_at<now()-interval '2 minutes' ORDER BY created_at,id LIMIT 100")
|
||||
for (const row of rows) {
|
||||
await withPartitionWrite(async ({ client }) => {
|
||||
await withBackgroundPartitionWrite(async ({ client }) => {
|
||||
const { rows: [lock] } = await client.query("SELECT pg_try_advisory_xact_lock(hashtextextended($1,0)) AS locked", ["mcp-operation:" + row.id])
|
||||
if (!lock.locked) return
|
||||
// Defer the whole transition when either row is busy.
|
||||
@@ -796,18 +796,25 @@ export async function recoverMcpOperations() {
|
||||
await audit(client, "mcp-recovery", "operation.interrupted", target.connection_id, { operationId: row.id, retryAllowed: false })
|
||||
})
|
||||
}
|
||||
await withPartitionWrite(({ client }) => client.query(`WITH expired AS (
|
||||
await withBackgroundPartitionWrite(({ client }) => client.query(`WITH expired AS (
|
||||
SELECT id FROM mcp_plans
|
||||
WHERE expires_at<now() AND state IN ('draft','approved','awaiting-approval')
|
||||
ORDER BY expires_at,id LIMIT 1000 FOR UPDATE SKIP LOCKED
|
||||
)
|
||||
UPDATE mcp_plans SET state='expired' FROM expired WHERE mcp_plans.id=expired.id`))
|
||||
await withPartitionWrite(({ client }) => client.query(`WITH expired AS (
|
||||
UPDATE mcp_plans SET state='expired' FROM expired WHERE mcp_plans.id=expired.id`)).catch(error => {
|
||||
const diagnostic = partitionWriteFailureDiagnostic(error)
|
||||
if (diagnostic) {
|
||||
try { console.error(JSON.stringify({ event: "mcp.worker.plan-expiry-failed", ...diagnostic })) }
|
||||
catch { /* Preserve the original transaction failure. */ }
|
||||
}
|
||||
throw error
|
||||
})
|
||||
await withBackgroundPartitionWrite(({ client }) => client.query(`WITH expired AS (
|
||||
SELECT id FROM mcp_read_leases WHERE expires_at<now()
|
||||
ORDER BY expires_at,id LIMIT 1000 FOR UPDATE SKIP LOCKED
|
||||
)
|
||||
DELETE FROM mcp_read_leases USING expired WHERE mcp_read_leases.id=expired.id`))
|
||||
await withPartitionWrite(({ client }) => client.query(`WITH interrupted AS (
|
||||
await withBackgroundPartitionWrite(({ client }) => client.query(`WITH interrupted AS (
|
||||
SELECT intent.operation_id FROM mcp_recovery_intents intent
|
||||
JOIN mcp_operations operation ON operation.id=intent.operation_id
|
||||
WHERE intent.state='intent' AND operation.state NOT IN ('pending','running')
|
||||
|
||||
@@ -17,7 +17,7 @@ import { audit, getMcpSettings } from "./store.js"
|
||||
import { accessGeneration, decryptSecretExisting, redact } from "./security.js"
|
||||
import { tickRoutingProfiles } from "./routing-profiles.js"
|
||||
import type { ChangeEvidence } from "./action-types.js"
|
||||
import { withPartitionWrite } from "../../services/partition-write-transaction.js"
|
||||
import { PartitionWriterBusy, withBackgroundPartitionWrite } from "../../services/partition-write-transaction.js"
|
||||
import { retainManualInverseChargesWithinOwner } from "./routing-manual-inverse-store.js"
|
||||
import { recoverScopedRecoveryIntents } from "./routing-community-scoped-recovery.js"
|
||||
|
||||
@@ -86,7 +86,7 @@ export async function detectMcpDrift() {
|
||||
if (expectation.matches(current) !== false) return
|
||||
const observedAt = new Date().toISOString()
|
||||
|
||||
await withPartitionWrite(async ({ client }) => {
|
||||
await withBackgroundPartitionWrite(async ({ client }) => {
|
||||
const { rows: [locked] } = await client.query<DriftCheckpoint>(
|
||||
`${driftCheckpointSql} FOR SHARE OF o,p,c,s,k,device`, args)
|
||||
if (!isDeepStrictEqual(locked, checkpoint) || await accessGeneration() !== generation) return
|
||||
@@ -107,28 +107,28 @@ export async function detectMcpDrift() {
|
||||
const mcpRetentionBatchSize = 1000
|
||||
|
||||
async function retainMcpData() {
|
||||
await withPartitionWrite(async ({ client }) => {
|
||||
await withBackgroundPartitionWrite(async ({ client }) => {
|
||||
const { requestsRetentionDays } = await getMcpSettings(client)
|
||||
await client.query(`WITH expired AS (
|
||||
SELECT id FROM mcp_requests WHERE created_at<now()-make_interval(days=>$1) AND status NOT IN ('running')
|
||||
ORDER BY created_at,id LIMIT $2 FOR UPDATE SKIP LOCKED
|
||||
) DELETE FROM mcp_requests USING expired WHERE mcp_requests.id=expired.id`, [requestsRetentionDays, mcpRetentionBatchSize])
|
||||
})
|
||||
await withPartitionWrite(async ({ client }) => {
|
||||
await withBackgroundPartitionWrite(async ({ client }) => {
|
||||
const { auditRetentionDays } = await getMcpSettings(client)
|
||||
await client.query(`WITH expired AS (
|
||||
SELECT id FROM mcp_audit WHERE created_at<now()-make_interval(days=>$1)
|
||||
ORDER BY created_at,id LIMIT $2 FOR UPDATE SKIP LOCKED
|
||||
) DELETE FROM mcp_audit USING expired WHERE mcp_audit.id=expired.id`, [auditRetentionDays, mcpRetentionBatchSize])
|
||||
})
|
||||
await withPartitionWrite(({ client }) => client.query(`WITH expired AS (
|
||||
await withBackgroundPartitionWrite(({ client }) => client.query(`WITH expired AS (
|
||||
SELECT id FROM mcp_artifacts WHERE expires_at<now()
|
||||
ORDER BY expires_at,id LIMIT $1 FOR UPDATE
|
||||
) DELETE FROM mcp_artifacts USING expired WHERE mcp_artifacts.id=expired.id`, [mcpRetentionBatchSize]))
|
||||
await withPartitionWrite(({ client }) => client.query("DELETE FROM mcp_oauth_requests WHERE expires_at<now()-interval '1 day'"))
|
||||
await withPartitionWrite(({ client }) => client.query("DELETE FROM mcp_budget_buckets WHERE bucket<now()-interval '2 days'"))
|
||||
await withPartitionWrite(({ client }) => retainManualInverseChargesWithinOwner(client))
|
||||
await withPartitionWrite(async ({ client }) => {
|
||||
await withBackgroundPartitionWrite(({ client }) => client.query("DELETE FROM mcp_oauth_requests WHERE expires_at<now()-interval '1 day'"))
|
||||
await withBackgroundPartitionWrite(({ client }) => client.query("DELETE FROM mcp_budget_buckets WHERE bucket<now()-interval '2 days'"))
|
||||
await withBackgroundPartitionWrite(({ client }) => retainManualInverseChargesWithinOwner(client))
|
||||
await withBackgroundPartitionWrite(async ({ client }) => {
|
||||
const { requestsRetentionDays } = await getMcpSettings(client)
|
||||
await client.query(`WITH candidates AS MATERIALIZED (SELECT source_operation_id FROM mcp_routing_forward_charges
|
||||
WHERE created_at<now()-make_interval(days=>$1) ORDER BY created_at,source_operation_id LIMIT 1000),
|
||||
@@ -136,10 +136,10 @@ async function retainMcpData() {
|
||||
WHERE o.id=c.source_operation_id AND o.state IN('pending','running')))
|
||||
DELETE FROM mcp_routing_forward_charges c USING expired e WHERE c.source_operation_id=e.source_operation_id`, [requestsRetentionDays])
|
||||
})
|
||||
await withPartitionWrite(({ client }) => client.query("DELETE FROM mcp_routing_observations WHERE created_at<now()-interval '7 days'"))
|
||||
await withPartitionWrite(({ client }) => client.query("DELETE FROM mcp_path_checks WHERE created_at<now()-interval '7 days'"))
|
||||
await withPartitionWrite(({ client }) => client.query("DELETE FROM mcp_unauthorized_buckets WHERE bucket<now()-interval '7 days'"))
|
||||
await withPartitionWrite(({ client }) => client.query("DELETE FROM mcp_portal_events WHERE created_at<now()-interval '7 days'"))
|
||||
await withBackgroundPartitionWrite(({ client }) => client.query("DELETE FROM mcp_routing_observations WHERE created_at<now()-interval '7 days'"))
|
||||
await withBackgroundPartitionWrite(({ client }) => client.query("DELETE FROM mcp_path_checks WHERE created_at<now()-interval '7 days'"))
|
||||
await withBackgroundPartitionWrite(({ client }) => client.query("DELETE FROM mcp_unauthorized_buckets WHERE bucket<now()-interval '7 days'"))
|
||||
await withBackgroundPartitionWrite(({ client }) => client.query("DELETE FROM mcp_portal_events WHERE created_at<now()-interval '7 days'"))
|
||||
}
|
||||
|
||||
export async function tickMcpWorkers() {
|
||||
@@ -159,6 +159,8 @@ export async function tickMcpWorkers() {
|
||||
await retireOrdinaryMcpOperations()
|
||||
if ((await getMcpSettings()).enabled) { await detectMcpDrift(); await tickRoutingProfiles() }
|
||||
await maintainServerProbeOperationRetirement()
|
||||
} catch (error) {
|
||||
if (!(error instanceof PartitionWriterBusy)) throw error
|
||||
} finally {
|
||||
if (locked) try { await lease.query("SELECT pg_advisory_unlock(hashtextextended('mcp-workers',0))") } catch { unsafe = true }
|
||||
lease.release(unsafe)
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
import assert from "node:assert/strict"
|
||||
import type { PoolClient } from "pg"
|
||||
import { sql } from "drizzle-orm"
|
||||
import { admitPartitionWriter, withPartitionWrite, PartitionWriteError, isPartitionWriteOutcomeUnknown, isPartitionWriteClient, PARTITION_WRITER_ADVISORY_KEY, type PartitionWriteOwner } from "./partition-write-transaction.js"
|
||||
import { admitPartitionWriter, withBackgroundPartitionWrite, withPartitionWrite, PartitionWriteError, PartitionWriterBusy, isPartitionWriteOutcomeUnknown, isPartitionWriteClient, partitionWriteFailureDiagnostic, PARTITION_WRITER_ADVISORY_KEY, type PartitionWriteOwner } from "./partition-write-transaction.js"
|
||||
|
||||
function fixture() {
|
||||
const tryAdmission = "SELECT pg_catalog.pg_try_advisory_xact_lock_shared($1::bigint) AS acquired"
|
||||
function fixture(acquired?: boolean) {
|
||||
const statements: Array<{ text: string; values: unknown[] }> = [], releases: boolean[] = []
|
||||
let fail: string | undefined, committed = 0, connects = 0
|
||||
let commitCommand: string | undefined = "COMMIT"
|
||||
@@ -12,8 +13,8 @@ function fixture() {
|
||||
const text = typeof raw === "string" ? raw : raw.text
|
||||
statements.push({ text, values: typeof raw === "string" ? values : raw.values ?? values })
|
||||
if (text === "COMMIT") committed++
|
||||
if (text === fail) throw Error("Injected uncertain SQL ACK")
|
||||
return { rows: [{ value: 1 }], rowCount: 1, command: text === "COMMIT" ? commitCommand : "SELECT" }
|
||||
if (text === fail) throw Object.assign(Error("Injected uncertain SQL ACK"), { code: "55P03" })
|
||||
return { rows: [text === tryAdmission ? { acquired } : { value: 1 }], rowCount: 1, command: text === "COMMIT" ? commitCommand : text === "ROLLBACK" ? "ROLLBACK" : "SELECT" }
|
||||
},
|
||||
release: (destroy: boolean) => { releases.push(destroy) },
|
||||
} as unknown as PoolClient
|
||||
@@ -35,6 +36,23 @@ await test("caller-owned transaction admission never starts or commits lifecycle
|
||||
const f = fixture(); await admitPartitionWriter(f.client)
|
||||
assert.equal(f.statements.length, 3); assert(f.statements[2].text.includes("pg_advisory_xact_lock_shared")); assert.equal(f.releases.length, 0)
|
||||
})
|
||||
await test("a definite busy writer rolls back without running scheduled work", async () => {
|
||||
const f = fixture(false)
|
||||
await assert.rejects(withBackgroundPartitionWrite(async () => assert.fail("scheduled SQL ran"), { connect: f.connect }), PartitionWriterBusy)
|
||||
assert.deepEqual(f.statements.slice(3).map(statement => statement.text), [tryAdmission, "ROLLBACK"])
|
||||
assert.deepEqual(f.releases, [false])
|
||||
})
|
||||
await test("background admission validates the lock reply and surfaces SQL failures", async () => {
|
||||
for (const failure of [undefined, tryAdmission, "ROLLBACK"]) {
|
||||
const f = fixture(failure === undefined ? undefined : false)
|
||||
if (failure) f.fail(failure)
|
||||
await assert.rejects(withBackgroundPartitionWrite(async () => assert.fail("scheduled SQL ran"), { connect: f.connect }), PartitionWriteError)
|
||||
assert.deepEqual(f.releases, [true])
|
||||
}
|
||||
const f = fixture(true)
|
||||
await withBackgroundPartitionWrite(async owner => { await owner.client.query("SELECT 1") }, { connect: f.connect })
|
||||
assert.deepEqual(f.statements.slice(3).map(statement => statement.text), [tryAdmission, "SELECT 1", "COMMIT"])
|
||||
})
|
||||
for (const boundary of ["BEGIN", "COMMIT"] as const) await test("lost " + boundary + " ACK destroys without rollback or retry", async () => {
|
||||
const f = fixture(); f.fail(boundary)
|
||||
await assert.rejects(withPartitionWrite(async owner => { await owner.client.query("INSERT INTO traffic_samples VALUES (1)") }, { connect: f.connect }), error => isPartitionWriteOutcomeUnknown(error) && error.automaticReplay === false && error.code === "PARTITION_WRITE_UNKNOWN")
|
||||
@@ -59,6 +77,23 @@ await test("caught query failure still prevents commit", async () => {
|
||||
await assert.rejects(withPartitionWrite(async owner => { await owner.client.query("SELECT fail").catch(() => {}) }, { connect: f.connect }), error => error instanceof PartitionWriteError && !error.outcomeUnknown)
|
||||
assert(!f.statements.some(row => row.text === "COMMIT")); assert.deepEqual(f.releases, [true])
|
||||
})
|
||||
for (const [statement, stage] of [
|
||||
["SELECT pg_advisory_xact_lock_shared($1::bigint)", "admission"],
|
||||
["SELECT 1", "work"],
|
||||
["COMMIT", "commit"],
|
||||
] as const) await test("private SQLSTATE identifies transaction stage " + stage, async () => {
|
||||
const f = fixture(); f.fail(statement)
|
||||
await assert.rejects(withPartitionWrite(async owner => {
|
||||
await owner.client.query("SELECT 1")
|
||||
}, { connect: f.connect }), (error: unknown) => {
|
||||
assert(error instanceof PartitionWriteError)
|
||||
assert.deepEqual(partitionWriteFailureDiagnostic(error), { stage, sqlState: "55P03" })
|
||||
assert.equal(error.outcomeUnknown, stage === "commit")
|
||||
assert.equal(error.automaticReplay, false)
|
||||
assert.doesNotMatch(JSON.stringify(error), /55P03|Injected|SELECT/)
|
||||
return true
|
||||
})
|
||||
})
|
||||
await test("callback raw query remains guarded and completes before commit", async () => {
|
||||
const f = fixture()
|
||||
await withPartitionWrite(async owner => { await new Promise<void>((resolve, reject) => { owner.client.query("SELECT 1", error => error ? reject(error) : resolve()) }) }, { connect: f.connect })
|
||||
|
||||
@@ -18,6 +18,9 @@ export interface PartitionWriteOptions {
|
||||
connect?: () => Promise<PoolClient>
|
||||
}
|
||||
export interface PartitionWriteOwner { client: PoolClient; db: NodePgDatabase<typeof schema> }
|
||||
export class PartitionWriterBusy extends Error {
|
||||
constructor() { super("Partition writer lock is busy") }
|
||||
}
|
||||
export class PartitionWriteError extends Error {
|
||||
readonly automaticReplay = false
|
||||
readonly code: "PARTITION_WRITE_UNAVAILABLE" | "PARTITION_WRITE_UNKNOWN"
|
||||
@@ -27,6 +30,19 @@ export class PartitionWriteError extends Error {
|
||||
}
|
||||
}
|
||||
export const isPartitionWriteOutcomeUnknown = (error: unknown): error is PartitionWriteError => error instanceof PartitionWriteError && error.outcomeUnknown
|
||||
type PartitionWriteStage = "checkout" | "begin" | "admission" | "work" | "commit"
|
||||
const failureDiagnostics = new WeakMap<PartitionWriteError, { stage: PartitionWriteStage; sqlState?: string }>()
|
||||
|
||||
export function partitionWriteFailureDiagnostic(error: unknown) {
|
||||
return error instanceof PartitionWriteError ? failureDiagnostics.get(error) : undefined
|
||||
}
|
||||
|
||||
function safeSqlState(error: unknown): string | undefined {
|
||||
if (!(error instanceof Error)) return undefined
|
||||
const code = Object.getOwnPropertyDescriptor(error, "code")?.value
|
||||
// SQLSTATE excludes transport codes such as EPIPE and ECONNRESET.
|
||||
return typeof code === "string" && /^[0-9A-Z]{5}$/.test(code) && /[0-9]/.test(code) ? code : undefined
|
||||
}
|
||||
const boundary = /^\s*(?:BEGIN|COMMIT|END|ROLLBACK|SAVEPOINT|RELEASE|START\s+TRANSACTION|PREPARE\s+TRANSACTION)\b/i
|
||||
function settings(options: PartitionWriteOptions) {
|
||||
const now = Date.now(), deadlineAt = options.deadlineAt ?? now + 30000
|
||||
@@ -45,7 +61,9 @@ function remaining(deadlineAt: number) {
|
||||
* Its bounded server timeout also applies when a surrounding Drizzle tx owns
|
||||
* the lifecycle. The surrounding owner must roll back/destroy on rejection.
|
||||
*/
|
||||
export async function admitPartitionWriter(client: { query(text: string, values?: [number]): Promise<unknown> }, options: Omit<PartitionWriteOptions, "connect"> = {}): Promise<void> {
|
||||
type PartitionAdmissionClient = { query(text: string, values?: [number]): Promise<unknown> }
|
||||
|
||||
async function acquirePartitionWriter(client: PartitionAdmissionClient, options: Omit<PartitionWriteOptions, "connect">, acquire: () => Promise<boolean>) {
|
||||
const current = settings(options)
|
||||
const check = async () => {
|
||||
if (current.signal?.aborted) throw new PartitionWriteError()
|
||||
@@ -58,8 +76,16 @@ export async function admitPartitionWriter(client: { query(text: string, values?
|
||||
await check()
|
||||
await client.query(`SET LOCAL statement_timeout='${remaining(current.deadlineAt)}ms'`)
|
||||
await check()
|
||||
await client.query("SELECT pg_advisory_xact_lock_shared($1::bigint)", [PARTITION_WRITER_ADVISORY_KEY])
|
||||
const acquired = await acquire()
|
||||
await check()
|
||||
return acquired
|
||||
}
|
||||
|
||||
export async function admitPartitionWriter(client: PartitionAdmissionClient, options: Omit<PartitionWriteOptions, "connect"> = {}): Promise<void> {
|
||||
await acquirePartitionWriter(client, options, async () => {
|
||||
await client.query("SELECT pg_advisory_xact_lock_shared($1::bigint)", [PARTITION_WRITER_ADVISORY_KEY])
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
/** Owns exactly one checked-out backend for raw SQL AND Drizzle. Admission is
|
||||
@@ -68,11 +94,13 @@ export async function admitPartitionWriter(client: { query(text: string, values?
|
||||
* after release. Any failed/uncertain boundary destroys the owner. No retries
|
||||
* and no ROLLBACK is sent after an uncertain BEGIN/COMMIT acknowledgement.
|
||||
*/
|
||||
export async function withPartitionWrite<T>(work: (owner: PartitionWriteOwner) => Promise<T>, options: PartitionWriteOptions = {}): Promise<T> {
|
||||
async function runPartitionWrite<T>(work: (owner: PartitionWriteOwner) => Promise<T>, options: PartitionWriteOptions, deferIfBusy: boolean): Promise<T> {
|
||||
const current = settings(options)
|
||||
let client: PoolClient | undefined, released = false, active = false, accepting = false
|
||||
let beginSent = false, beginAcknowledged = false, commitSent = false, committed = false, ended = false
|
||||
let queryFailed = false
|
||||
let stage: PartitionWriteStage = "checkout"
|
||||
let querySqlState: string | undefined
|
||||
const inflight = new Set<Promise<unknown>>()
|
||||
const release = (destroy: boolean) => { active = false; accepting = false; if (client && !released) { released = true; client.release(destroy) } }
|
||||
const check = async () => {
|
||||
@@ -96,11 +124,29 @@ export async function withPartitionWrite<T>(work: (owner: PartitionWriteOwner) =
|
||||
})]).catch(error => { ended = true; void pending.then(late => late.release(true), () => {}); throw error })
|
||||
if (ended) { release(true); throw new PartitionWriteError() }
|
||||
await check()
|
||||
stage = "begin"
|
||||
beginSent = true
|
||||
await Promise.race([client.query("BEGIN"), stopped])
|
||||
beginAcknowledged = true
|
||||
stage = "admission"
|
||||
await check()
|
||||
await Promise.race([admitPartitionWriter(client, current), stopped])
|
||||
const admissionClient = client
|
||||
if (deferIfBusy) {
|
||||
const acquired = await Promise.race([acquirePartitionWriter(client, current, async () => {
|
||||
const { rows } = await admissionClient.query<{ acquired: boolean }>(
|
||||
"SELECT pg_catalog.pg_try_advisory_xact_lock_shared($1::bigint) AS acquired", [PARTITION_WRITER_ADVISORY_KEY])
|
||||
if (rows.length !== 1 || typeof rows[0].acquired !== "boolean") throw new PartitionWriteError()
|
||||
return rows[0].acquired
|
||||
}), stopped])
|
||||
if (!acquired) {
|
||||
const rollback = await Promise.race([client.query("ROLLBACK"), stopped])
|
||||
if (rollback.command !== "ROLLBACK") throw new PartitionWriteError()
|
||||
await check()
|
||||
release(false)
|
||||
throw new PartitionWriterBusy()
|
||||
}
|
||||
} else await Promise.race([admitPartitionWriter(client, current), stopped])
|
||||
stage = "work"
|
||||
active = true; accepting = true
|
||||
const rawClient = client
|
||||
const guardedQuery = (...args: unknown[]) => {
|
||||
@@ -114,7 +160,9 @@ export async function withPartitionWrite<T>(work: (owner: PartitionWriteOwner) =
|
||||
if (!statement || boundary.test(statement) || /;\s*\S/.test(statement)) throw new PartitionWriteError()
|
||||
await check()
|
||||
if (!active || !accepting || released) throw new PartitionWriteError()
|
||||
const result = await Reflect.apply(rawClient.query, rawClient, args)
|
||||
let result: unknown
|
||||
try { result = await Reflect.apply(rawClient.query, rawClient, args) }
|
||||
catch (error) { querySqlState ??= safeSqlState(error); throw error }
|
||||
await check()
|
||||
return result
|
||||
}
|
||||
@@ -136,19 +184,30 @@ export async function withPartitionWrite<T>(work: (owner: PartitionWriteOwner) =
|
||||
// Await already-started work, then close both SQL handles BEFORE COMMIT.
|
||||
await Promise.race([Promise.all([...inflight]), stopped]); active = false
|
||||
if (queryFailed) throw new PartitionWriteError()
|
||||
await check(); commitSent = true
|
||||
await check(); stage = "commit"; commitSent = true
|
||||
const commitResult = await Promise.race([client.query("COMMIT"), stopped])
|
||||
if (commitResult.command !== "COMMIT") throw new PartitionWriteError(true)
|
||||
committed = true
|
||||
await check(); release(false)
|
||||
return result
|
||||
} catch (error) {
|
||||
if (error instanceof PartitionWriterBusy && released && stage === "admission") throw error
|
||||
const unknown = beginSent && !beginAcknowledged || commitSent || committed
|
||||
destroy()
|
||||
if (error instanceof PartitionWriteError && error.outcomeUnknown === unknown) throw error
|
||||
throw new PartitionWriteError(unknown)
|
||||
const failure = error instanceof PartitionWriteError && error.outcomeUnknown === unknown ? error : new PartitionWriteError(unknown)
|
||||
failureDiagnostics.set(failure, { stage, sqlState: querySqlState ?? safeSqlState(error) })
|
||||
throw failure
|
||||
} finally {
|
||||
clearTimeout(timer); current.signal?.removeEventListener("abort", destroy)
|
||||
if (!released) release(true)
|
||||
}
|
||||
}
|
||||
|
||||
export function withPartitionWrite<T>(work: (owner: PartitionWriteOwner) => Promise<T>, options: PartitionWriteOptions = {}): Promise<T> {
|
||||
return runPartitionWrite(work, options, false)
|
||||
}
|
||||
|
||||
/** Defer scheduled maintenance only on a definite negative lock result. */
|
||||
export function withBackgroundPartitionWrite<T>(work: (owner: PartitionWriteOwner) => Promise<T>, options: PartitionWriteOptions = {}): Promise<T> {
|
||||
return runPartitionWrite(work, options, true)
|
||||
}
|
||||
@@ -6,7 +6,7 @@ import { fingerprint } from "../modules/mcp/security.js"
|
||||
import { assertServerPollArchiveAccounting } from "./server-poll-archive-accounting.js"
|
||||
import { SERVER_POLL_ARCHIVE_READ_SQL } from "./server-poll-terminal-compaction.js"
|
||||
import { openServerPollRetiredArchive, projectServerPollRetiredArchive, type LockedServerPollRetirementSource } from "./server-poll-retired-archive.js"
|
||||
import { PollStorageRefused, runPollStorageTransaction } from "./server-poll-storage.js"
|
||||
import { PollStorageRefused, runPollStorageTransaction, tryAdmitPollStorageMaintenanceWriter } from "./server-poll-storage.js"
|
||||
import type { ProbeStorageSql } from "./server-probe-journal-maintenance.js"
|
||||
import { connectProbeStorage } from "./server-probe-storage-connection.js"
|
||||
import { sharedDatabaseGate } from "./shared-database-gate.js"
|
||||
@@ -59,8 +59,8 @@ async function admitRetirementWriter(sql:ProbeStorageSql,deadlineAt:number) {
|
||||
await admitPartitionWriter(sql,{deadlineAt:Math.min(deadlineAt,Date.now()+5000)})
|
||||
await sql.query("SELECT pg_catalog.set_config('lock_timeout',$1,true)",[settings[0].lock_timeout])
|
||||
}
|
||||
const transaction = <T>(deadlineAt:number,mode:"scan"|"write",work:(sql:ProbeStorageSql)=>Promise<T>) =>
|
||||
runPollStorageTransaction(deps,{deadlineAt},async()=>{},work,mode==="write"?sql=>admitRetirementWriter(sql,deadlineAt):undefined)
|
||||
const writeTransaction = <T>(deadlineAt:number,work:(sql:ProbeStorageSql)=>Promise<T>) =>
|
||||
runPollStorageTransaction(deps,{deadlineAt},async()=>{},work,sql=>admitRetirementWriter(sql,deadlineAt))
|
||||
type Target = z.infer<typeof request>
|
||||
|
||||
async function prepare(sql: ProbeStorageSql,{receiptId,oldArchiveProof}:Target) {
|
||||
@@ -93,7 +93,7 @@ async function prepare(sql: ProbeStorageSql,{receiptId,oldArchiveProof}:Target)
|
||||
async function retire(target:Target,deadlineAt:number) {
|
||||
const {receiptId,oldArchiveProof}=target
|
||||
const token=randomUUID()
|
||||
await transaction(deadlineAt,"write",async sql=>{
|
||||
await writeTransaction(deadlineAt,async sql=>{
|
||||
// Check the live hold predicates, all locks and both archive proofs before
|
||||
// spending the irreversible one-shot claim.
|
||||
await prepare(sql,target)
|
||||
@@ -103,7 +103,7 @@ async function retire(target:Target,deadlineAt:number) {
|
||||
})
|
||||
// No destructive transaction after an unacknowledged claim. A claim is
|
||||
// never cleared, even when a later checkout/hold/COMMIT is inconclusive.
|
||||
return transaction(deadlineAt,"write",async sql=>{
|
||||
return writeTransaction(deadlineAt,async sql=>{
|
||||
const claimed=await sql.query(`SELECT receipt_id FROM public.server_poll_retirement_claims
|
||||
WHERE receipt_id=$1 AND old_archive_proof=$2 AND claim_token=$3 AND claim_xid<>pg_current_xact_id()`,
|
||||
[receiptId,oldArchiveProof,token])
|
||||
@@ -150,8 +150,12 @@ let cursor:{finishedAt:string;receiptId:string}|undefined
|
||||
export function retireCompletedServerPollOperations() {
|
||||
return sharedDatabaseGate.request(async()=>{
|
||||
const deadlineAt=Date.now()+20000
|
||||
const rows=await transaction(deadlineAt,"scan",async sql=>(await sql.query(CANDIDATES,
|
||||
[cursor?.finishedAt??null,cursor?.receiptId??null])).rows)
|
||||
// A live POLL may hold both locks. Skip before requesting its storage key.
|
||||
const rows=await runPollStorageTransaction<Record<string,unknown>[]|null>(deps,{deadlineAt},async()=>{},async sql=>(await sql.query(CANDIDATES,
|
||||
[cursor?.finishedAt??null,cursor?.receiptId??null])).rows,async sql=>{
|
||||
if(!await tryAdmitPollStorageMaintenanceWriter(sql,"transaction"))return {deferred:null}
|
||||
})
|
||||
if(rows===null)return {scanned:0,retired:0,automaticReplay:false as const}
|
||||
requireThat(rows.length<=8)
|
||||
let retired=0,scanned=0
|
||||
for(const row of rows){
|
||||
|
||||
@@ -32,7 +32,7 @@ function fixture() {
|
||||
const now = Date.now(), receiptId = randomUUID(), operationId = randomUUID(), connectionId = randomUUID()
|
||||
let state: State = { claim: null, archive: null, audits: [] }, connects = 0, transactions=0, commits = 0, createCalls = 0, openCalls = 0, accountCalls = 0
|
||||
let eligible = true, sourceMissing = false, accountingFails = false, preflight: "valid" | "oversize" | "invalid" = "valid", secondOversize = false
|
||||
let lostCommit = 0, commitPersists = true, badPublishedSize = false, badPublishedText = false, auditFails = false, writer: boolean|Error=true
|
||||
let lostCommit = 0, commitPersists = true, badPublishedSize = false, badPublishedText = false, badPublishedId = false, auditFails = false, writer: boolean|Error=true
|
||||
const commands: Command[] = [], releases: { tx: number; unsafe: boolean }[] = [], events: string[] = []
|
||||
class ArchiveRefused extends Error {}
|
||||
const forbiddenConnector = async (): Promise<never> => { throw Error("Production connector must never run") }
|
||||
@@ -64,7 +64,7 @@ function fixture() {
|
||||
async openServerPollTerminalArchive(archive: Row) {
|
||||
openCalls++; events.push("open:" + openCalls)
|
||||
assert(archives.has(String(archive.ciphertext)))
|
||||
return { sourceText: badPublishedText && openCalls === 2 ? "altered" : archives.get(String(archive.ciphertext)), metadata: { receipt_id: receiptId },
|
||||
return { sourceText: badPublishedText ? "altered" : archives.get(String(archive.ciphertext)), metadata: { receipt_id: badPublishedId ? operationId : receiptId },
|
||||
source: { journal: { intents: [{}, {}], acks: [{}, {}] } } }
|
||||
},
|
||||
},
|
||||
@@ -124,10 +124,11 @@ function fixture() {
|
||||
} }
|
||||
return { run: () => api.compactCompletedServerPollStorageForTests(deps), commands, releases, events, state: () => structuredClone(state),
|
||||
counts: () => ({ connects, commits, createCalls, openCalls, accountCalls }), Unknown: storage.PollStorageUnknown as new () => Error,setWriter:(value:boolean|Error)=>{writer=value},
|
||||
configure(options: { eligible?: boolean; sourceMissing?: boolean; accountingFails?: boolean; preflight?: typeof preflight; secondOversize?: boolean; lostCommit?: number; commitPersists?: boolean; badPublishedSize?: boolean; badPublishedText?: boolean; auditFails?: boolean }) {
|
||||
configure(options: { eligible?: boolean; sourceMissing?: boolean; accountingFails?: boolean; preflight?: typeof preflight; secondOversize?: boolean; lostCommit?: number; commitPersists?: boolean; badPublishedSize?: boolean; badPublishedText?: boolean; badPublishedId?: boolean; auditFails?: boolean }) {
|
||||
eligible = options.eligible ?? eligible; sourceMissing = options.sourceMissing ?? sourceMissing; accountingFails = options.accountingFails ?? accountingFails
|
||||
preflight = options.preflight ?? preflight; secondOversize = options.secondOversize ?? secondOversize; lostCommit = options.lostCommit ?? lostCommit
|
||||
commitPersists = options.commitPersists ?? commitPersists; badPublishedSize = options.badPublishedSize ?? badPublishedSize; badPublishedText = options.badPublishedText ?? badPublishedText; auditFails = options.auditFails ?? auditFails
|
||||
badPublishedId = options.badPublishedId ?? badPublishedId
|
||||
} }
|
||||
}
|
||||
const claimed = (f: ReturnType<typeof fixture>) => f.commands.filter(command => command.text.startsWith("UPDATE server_poll_storage_admissions"))
|
||||
@@ -179,7 +180,7 @@ test("phase1 COMMIT ACK loss never starts phase2 or inserts archive, for either
|
||||
test("normal two acknowledged COMMITs publish exact readback and audit in the same transaction", async () => {
|
||||
const f = fixture(), result = await f.run()
|
||||
assert.equal(result.compacted, 1); assert.equal(result.automaticReplay, false)
|
||||
assert.deepEqual(f.counts(), { connects: 1, commits: 2, createCalls: 2, openCalls: 2, accountCalls: 3 })
|
||||
assert.deepEqual(f.counts(), { connects: 1, commits: 2, createCalls: 2, openCalls: 1, accountCalls: 3 })
|
||||
const firstCommit = f.events.indexOf("1:COMMIT"), secondBegin = f.events.indexOf("2:BEGIN")
|
||||
assert(firstCommit >= 0 && firstCommit < secondBegin)
|
||||
for (const tx of [1, 2]) {
|
||||
@@ -190,7 +191,9 @@ test("normal two acknowledged COMMITs publish exact readback and audit in the sa
|
||||
const second = f.commands.filter(command => command.tx === 2).map(command => command.text)
|
||||
const insert = second.findIndex(text => text.startsWith("INSERT INTO server_poll_archives")), read = second.findIndex(text => text.startsWith("SELECT r.*,")), audit = second.findIndex(text => text.startsWith("INSERT INTO mcp_audit"))
|
||||
assert(insert < read && read < audit && audit < second.lastIndexOf(ACCOUNTING) && second.lastIndexOf(ACCOUNTING) < second.indexOf("COMMIT"))
|
||||
assert(f.events.indexOf("open:2") < f.events.findIndex(event => event.startsWith("2:INSERT INTO mcp_audit")))
|
||||
const opened = f.events.indexOf("open:1")
|
||||
assert(f.events.findIndex(event => event.startsWith("2:SELECT r.*,")) < opened &&
|
||||
opened < f.events.findIndex(event => event.startsWith("2:INSERT INTO mcp_audit")))
|
||||
assert.equal(f.state().audits.length, 1); assert.deepEqual(f.state().audits[0].payload, { operationId: f.state().archive!.execution_id, receiptId: f.state().archive!.receipt_id,
|
||||
archiveBytes: f.state().archive!.ciphertext_bytes, sourceJournalRows: 6, originalRetentionDays: 30, automaticReplay: false })
|
||||
assert.deepEqual(f.releases, [{ tx: 1, unsafe: false }])
|
||||
@@ -216,8 +219,8 @@ test("phase2 size growth preserves spent full charge and subsequent candidate ex
|
||||
assert(!f.commands.some(command => /^(?:DELETE|UPDATE server_poll_storage_usage)/.test(command.text)))
|
||||
})
|
||||
|
||||
test("published size/text mismatch or audit failure rolls back archive but keeps acknowledged claim", async () => {
|
||||
for (const options of [{ badPublishedSize: true }, { badPublishedText: true }, { auditFails: true }]) {
|
||||
test("published size, text or identity mismatch or audit failure rolls back archive but keeps acknowledged claim", async () => {
|
||||
for (const options of [{ badPublishedSize: true }, { badPublishedText: true }, { badPublishedId: true }, { auditFails: true }]) {
|
||||
const f = fixture(); f.configure(options); await assert.rejects(f.run(), model.PollStorageRefused)
|
||||
assert(f.state().claim); assert.equal(f.state().archive, null); assert.equal(f.state().audits.length, 0)
|
||||
assert.equal(f.commands.filter(command => command.tx === 2 && command.text === "ROLLBACK").length, 1)
|
||||
|
||||
@@ -141,14 +141,15 @@ async function compactOnWriter(writer:ReturnType<typeof compactionWriter>,firstD
|
||||
const text = rows[0].source_text, now = new Date(rows[0].observed_at as string | Date).getTime()
|
||||
const archive = await createServerPollTerminalArchive(text, now)
|
||||
if (!archive) return 0 // Oversize retains the FULL charge and spent claim.
|
||||
const check = await openServerPollTerminalArchive(archive, now)
|
||||
if (check.sourceText !== text || check.metadata.receipt_id !== claimed.receiptId) fail()
|
||||
if (archive.receipt_id !== claimed.receiptId) fail()
|
||||
const command = insert(archive), inserted = (await sql.query(command.text, command.values)).rows
|
||||
if (inserted.length !== 1 || inserted[0].receipt_id !== claimed.receiptId) fail()
|
||||
const published = (await sql.query(SERVER_POLL_ARCHIVE_READ_SQL, [claimed.receiptId])).rows
|
||||
if (published.length !== 1 || published[0].actual_ciphertext_bytes !== archive.ciphertext_bytes) fail()
|
||||
// Authenticate the persisted representation before commit. A refusal rolls
|
||||
// back both the archive insert and its trigger's journal/charge transition.
|
||||
const authenticated = await openServerPollTerminalArchive(published[0], now)
|
||||
if (authenticated.sourceText !== text) fail()
|
||||
if (authenticated.sourceText !== text || authenticated.metadata.receipt_id !== claimed.receiptId) fail()
|
||||
await sql.query("INSERT INTO mcp_audit(actor_id,action,connection_id,payload) VALUES($1,$2,$3,$4::jsonb)",
|
||||
["mcp-worker", "poll.history.compacted", archive.connection_id, JSON.stringify({ operationId: archive.execution_id,
|
||||
receiptId: archive.receipt_id, archiveBytes: archive.ciphertext_bytes,
|
||||
|
||||
Reference in new issue
Block a user