diff --git a/src/common/kafka-events.ts b/src/common/kafka-events.ts index f5492b5..912633f 100644 --- a/src/common/kafka-events.ts +++ b/src/common/kafka-events.ts @@ -89,6 +89,20 @@ export async function publishAgentRunKafkaMessage(config: AgentRunKafkaConfig, m } } +export async function publishAgentRunKafkaMessageBatch(config: AgentRunKafkaConfig, messages: AgentRunKafkaMessage[]): Promise { + if (!config.enabled || messages.length === 0) return; + if (config.brokers.length === 0) throw new Error("AGENTRUN_KAFKA_BOOTSTRAP_SERVERS is empty"); + const topics = new Set(messages.map((message) => message.topic)); + if (topics.size !== 1) throw new Error("Kafka message batch must use one topic"); + const producer = await producerForConfig(config); + try { + await sendAgentRunKafkaMessageBatch(producer, messages); + } catch (error) { + await invalidateProducer(producer); + throw error; + } +} + export async function publishAgentRunKafkaMessagesOnce(config: AgentRunKafkaConfig, messages: AgentRunKafkaMessage[]): Promise { if (!config.enabled) return; if (messages.length === 0) return; @@ -281,16 +295,22 @@ function defaultProducerFactory(config: AgentRunKafkaConfig): Producer { } async function sendAgentRunKafkaMessage(producer: Producer, message: AgentRunKafkaMessage): Promise { + await sendAgentRunKafkaMessageBatch(producer, [message]); +} + +async function sendAgentRunKafkaMessageBatch(producer: Producer, messages: AgentRunKafkaMessage[]): Promise { + const topic = messages[0]?.topic; + if (!topic) return; await producer.send({ - topic: message.topic, - messages: [{ + topic, + messages: messages.map((message) => ({ key: message.key, value: JSON.stringify(message.value), headers: { "x-values-printed": "false", ...(message.headers ?? {}), }, - }], + })), }); } diff --git a/src/mgr/kafka-outbox-relay.ts b/src/mgr/kafka-outbox-relay.ts index b469692..caccd33 100644 --- a/src/mgr/kafka-outbox-relay.ts +++ b/src/mgr/kafka-outbox-relay.ts @@ -1,5 +1,5 @@ import type { JsonRecord, KafkaEventOutboxRecord } from "../common/types.js"; -import { closeAgentRunKafkaProducer, publishAgentRunKafkaMessage, type AgentRunKafkaConfig, type AgentRunKafkaMessage } from "../common/kafka-events.js"; +import { closeAgentRunKafkaProducer, publishAgentRunKafkaMessageBatch, type AgentRunKafkaConfig, type AgentRunKafkaMessage } from "../common/kafka-events.js"; import { redactText } from "../common/redaction.js"; import type { AgentRunStore } from "./store.js"; @@ -19,53 +19,74 @@ export function kafkaOutboxRelayStatus(): JsonRecord { return { lastCycle: lastRelayCycle, valuesPrinted: false }; } -export async function relayKafkaEventOutboxOnce(input: { store: AgentRunStore; options: KafkaOutboxRelayOptions; publish?: (config: AgentRunKafkaConfig, message: AgentRunKafkaMessage) => Promise }): Promise { +export async function relayKafkaEventOutboxOnce(input: { store: AgentRunStore; options: KafkaOutboxRelayOptions; publishBatch?: (config: AgentRunKafkaConfig, messages: AgentRunKafkaMessage[]) => Promise }): Promise { if (!input.options.enabled) return { action: "kafka-event-outbox-relay", enabled: false, claimedCount: 0, deliveredCount: 0, retryCount: 0, staleCount: 0, valuesPrinted: false }; - const claimed = await input.store.claimKafkaEventOutbox({ owner: input.options.owner, leaseMs: input.options.leaseMs, limit: input.options.batchSize }); + const cycleStartedAt = Date.now(); + const claimStartedAt = Date.now(); + let claimed: KafkaEventOutboxRecord[]; + try { + claimed = await input.store.claimKafkaEventOutbox({ owner: input.options.owner, leaseMs: input.options.leaseMs, limit: input.options.batchSize }); + } catch (error) { + const claimDurationMs = Date.now() - claimStartedAt; + lastRelayCycle = { at: new Date().toISOString(), claimedCount: 0, partitionKeyCount: 0, publishedCount: 0, deliveredCount: 0, retryCount: 0, staleCount: 0, durationMs: Math.max(1, Date.now() - cycleStartedAt), claimDurationMs, publishDurationMs: 0, settleDurationMs: 0, firstError: { code: "kafka-outbox-claim-failed", message: redactText(error instanceof Error ? error.message : String(error)).slice(0, 500), valuesPrinted: false }, valuesPrinted: false }; + throw error; + } + const claimDurationMs = Date.now() - claimStartedAt; const items = claimed.slice().sort((a, b) => a.outboxSeq - b.outboxSeq); - const startedAt = Date.now(); const groups = groupByPartitionKey(items); let deliveredCount = 0; + let publishedCount = 0; let retryCount = 0; let staleCount = 0; let firstError: JsonRecord | null = null; - await Promise.all(groups.map(async (group) => { - let blocked = false; - for (const item of group) { - if (blocked) { - if (await retry(input.store, input.options, item, { code: "partition-key-blocked", message: "an earlier outbox record for this partition key failed", valuesPrinted: false })) retryCount++; - else staleCount++; - continue; - } - try { - await (input.publish ?? publishAgentRunKafkaMessage)(input.options.config, { - topic: item.topic, - key: item.partitionKey, - value: item.value, - headers: stringHeaders(item.headers), - }); - await input.store.completeKafkaEventOutbox(item); - deliveredCount++; - } catch (error) { - blocked = true; - const typedError = { - code: "kafka-produce-failed", - message: redactText(error instanceof Error ? error.message : String(error)).slice(0, 500), - topic: item.topic, - partitionKey: item.partitionKey, - outboxSeq: item.outboxSeq, - valuesPrinted: false, - }; - firstError ??= typedError; + const publishStartedAt = Date.now(); + const publishResults = await Promise.all(groups.map(async (group) => { + try { + await (input.publishBatch ?? publishAgentRunKafkaMessageBatch)(input.options.config, group.map((item) => ({ + topic: item.topic, + key: item.partitionKey, + value: item.value, + headers: stringHeaders(item.headers), + }))); + publishedCount += group.length; + return { group, error: null }; + } catch (error) { + return { group, error }; + } + })); + const publishDurationMs = Date.now() - publishStartedAt; + const settleStartedAt = Date.now(); + await Promise.all(publishResults.map(async ({ group, error }) => { + if (error) { + const typedError = batchError("kafka-produce-failed", error, group); + firstError ??= typedError; + for (const item of group) { const retried = await retry(input.store, input.options, item, typedError); if (retried) retryCount++; else staleCount++; } + return; + } + for (let index = 0; index < group.length; index++) { + const item = group[index]!; + try { + await input.store.completeKafkaEventOutbox(item); + deliveredCount++; + } catch (settleError) { + if (isStaleClaim(settleError)) { + const unsettled = group.slice(index); + staleCount += unsettled.length; + firstError ??= batchError("kafka-outbox-settle-stale", settleError, unsettled); + break; + } + throw settleError; + } } })); - const durationMs = Math.max(1, Date.now() - startedAt); - lastRelayCycle = { at: new Date().toISOString(), claimedCount: items.length, partitionKeyCount: groups.length, deliveredCount, retryCount, staleCount, durationMs, claimRatePerSecond: Number((items.length * 1_000 / durationMs).toFixed(2)), publishRatePerSecond: Number((deliveredCount * 1_000 / durationMs).toFixed(2)), firstError, firstOutboxSeq: items[0]?.outboxSeq ?? null, lastOutboxSeq: items.at(-1)?.outboxSeq ?? null, valuesPrinted: false }; + const settleDurationMs = Date.now() - settleStartedAt; + const durationMs = Math.max(1, Date.now() - cycleStartedAt); + lastRelayCycle = { at: new Date().toISOString(), claimedCount: items.length, partitionKeyCount: groups.length, publishedCount, deliveredCount, retryCount, staleCount, durationMs, claimDurationMs, publishDurationMs, settleDurationMs, claimRatePerSecond: Number((items.length * 1_000 / Math.max(1, claimDurationMs)).toFixed(2)), publishRatePerSecond: Number((publishedCount * 1_000 / Math.max(1, publishDurationMs)).toFixed(2)), firstError, firstOutboxSeq: items[0]?.outboxSeq ?? null, lastOutboxSeq: items.at(-1)?.outboxSeq ?? null, valuesPrinted: false }; return { action: "kafka-event-outbox-relay", enabled: true, ...lastRelayCycle }; } @@ -109,6 +130,24 @@ async function retry(store: AgentRunStore, options: KafkaOutboxRelayOptions, ite } } +function batchError(code: string, error: unknown, group: KafkaEventOutboxRecord[]): JsonRecord { + const first = group[0]; + return { + code, + message: redactText(error instanceof Error ? error.message : String(error)).slice(0, 500), + topic: first?.topic ?? null, + partitionKey: first?.partitionKey ?? null, + firstOutboxSeq: first?.outboxSeq ?? null, + lastOutboxSeq: group.at(-1)?.outboxSeq ?? null, + batchSize: group.length, + valuesPrinted: false, + }; +} + +function isStaleClaim(error: unknown): boolean { + return typeof error === "object" && error !== null && (error as { failureKind?: unknown }).failureKind === "runner-lease-conflict"; +} + function stringHeaders(headers: JsonRecord): Record { return Object.fromEntries(Object.entries(headers).map(([key, value]) => [key, typeof value === "string" ? value : JSON.stringify(value)])); } diff --git a/src/mgr/postgres-store.ts b/src/mgr/postgres-store.ts index a6cf049..9c6cb70 100644 --- a/src/mgr/postgres-store.ts +++ b/src/mgr/postgres-store.ts @@ -825,53 +825,58 @@ CREATE TABLE IF NOT EXISTS agentrun_schema_migrations ( } async claimKafkaEventOutbox(input: { owner: string; leaseMs: number; limit: number }): Promise { - const leaseExpiresAt = new Date(Date.now() + input.leaseMs).toISOString(); const result = await this.pool.query( - `WITH heads AS MATERIALIZED ( - SELECT candidate.id, candidate.topic, candidate.partition_key + `WITH head_ids AS MATERIALIZED ( + SELECT DISTINCT ON (topic, partition_key) id + FROM agentrun_kafka_event_outbox + WHERE delivered_at IS NULL + ORDER BY topic, partition_key, outbox_seq + ), heads AS MATERIALIZED ( + SELECT candidate.topic, candidate.partition_key, candidate.outbox_seq FROM agentrun_kafka_event_outbox candidate - WHERE candidate.delivered_at IS NULL - AND candidate.next_attempt_at <= now() + JOIN head_ids ON head_ids.id = candidate.id + WHERE candidate.next_attempt_at <= now() AND (candidate.lease_expires_at IS NULL OR candidate.lease_expires_at <= now()) - AND NOT EXISTS ( - SELECT 1 FROM agentrun_kafka_event_outbox earlier - WHERE earlier.topic = candidate.topic - AND earlier.partition_key = candidate.partition_key - AND earlier.delivered_at IS NULL - AND earlier.outbox_seq < candidate.outbox_seq - ) - ORDER BY candidate.outbox_seq ASC + ORDER BY candidate.outbox_seq FOR UPDATE OF candidate SKIP LOCKED LIMIT $1 - ), ranked AS MATERIALIZED ( + ), prefix_candidates AS MATERIALIZED ( SELECT candidate.id, - row_number() OVER (PARTITION BY candidate.topic, candidate.partition_key ORDER BY candidate.outbox_seq ASC) AS key_offset - FROM agentrun_kafka_event_outbox candidate - JOIN heads ON heads.topic = candidate.topic AND heads.partition_key = candidate.partition_key - WHERE candidate.delivered_at IS NULL - AND candidate.next_attempt_at <= now() - AND (candidate.lease_expires_at IS NULL OR candidate.lease_expires_at <= now()) - AND NOT EXISTS ( - SELECT 1 FROM agentrun_kafka_event_outbox blocker - WHERE blocker.topic = candidate.topic - AND blocker.partition_key = candidate.partition_key - AND blocker.delivered_at IS NULL - AND blocker.outbox_seq < candidate.outbox_seq - AND (blocker.next_attempt_at > now() OR blocker.lease_expires_at > now()) - ) + candidate.topic, + candidate.partition_key, + candidate.outbox_seq, + candidate.next_attempt_at, + candidate.lease_expires_at + FROM heads + CROSS JOIN LATERAL ( + SELECT item.* + FROM agentrun_kafka_event_outbox item + WHERE item.topic = heads.topic + AND item.partition_key = heads.partition_key + AND item.delivered_at IS NULL + ORDER BY item.outbox_seq + LIMIT $1 + ) candidate + ), ranked AS MATERIALIZED ( + SELECT prefix_candidates.id, + prefix_candidates.outbox_seq, + row_number() OVER (PARTITION BY prefix_candidates.topic, prefix_candidates.partition_key ORDER BY prefix_candidates.outbox_seq) AS key_offset, + bool_and(prefix_candidates.next_attempt_at <= now() AND (prefix_candidates.lease_expires_at IS NULL OR prefix_candidates.lease_expires_at <= now())) OVER (PARTITION BY prefix_candidates.topic, prefix_candidates.partition_key ORDER BY prefix_candidates.outbox_seq) AS prefix_due + FROM prefix_candidates ), due AS MATERIALIZED ( SELECT candidate.id FROM agentrun_kafka_event_outbox candidate JOIN ranked ON ranked.id = candidate.id + WHERE ranked.prefix_due ORDER BY ranked.key_offset ASC, candidate.outbox_seq ASC FOR UPDATE OF candidate SKIP LOCKED LIMIT $1 ) UPDATE agentrun_kafka_event_outbox o - SET attempt_count = o.attempt_count + 1, lease_owner = $2, lease_expires_at = $3, updated_at = now() + SET attempt_count = o.attempt_count + 1, lease_owner = $2, lease_expires_at = clock_timestamp() + ($3 * interval '1 millisecond'), updated_at = clock_timestamp() FROM due WHERE o.id = due.id RETURNING o.*`, - [clamp(input.limit, 1, 500), input.owner, leaseExpiresAt], + [clamp(input.limit, 1, 500), input.owner, input.leaseMs], ); return result.rows.map(kafkaEventOutboxFromRow); } diff --git a/src/selftest/cases/35-kafka-durable-outbox.ts b/src/selftest/cases/35-kafka-durable-outbox.ts index a117e0e..6bf42d2 100644 --- a/src/selftest/cases/35-kafka-durable-outbox.ts +++ b/src/selftest/cases/35-kafka-durable-outbox.ts @@ -1,5 +1,5 @@ import assert from "node:assert/strict"; -import { closeAgentRunKafkaProducer, kafkaMessageMetadata, publishAgentRunKafkaMessage, rawFrameValue, setAgentRunKafkaProducerFactoryForSelfTest, type AgentRunKafkaConfig } from "../../common/kafka-events.js"; +import { closeAgentRunKafkaProducer, kafkaMessageMetadata, publishAgentRunKafkaMessage, publishAgentRunKafkaMessageBatch, rawFrameValue, setAgentRunKafkaProducerFactoryForSelfTest, type AgentRunKafkaConfig } from "../../common/kafka-events.js"; import { AgentRunError } from "../../common/errors.js"; import { agentRunOtelTraceContext } from "../../common/otel-trace.js"; import type { CreateRunInput, JsonRecord, KafkaEventOutboxRecord } from "../../common/types.js"; @@ -99,14 +99,16 @@ const selfTest: SelfTestCase = async () => { await delay(25); const published: JsonRecord[] = []; - const publish = async (_config: AgentRunKafkaConfig, message: { value: JsonRecord }): Promise => { published.push(message.value); }; - const firstCycle = await relayKafkaEventOutboxOnce({ store, options: relayOptions, publish }); - const secondCycle = await relayKafkaEventOutboxOnce({ store, options: relayOptions, publish }); - const thirdCycle = await relayKafkaEventOutboxOnce({ store, options: relayOptions, publish }); + let batchCalls = 0; + const publishBatch = async (_config: AgentRunKafkaConfig, messages: Array<{ value: JsonRecord }>): Promise => { batchCalls++; published.push(...messages.map((message) => message.value)); }; + const firstCycle = await relayKafkaEventOutboxOnce({ store, options: relayOptions, publishBatch }); + const secondCycle = await relayKafkaEventOutboxOnce({ store, options: relayOptions, publishBatch }); + const thirdCycle = await relayKafkaEventOutboxOnce({ store, options: relayOptions, publishBatch }); assert.equal(firstCycle.deliveredCount, 3); assert.equal(firstCycle.partitionKeyCount, 1); assert.equal(secondCycle.deliveredCount, 0); assert.equal(thirdCycle.deliveredCount, 0); + assert.equal(batchCalls, 1); assert.deepEqual(published.map((item) => item.outboxSeq), [1, 2, 3]); assert.deepEqual(published.map((item) => item.sessionId), ["ses_kafka_durable", "ses_kafka_durable", "ses_kafka_durable"]); assert.equal(published[2]?.schema, "agentrun.event.v1"); @@ -135,6 +137,10 @@ const selfTest: SelfTestCase = async () => { assert.equal(stdioMetadata.commandId, command.id); await assertRelayParallelismAndFailureIsolation(); await assertRelayParallelismBoundedByBatchSize(); + await assertNativeBatchProducerOrder(); + await assertDeliveredOnlyAfterBatchAck(); + await assertClaimDurationAndMonotonicDrain(); + await assertSettleStaleIsNotProduceFailure(); assertWarmRunCommandTraceAuthority(); const claimedIntent = store.claimRunnerDispatchIntents({ owner: "terminalizer", leaseMs: 60_000, limit: 1 })[0]; @@ -174,7 +180,7 @@ const selfTest: SelfTestCase = async () => { await assertStaleClaimsCannotOverwrite(); await assertProducerRecreatedAfterDisconnected(); - return { name: "kafka-durable-outbox", tests: ["explicit-enable", "canonical-event", "warm-run-command-trace-authority", "run-level-trace-fallback", "stdio-command-trace-context", "kafka-tail-correlation-summary", "durable-key-batch", "cross-key-parallelism", "batch-bounded-parallelism", "failed-key-isolation", "atomic-dispatch-terminal", "terminal-command-idempotent-replay", "terminal-command-replay-no-new-events-or-outbox", "terminal-command-same-key-payload-conflict", "terminal-command-different-key-fails-closed", "stale-claim-fencing", "producer-reconnect"] }; + return { name: "kafka-durable-outbox", tests: ["explicit-enable", "canonical-event", "warm-run-command-trace-authority", "run-level-trace-fallback", "stdio-command-trace-context", "kafka-tail-correlation-summary", "durable-key-batch", "native-broker-batch", "ack-before-delivered", "claim-duration-visible", "monotonic-backlog-drain", "settle-stale-classification", "cross-key-parallelism", "batch-bounded-parallelism", "failed-key-isolation", "atomic-dispatch-terminal", "terminal-command-idempotent-replay", "terminal-command-replay-no-new-events-or-outbox", "terminal-command-same-key-payload-conflict", "terminal-command-different-key-fails-closed", "stale-claim-fencing", "producer-reconnect"] }; }; function assertWarmRunCommandTraceAuthority(): void { @@ -340,13 +346,13 @@ async function assertRelayParallelismAndFailureIsolation(): Promise { const result = await relayKafkaEventOutboxOnce({ store, options: { ...relayOptions, batchSize: 6 }, - publish: async (_config, message) => { + publishBatch: async (_config, messages) => { active++; maxActive = Math.max(maxActive, active); try { await delay(10); - const key = String(message.value.sessionId); - publishedByKey.set(key, [...(publishedByKey.get(key) ?? []), Number(message.value.outboxSeq)]); + const key = String(messages[0]?.value.sessionId); + publishedByKey.set(key, messages.map((message) => Number(message.value.outboxSeq))); if (key === "ses_kafka_blocked") throw new Error("selftest broker rejection"); } finally { active--; @@ -359,7 +365,7 @@ async function assertRelayParallelismAndFailureIsolation(): Promise { assert.equal(result.retryCount, 3); assert.equal((result.firstError as JsonRecord).code, "kafka-produce-failed"); assert.equal(maxActive, 2); - assert.deepEqual(publishedByKey.get("ses_kafka_blocked"), [1]); + assert.deepEqual(publishedByKey.get("ses_kafka_blocked"), [1, 2, 3]); assert.deepEqual(publishedByKey.get("ses_kafka_healthy"), [4, 5, 6]); assert.equal(store.kafkaEventOutboxStatus().backlogCount, 3); } @@ -375,7 +381,7 @@ async function assertRelayParallelismBoundedByBatchSize(): Promise { const result = await relayKafkaEventOutboxOnce({ store, options: { ...relayOptions, batchSize: 8 }, - publish: async () => { + publishBatch: async () => { active++; maxActive = Math.max(maxActive, active); await delay(10); @@ -388,6 +394,89 @@ async function assertRelayParallelismBoundedByBatchSize(): Promise { assert.equal(store.kafkaEventOutboxStatus().backlogCount, 12); } +async function assertNativeBatchProducerOrder(): Promise { + const sends: JsonRecord[] = []; + await setAgentRunKafkaProducerFactoryForSelfTest(() => ({ + connect: async () => undefined, + send: async (payload: JsonRecord) => { sends.push(payload); return []; }, + disconnect: async () => undefined, + }) as never); + try { + await publishAgentRunKafkaMessageBatch(config, [1, 2, 3].map((outboxSeq) => ({ + topic: config.agentrunEventTopic, + key: "ses_native_batch", + value: { outboxSeq }, + }))); + assert.equal(sends.length, 1); + assert.equal(sends[0]?.topic, config.agentrunEventTopic); + assert.deepEqual(((sends[0]?.messages as JsonRecord[]) ?? []).map((message) => JSON.parse(String(message.value)).outboxSeq), [1, 2, 3]); + assert.deepEqual(((sends[0]?.messages as JsonRecord[]) ?? []).map((message) => message.key), ["ses_native_batch", "ses_native_batch", "ses_native_batch"]); + } finally { + await closeAgentRunKafkaProducer(); + await setAgentRunKafkaProducerFactoryForSelfTest(null); + } +} + +async function assertDeliveredOnlyAfterBatchAck(): Promise { + const store = new MemoryAgentRunStore({ eventOutbox: { enabled: true, topic: config.agentrunEventTopic, source: config.clientId } }); + const run = store.createRun(runInput("ses_kafka_ack_gate")); + store.appendEvent(run.id, "backend_status", { phase: "ack-1" }); + store.appendEvent(run.id, "backend_status", { phase: "ack-2" }); + store.appendEvent(run.id, "backend_status", { phase: "ack-3" }); + let releaseAck: (() => void) | null = null; + const ack = new Promise((resolve) => { releaseAck = resolve; }); + let batchObserved = false; + const relay = relayKafkaEventOutboxOnce({ + store, + options: relayOptions, + publishBatch: async (_config, messages) => { + assert.deepEqual(messages.map((message) => message.value.outboxSeq), [1, 2, 3]); + batchObserved = true; + await ack; + }, + }); + while (!batchObserved) await delay(1); + assert.equal(store.kafkaEventOutboxStatus().backlogCount, 3); + releaseAck?.(); + const result = await relay; + assert.equal(result.deliveredCount, 3); + assert.equal(store.kafkaEventOutboxStatus().backlogCount, 0); +} + +async function assertClaimDurationAndMonotonicDrain(): Promise { + const store = new MemoryAgentRunStore({ eventOutbox: { enabled: true, topic: config.agentrunEventTopic, source: config.clientId } }); + const run = store.createRun(runInput("ses_kafka_monotonic_drain")); + for (let index = 1; index <= 6; index++) store.appendEvent(run.id, "backend_status", { phase: `drain-${index}` }); + const claim = store.claimKafkaEventOutbox.bind(store); + store.claimKafkaEventOutbox = async (input) => { + await delay(10); + return claim(input); + }; + const backlogCounts = [Number(store.kafkaEventOutboxStatus().backlogCount)]; + const cycleOptions = { ...relayOptions, batchSize: 2 }; + while (backlogCounts.at(-1)! > 0) { + const result = await relayKafkaEventOutboxOnce({ store, options: cycleOptions, publishBatch: async () => undefined }); + assert.ok(Number(result.claimDurationMs) >= 8); + backlogCounts.push(Number(store.kafkaEventOutboxStatus().backlogCount)); + } + assert.deepEqual(backlogCounts, [6, 4, 2, 0]); + assert.equal(backlogCounts.every((count, index) => index === 0 || count < backlogCounts[index - 1]!), true); +} + +async function assertSettleStaleIsNotProduceFailure(): Promise { + const store = new MemoryAgentRunStore({ eventOutbox: { enabled: true, topic: config.agentrunEventTopic, source: config.clientId } }); + const run = store.createRun(runInput("ses_kafka_settle_stale")); + store.appendEvent(run.id, "backend_status", { phase: "settle-stale-1" }); + store.appendEvent(run.id, "backend_status", { phase: "settle-stale-2" }); + store.completeKafkaEventOutbox = () => { throw new AgentRunError("runner-lease-conflict", "Kafka event outbox claim is stale", { httpStatus: 409 }); }; + const result = await relayKafkaEventOutboxOnce({ store, options: relayOptions, publishBatch: async () => undefined }); + assert.equal(result.deliveredCount, 0); + assert.equal(result.publishedCount, 2); + assert.equal(result.retryCount, 0); + assert.equal(result.staleCount, 2); + assert.equal((result.firstError as JsonRecord).code, "kafka-outbox-settle-stale"); +} + function delay(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } diff --git a/src/selftest/integration/postgres-kafka-order.ts b/src/selftest/integration/postgres-kafka-order.ts index 634bda3..58adb49 100644 --- a/src/selftest/integration/postgres-kafka-order.ts +++ b/src/selftest/integration/postgres-kafka-order.ts @@ -210,34 +210,60 @@ async function assertPostgresBatchRelay(store: Awaited { + publishBatch: async (_config, messages) => { active++; maxActive = Math.max(maxActive, active); try { await delay(1); - const key = String(message.value.sessionId); - publishedByKey.set(key, [...(publishedByKey.get(key) ?? []), Number(message.value.outboxSeq)]); + const key = String(messages[0]?.value.sessionId); + publishedByKey.set(key, messages.map((message) => Number(message.value.outboxSeq))); if (key === "ses_pg_kafka_blocked") throw new Error("postgres selftest broker rejection"); } finally { active--; } }, }); + const backlogCounts = [200, Number((await store.kafkaEventOutboxStatus()).backlogCount)]; assert.equal(result.claimedCount, 200); assert.equal(result.partitionKeyCount, 2); assert.equal(result.deliveredCount, 100); assert.equal(result.retryCount, 100); assert.equal(maxActive, 2); - assert.deepEqual(publishedByKey.get("ses_pg_kafka_blocked"), [1]); + assert.deepEqual(publishedByKey.get("ses_pg_kafka_blocked"), Array.from({ length: 100 }, (_, index) => index * 2 + 1)); assert.deepEqual(publishedByKey.get("ses_pg_kafka_healthy"), Array.from({ length: 100 }, (_, index) => (index + 1) * 2)); const status = await store.kafkaEventOutboxStatus(); assert.equal(status.backlogCount, 100); assert.equal((status.firstError as JsonRecord).code, "kafka-produce-failed"); await delay(5); - const retryResult = await relayKafkaEventOutboxOnce({ store, options, publish: async () => undefined }); + const retryResult = await relayKafkaEventOutboxOnce({ store, options, publishBatch: async () => undefined }); assert.equal(retryResult.claimedCount, 100); assert.equal(retryResult.deliveredCount, 100); - assert.equal((await store.kafkaEventOutboxStatus()).backlogCount, 0); + backlogCounts.push(Number((await store.kafkaEventOutboxStatus()).backlogCount)); + assert.deepEqual(backlogCounts, [200, 100, 0]); + assert.equal(backlogCounts.every((count, index) => index === 0 || count <= backlogCounts[index - 1]!), true); + await assertLeaseStartsAtDatabaseClaim(store, pool); +} + +async function assertLeaseStartsAtDatabaseClaim(store: Awaited>, pool: Pool): Promise { + const run = await store.createRun(runInput("ses_pg_kafka_lease_clock")); + await store.appendEvent(run.id, "backend_status", { phase: "lease-clock" }); + const row = await pool.query<{ id: string }>("SELECT id FROM agentrun_kafka_event_outbox WHERE run_id = $1 ORDER BY outbox_seq LIMIT 1", [run.id]); + const lock = await pool.connect(); + try { + await lock.query("BEGIN"); + await lock.query("SELECT id FROM agentrun_kafka_event_outbox WHERE id = $1 FOR UPDATE", [row.rows[0]?.id]); + const claimPromise = store.claimKafkaEventOutbox({ owner: "lease-clock", leaseMs: 50, limit: 1 }); + await delay(100); + await lock.query("COMMIT"); + const claim = singleClaim(await claimPromise); + assert.ok(Date.parse(claim.leaseExpiresAt ?? "") - Date.now() > 20, "lease expired while claim query was blocked before the database update"); + await store.completeKafkaEventOutbox(claim); + } catch (error) { + await lock.query("ROLLBACK").catch(() => undefined); + throw error; + } finally { + lock.release(); + } } async function assertForwardMigrationFromExistingV02Ledger(store: Awaited>, pool: Pool): Promise>> {