diff --git a/.changeset/cosmos-cluster-reply-uniqueness.md b/.changeset/cosmos-cluster-reply-uniqueness.md new file mode 100644 index 000000000..6d90a3d20 --- /dev/null +++ b/.changeset/cosmos-cluster-reply-uniqueness.md @@ -0,0 +1,5 @@ +--- +"@effect-app/infra": patch +--- + +Harden Cosmos cluster reply persistence against duplicate terminal replies and duplicate chunk sequences. diff --git a/packages/infra/src/ClusterCosmos.ts b/packages/infra/src/ClusterCosmos.ts index e4701e61f..f450a07cd 100644 --- a/packages/infra/src/ClusterCosmos.ts +++ b/packages/infra/src/ClusterCosmos.ts @@ -1,4 +1,4 @@ -import type { OperationInput, PatchRequestBody } from "@azure/cosmos" +import type { JSONObject, JSONValue, OperationInput, OperationResponse, PatchRequestBody } from "@azure/cosmos" import * as Arr from "effect-app/Array" import * as Effect from "effect-app/Effect" import * as Layer from "effect-app/Layer" @@ -27,6 +27,7 @@ export interface ClusterCosmosConfig { type MessageKind = "Request" | "AckChunk" | "Interrupt" type CosmosQueryValue = string | number | boolean | null | Array type CosmosParameter = { readonly name: string; readonly value: CosmosQueryValue } +type JsonNonEmptyArray = readonly [JSONValue, ...Array] interface MessageDoc { readonly id: string @@ -54,8 +55,9 @@ interface MessageDoc { } type ReplyDoc = WithExitReplyDoc | ChunkReplyDoc +type ReplyItemDoc = ReplyDoc | ReplyIndexDoc -interface ReplyDocBase { +interface ReplyDocBase extends JSONObject { readonly id: string readonly _partitionKey: string readonly type: "reply" @@ -66,16 +68,26 @@ interface ReplyDocBase { interface WithExitReplyDoc extends ReplyDocBase { readonly kind: "WithExit" - readonly payload: Reply.WithExitEncoded["exit"] + readonly payload: Reply.WithExitEncoded["exit"] & JSONValue readonly sequence: null } interface ChunkReplyDoc extends ReplyDocBase { readonly kind: "Chunk" - readonly payload: Reply.ChunkEncoded["values"] + readonly payload: JsonNonEmptyArray readonly sequence: number } +interface ReplyIndexDoc extends JSONObject { + readonly id: string + readonly _partitionKey: string + readonly type: "reply-index" + readonly requestId: string + readonly replyId: string + readonly kind: "WithExit" | "ChunkSequence" + readonly sequence: number | null +} + interface RunnerDoc { readonly id: string readonly _partitionKey: "runner" @@ -114,6 +126,9 @@ const messagePartition = (shardId: string) => `message::${shardId}` const messageDocId = (envelope: Envelope.Encoded, primaryKey: string | null) => cosmosId(primaryKey === null ? envelopeId(envelope) : `primary::${primaryKey}`) const replyPartition = (requestId: string) => `reply::${requestId}` +const terminalReplyIndexDocId = (requestId: string) => cosmosId(`reply::${requestId}::exit`) +const chunkReplyIndexDocId = (requestId: string, sequence: number) => + cosmosId(`reply::${requestId}::chunk::${sequence}`) const runnerDocId = (address: string) => cosmosId(`runner::${address}`) const lockDocId = (shardId: string) => cosmosId(`lock::${shardId}`) const tenMinutes = Duration.toMillis(Duration.minutes(10)) @@ -121,6 +136,20 @@ const maxCosmosBatchOperations = 100 const isSuccessfulStatus = (statusCode: number | undefined): boolean => statusCode !== undefined && statusCode >= 200 && statusCode < 300 +class CosmosBatchOperationError extends Error { + readonly _tag = "CosmosBatchOperationError" + constructor(readonly statusCode: number | undefined) { + super(`Cosmos batch operation failed with status ${statusCode ?? "unknown"}`) + } +} + +const ensureBatchSucceeded = (results: ReadonlyArray | undefined) => { + const failed = results?.find((result) => !isSuccessfulStatus(result.statusCode)) + return failed === undefined && results !== undefined + ? Effect.void + : Effect.fail(new CosmosBatchOperationError(failed?.statusCode)) +} + const isCosmosStatus = (u: unknown, code: number): boolean => Cause.isUnknownError(u) ? isCosmosStatus(u.cause, code) @@ -130,6 +159,31 @@ const isConflict = (u: unknown) => isCosmosStatus(u, 409) const isNotFound = (u: unknown) => isCosmosStatus(u, 404) const isPreconditionFailed = (u: unknown) => isCosmosStatus(u, 412) +const isJsonRecord = (value: unknown): value is Readonly> => + typeof value === "object" && value !== null && !Array.isArray(value) && !(value instanceof Date) + +const isJsonValue = (value: unknown): value is JSONValue => + value === null + || typeof value === "string" + || typeof value === "boolean" + || value instanceof Date + || (typeof value === "number" && Number.isFinite(value)) + || (Array.isArray(value) && value.every(isJsonValue)) + || (isJsonRecord(value) && Object.values(value).every(isJsonValue)) + +const jsonValue = (value: unknown, label: string) => + isJsonValue(value) ? Effect.succeed(value) : Effect.fail(new TypeError(`${label} is not JSON-encodable`)) + +const jsonReplyExit = (value: Reply.WithExitEncoded["exit"]) => + isJsonValue(value) ? Effect.succeed(value) : Effect.fail(new TypeError("reply exit is not JSON-encodable")) + +const jsonNonEmptyArray = Effect.fnUntraced(function*(values: Reply.ChunkEncoded["values"], label: string) { + const head = yield* jsonValue(values[0], label) + const tail = yield* Effect.forEach(values.slice(1), (value) => jsonValue(value, label)) + const result: JsonNonEmptyArray = [head, ...tail] + return result +}) + const respBytes = ( resp: { diagnostics?: { clientSideRequestStatistics?: { totalResponsePayloadLengthInBytes?: number } } } ) => resp.diagnostics?.clientSideRequestStatistics?.totalResponsePayloadLengthInBytes ?? 0 @@ -298,29 +352,51 @@ const envelopeFromDoc = ( } } -const replyToDoc = (reply: Reply.Encoded): ReplyDoc => - reply._tag === "WithExit" - ? { +const replyToDoc = Effect.fnUntraced(function*(reply: Reply.Encoded) { + return reply._tag === "WithExit" + ? ({ id: cosmosId(reply.id), _partitionKey: replyPartition(reply.requestId), type: "reply", rowid: reply.id, kind: "WithExit", requestId: reply.requestId, - payload: reply.exit, + payload: yield* jsonReplyExit(reply.exit), sequence: null, acked: false - } - : { + } satisfies WithExitReplyDoc) + : ({ id: cosmosId(reply.id), _partitionKey: replyPartition(reply.requestId), type: "reply", rowid: reply.id, kind: "Chunk", requestId: reply.requestId, - payload: reply.values, + payload: yield* jsonNonEmptyArray(reply.values, "reply chunk"), sequence: reply.sequence, acked: false + } satisfies ChunkReplyDoc) +}) + +const replyToIndexDoc = (reply: Reply.Encoded): ReplyIndexDoc => + reply._tag === "WithExit" + ? { + id: terminalReplyIndexDocId(reply.requestId), + _partitionKey: replyPartition(reply.requestId), + type: "reply-index", + requestId: reply.requestId, + replyId: reply.id, + kind: "WithExit", + sequence: null + } + : { + id: chunkReplyIndexDocId(reply.requestId, reply.sequence), + _partitionKey: replyPartition(reply.requestId), + type: "reply-index", + requestId: reply.requestId, + replyId: reply.id, + kind: "ChunkSequence", + sequence: reply.sequence } const replyFromDoc = (doc: ReplyDoc): Reply.Encoded => @@ -391,6 +467,14 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: { Effect.map((resp) => resp.resources) ) + const queryReplyItems = (query: string, parameters: ReadonlyArray) => + Effect + .tryPromise(() => container.items.query({ query, parameters: Array.from(parameters) }).fetchAll()) + .pipe( + Effect.tap(annotateFeed), + Effect.map((resp) => resp.resources) + ) + const lastReply = (replyId: string | null) => replyId === null ? Effect.succeed(Option.none()) @@ -443,10 +527,7 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: { .tryPromise(() => container.items.batch(operations, partitionKey)) .pipe( Effect.tap(annotateItem), - Effect.flatMap((resp) => { - const failed = resp.result?.find((result) => !isSuccessfulStatus(result.statusCode)) - return failed === undefined ? Effect.void : Effect.fail(failed.statusCode) - }), + Effect.flatMap((resp) => ensureBatchSucceeded(resp.result)), Effect.catchIf(() => true, () => Effect.forEach(chunk, fallback, { discard: true })) ) }, @@ -500,6 +581,27 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: { deleteDoc ) + const createReply = Effect.fnUntraced(function*(reply: Reply.Encoded) { + const doc = yield* replyToDoc(reply) + const indexDoc = replyToIndexDoc(reply) + const operations: Array = [ + { + operationType: "Create" as const, + resourceBody: doc + }, + { + operationType: "Create" as const, + resourceBody: indexDoc + } + ] + return yield* Effect + .tryPromise(() => container.items.batch(operations, doc._partitionKey)) + .pipe( + Effect.tap(annotateItem), + Effect.flatMap((resp) => ensureBatchSucceeded(resp.result)) + ) + }) + return yield* MessageStorage.makeEncoded({ saveEnvelope: ({ deliverAt, envelope, primaryKey }) => Effect @@ -540,11 +642,7 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: { saveReply: (reply) => Effect .gen(function*() { - const doc = replyToDoc(reply) - yield* Effect.tryPromise(() => container.items.create(doc)).pipe( - Effect.tap(annotateItem), - Effect.catchIf(isConflict, () => Effect.void) - ) + yield* createReply(reply) const messages = yield* queryMessages( "SELECT * FROM c WHERE c.type = 'message' AND c.requestId = @requestId", [{ name: "@requestId", value: reply.requestId }] @@ -566,8 +664,8 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: { Effect .gen(function*() { const id = String(requestId) - const replies = yield* queryReplies( - "SELECT * FROM c WHERE c.type = 'reply' AND c.requestId = @requestId AND c.kind = 'WithExit'", + const replies = yield* queryReplyItems( + "SELECT * FROM c WHERE c.requestId = @requestId AND ((c.type = 'reply' AND c.kind = 'WithExit') OR (c.type = 'reply-index' AND c.kind = 'WithExit'))", [ { name: "@requestId", value: id } ] @@ -683,8 +781,8 @@ export const makeMessageStorage = Effect.fnUntraced(function*(options?: { Effect.flatMap((messages) => { if (!Arr.isArrayNonEmpty(messages)) return Effect.void const requestIds = Array.from(new Set(messages.map((message) => message.requestId))) - return queryReplies( - "SELECT * FROM c WHERE c.type = 'reply' AND ARRAY_CONTAINS(@requestIds, c.requestId)", + return queryReplyItems( + "SELECT * FROM c WHERE ARRAY_CONTAINS(@requestIds, c.requestId) AND (c.type = 'reply' OR c.type = 'reply-index')", [{ name: "@requestIds", value: requestIds }] ) .pipe( diff --git a/packages/infra/test/cluster-cosmos.test.ts b/packages/infra/test/cluster-cosmos.test.ts index 9a4824819..cf2feccff 100644 --- a/packages/infra/test/cluster-cosmos.test.ts +++ b/packages/infra/test/cluster-cosmos.test.ts @@ -86,6 +86,33 @@ describe.skipIf(!cosmosUrl)("ClusterCosmos MessageStorage", () => { }) .pipe(Effect.provide(layerFor()))) + it.effect("rejects duplicate terminal replies and chunk sequences", () => + Effect + .gen(function*() { + const storage = yield* MessageStorage.MessageStorage + + const terminalRequest = yield* makeStreamRequest( + `duplicate-terminal/${testRunId}`, + testShardId("message-reply-unique") + ) + assert.strictEqual((yield* storage.saveRequest(terminalRequest))._tag, "Success") + yield* storage.saveReply(yield* makeStreamReply(terminalRequest)) + + const duplicateTerminal = yield* Effect.flip(storage.saveReply(yield* makeStreamReply(terminalRequest))) + assert.strictEqual(duplicateTerminal._tag, "PersistenceError") + + const chunkRequest = yield* makeStreamRequest( + `duplicate-chunk/${testRunId}`, + testShardId("message-reply-unique") + ) + assert.strictEqual((yield* storage.saveRequest(chunkRequest))._tag, "Success") + yield* storage.saveReply(yield* makeChunkReply(chunkRequest, 0)) + + const duplicateChunk = yield* Effect.flip(storage.saveReply(yield* makeChunkReply(chunkRequest, 0))) + assert.strictEqual(duplicateChunk._tag, "PersistenceError") + }) + .pipe(Effect.provide(layerFor())), 20000) + it.effect("returns each unprocessed message to only one concurrent node poll", () => Effect .gen(function*() {