1import { assert, describe, expect, it } from "@effect/vitest"
2import { Cause, Context, Effect, Exit, Fiber, Latch, Layer, Option, Schema } from "effect"
3import {
4 ClusterError,
5 ClusterSchema,
6 EntityAddress,
7 EntityId,
8 EntityType,
9 Envelope,
10 Message,
11 MessageStorage,
12 Reply,
13 ShardId,
14 ShardingConfig,
15 Snowflake
16} from "effect/cluster"
17import { Headers } from "effect/http"
18import { Rpc, RpcSchema } from "effect/rpc"
19import { TestClock } from "effect/testing"
20
21const MemoryLayer = MessageStorage.layerMemory.pipe(
22 Layer.provideMerge(Snowflake.layerGenerator),
23 Layer.provide(ShardingConfig.layerDefaults)
24)
25
26describe("MessageStorage", () => {
27 describe("memory", () => {
28 it.effect("removes the primary-key index when clearing an address", () =>
29 Effect.gen(function*() {
30 const driver = yield* MessageStorage.MemoryDriver
31 const address = EntityAddress.make({
32 shardId: ShardId.make("default", 1),
33 entityType: EntityType.make("Repro"),
34 entityId: EntityId.make("one")
35 })
36 const envelope: Envelope.PartialRequestEncoded = {
37 _tag: "Request",
38 requestId: "1",
39 address: { shardId: { group: "default", id: 1 }, entityType: "Repro", entityId: "one" },
40 tag: "Repro",
41 payload: {},
42 headers: {}
43 }
44 yield* driver.encoded.saveEnvelope({ envelope, primaryKey: "dedup-key", deliverAt: null })
45 yield* driver.encoded.clearAddress(address)
46 const result = yield* driver.encoded.saveEnvelope({
47 envelope: { ...envelope, requestId: "2" },
48 primaryKey: "dedup-key",
49 deliverAt: null
50 })
51 expect(result._tag).toEqual("Success")
52 }).pipe(Effect.provide(MessageStorage.MemoryDriver.layer)))
53
54 it.effect("removes a queued Interrupt when clearing an address", () =>
55 Effect.gen(function*() {
56 const storage = yield* MessageStorage.MessageStorage
57 const snowflake = yield* Snowflake.Generator
58 const request = yield* makeRequest()
59 yield* storage.saveRequest(request)
60 yield* storage.saveEnvelope(Message.OutgoingEnvelope.interrupt({
61 id: snowflake.nextUnsafe(),
62 requestId: request.envelope.requestId,
63 address: request.envelope.address
64 }))
65
66 yield* storage.clearAddress(request.envelope.address)
67
68 const messages = yield* storage.unprocessedMessages([request.envelope.address.shardId])
69 expect(messages.map(({ envelope }) => envelope._tag)).toEqual([])
70 }).pipe(Effect.provide(MemoryLayer)))
71
72 it.effect("encoded unprocessedMessages fails closed for an empty address filter", () =>
73 Effect.gen(function*() {
74 const driver = yield* MessageStorage.MemoryDriver
75 const envelope: Envelope.PartialRequestEncoded = {
76 _tag: "Request",
77 requestId: "1",
78 address: { shardId: { group: "default", id: 1 }, entityType: "Repro", entityId: "one" },
79 tag: "Repro",
80 payload: {},
81 headers: {}
82 }
83 yield* driver.encoded.saveEnvelope({ envelope, primaryKey: null, deliverAt: null })
84 const messages = yield* driver.encoded.unprocessedMessages(["default:1"], 0, { addresses: [] })
85 expect(messages).toHaveLength(0)
86 }).pipe(Effect.provide(MessageStorage.MemoryDriver.layer)))
87
88 it.effect("encoded resetAddresses fails closed for an empty address list", () =>
89 Effect.gen(function*() {
90 const driver = yield* MessageStorage.MemoryDriver
91 const envelope: Envelope.PartialRequestEncoded = {
92 _tag: "Request",
93 requestId: "1",
94 address: { shardId: { group: "default", id: 1 }, entityType: "Repro", entityId: "one" },
95 tag: "Repro",
96 payload: {},
97 headers: {}
98 }
99 yield* driver.encoded.saveEnvelope({ envelope, primaryKey: null, deliverAt: null })
100 const claimed = yield* driver.encoded.unprocessedMessages(["default:1"], 0)
101 expect(claimed).toHaveLength(1)
102 yield* driver.encoded.resetAddresses([])
103 const messages = yield* driver.encoded.unprocessedMessages(["default:1"], 1)
104 expect(messages).toHaveLength(0)
105 }).pipe(Effect.provide(MessageStorage.MemoryDriver.layer)))
106
107 it.effect("resetRequests with no IDs leaves existing claims untouched", () =>
108 Effect.gen(function*() {
109 const storage = yield* MessageStorage.MessageStorage
110 const request = yield* makeRequest()
111 yield* storage.saveRequest(request)
112 expect(yield* storage.unprocessedMessages([request.envelope.address.shardId])).toHaveLength(1)
113 yield* storage.resetRequests([])
114 const driver = yield* MessageStorage.MemoryDriver
115 yield* driver.encoded.resetRequests([])
116 expect(yield* storage.unprocessedMessages([request.envelope.address.shardId])).toHaveLength(0)
117 }).pipe(Effect.provide(MemoryLayer)))
118
119 it.effect("resetRequests releases only the selected request at a shared address", () =>
120 Effect.gen(function*() {
121 const storage = yield* MessageStorage.MessageStorage
122 const selected = yield* makeRequest()
123 const unrelated = yield* makeRequest()
124 yield* storage.saveRequest(selected)
125 yield* storage.saveRequest(unrelated)
126 const shards = [selected.envelope.address.shardId]
127 expect(yield* storage.unprocessedMessages(shards)).toHaveLength(2)
128 yield* storage.resetRequests([selected.envelope.requestId])
129 const messages = yield* storage.unprocessedMessages(shards)
130 expect(messages.map((message) => message.envelope.requestId)).toEqual([selected.envelope.requestId])
131 expect(yield* storage.unprocessedMessages(shards)).toHaveLength(0)
132 }).pipe(Effect.provide(MemoryLayer)))
133
134 it.effect("resetRequests preserves chunk replies, exit replies, and completed state", () =>
135 Effect.gen(function*() {
136 const storage = yield* MessageStorage.MessageStorage
137 const streaming = yield* makeRequest({ rpc: StreamRpc, payload: StreamRpc.payloadSchema.make({ id: 123 }) })
138 const completed = yield* makeRequest()
139 yield* storage.saveRequest(streaming)
140 yield* storage.saveRequest(completed)
141 yield* storage.saveReply(yield* makeChunkReply(streaming))
142 yield* storage.saveReply(yield* makeReply(completed))
143 const replies = yield* storage.repliesFor([streaming, completed])
144 expect(replies).toHaveLength(2)
145 const shards = [streaming.envelope.address.shardId]
146 expect(yield* storage.unprocessedMessages(shards)).toHaveLength(1)
147 yield* storage.resetRequests([streaming.envelope.requestId, completed.envelope.requestId])
148 const messages = yield* storage.unprocessedMessages(shards)
149 expect(messages.map((message) => message.envelope.requestId)).toEqual([streaming.envelope.requestId])
150 expect(yield* storage.repliesFor([streaming, completed])).toEqual(replies)
151 }).pipe(Effect.provide(MemoryLayer)))
152
153 it.effect("saves a request", () =>
154 Effect.gen(function*() {
155 const storage = yield* MessageStorage.MessageStorage
156 const request = yield* makeRequest()
157 const result = yield* storage.saveRequest(request)
158 expect(result._tag).toEqual("Success")
159 const messages = yield* storage.unprocessedMessages([request.envelope.address.shardId])
160 expect(messages).toHaveLength(1)
161 }).pipe(Effect.provide(MemoryLayer)))
162
163 it.effect("detects duplicates", () =>
164 Effect.gen(function*() {
165 const storage = yield* MessageStorage.MessageStorage
166 yield* storage.saveRequest(
167 yield* makeRequest({
168 rpc: PrimaryKeyTest,
169 payload: PrimaryKeyTest.payloadSchema.make({ id: 123 })
170 })
171 )
172 const result = yield* storage.saveRequest(
173 yield* makeRequest({
174 rpc: PrimaryKeyTest,
175 payload: PrimaryKeyTest.payloadSchema.make({ id: 123 })
176 })
177 )
178 expect(result._tag).toEqual("Duplicate")
179 }).pipe(Effect.provide(MemoryLayer)))
180
181 it.effect("unprocessedMessages excludes complete requests", () =>
182 Effect.gen(function*() {
183 const storage = yield* MessageStorage.MessageStorage
184 const request = yield* makeRequest()
185 yield* storage.saveRequest(request)
186 yield* storage.saveReply(yield* makeReply(request))
187 const messages = yield* storage.unprocessedMessages([request.envelope.address.shardId])
188 expect(messages).toHaveLength(0)
189 }).pipe(Effect.provide(MemoryLayer)))
190 it.effect("unprocessedMessages honors the limit option", () =>
191 Effect.gen(function*() {
192 const storage = yield* MessageStorage.MessageStorage
193 for (let i = 1; i <= 5; i++) {
194 yield* storage.saveRequest(yield* makeRequest({ payload: { id: i }, entityId: String(i) }))
195 }
196 const shardId = ShardId.make("default", 1)
197 const messages = yield* storage.unprocessedMessages([shardId], { limit: 3 })
198 expect(messages).toHaveLength(3)
199 expect(messages.map((m: any) => m.envelope.payload.id)).toEqual([1, 2, 3])
200 const remaining = yield* storage.unprocessedMessages([shardId])
201 expect(remaining.map((m: any) => m.envelope.payload.id)).toEqual([4, 5])
202 }).pipe(Effect.provide(MemoryLayer)))
203
204 it.effect("unprocessedMessages filters by address", () =>
205 Effect.gen(function*() {
206 const storage = yield* MessageStorage.MessageStorage
207 for (let i = 1; i <= 4; i++) {
208 yield* storage.saveRequest(yield* makeRequest({ payload: { id: i }, entityId: String(i) }))
209 }
210 const shardId = ShardId.make("default", 1)
211 const address = (entityId: string) =>
212 EntityAddress.make({
213 shardId,
214 entityType: EntityType.make("test"),
215 entityId: EntityId.make(entityId)
216 })
217 const messages = yield* storage.unprocessedMessages([shardId], {
218 addresses: [address("2"), address("4")]
219 })
220 expect(messages.map((m: any) => m.envelope.payload.id)).toEqual([2, 4])
221 // an empty address filter returns nothing
222 const none = yield* storage.unprocessedMessages([shardId], { addresses: [] })
223 expect(none).toHaveLength(0)
224 }).pipe(Effect.provide(MemoryLayer)))
225
226 it.effect("resetAddresses makes claimed messages eligible again", () =>
227 Effect.gen(function*() {
228 const storage = yield* MessageStorage.MessageStorage
229 for (let i = 1; i <= 4; i++) {
230 yield* storage.saveRequest(yield* makeRequest({ payload: { id: i }, entityId: String(i) }))
231 }
232 const shardId = ShardId.make("default", 1)
233 yield* storage.unprocessedMessages([shardId], { limit: 3 })
234 yield* storage.resetAddresses(["1", "3"].map((entityId) =>
235 EntityAddress.make({
236 shardId,
237 entityType: EntityType.make("test"),
238 entityId: EntityId.make(entityId)
239 })
240 ))
241
242 const messages = yield* storage.unprocessedMessages([shardId])
243 expect(messages.map((m: any) => m.envelope.payload.id)).toEqual([1, 3, 4])
244 }).pipe(Effect.provide(MemoryLayer)))
245
246 it.effect("repliesFor", () =>
247 Effect.gen(function*() {
248 const storage = yield* MessageStorage.MessageStorage
249 const request = yield* makeRequest()
250 yield* storage.saveRequest(request)
251 let replies = yield* storage.repliesFor([request])
252 expect(replies).toHaveLength(0)
253 yield* storage.saveReply(yield* makeReply(request))
254 replies = yield* storage.repliesFor([request])
255 expect(replies).toHaveLength(1)
256 expect(replies[0].requestId).toEqual(request.envelope.requestId)
257 }).pipe(Effect.provide(MemoryLayer)))
258
259 it.effect("registerReplyHandler", () =>
260 Effect.gen(function*() {
261 const storage = yield* MessageStorage.MessageStorage
262 const latch = yield* Latch.make()
263 const request = yield* makeRequest()
264 yield* storage.saveRequest(request)
265 const fiber = yield* storage.registerReplyHandler(
266 new Message.OutgoingRequest({
267 ...request,
268 respond: () => latch.open
269 })
270 ).pipe(Effect.forkChild)
271 yield* TestClock.adjust(1)
272 yield* storage.saveReply(yield* makeReply(request))
273 yield* latch.await
274 yield* Fiber.await(fiber)
275 }).pipe(Effect.provide(MemoryLayer)))
276
277 it.effect("unregisterShardReplyHandlers fails parked waiters during rebalance", () =>
278 Effect.gen(function*() {
279 const storage = yield* MessageStorage.MessageStorage
280 const request = yield* makeRequest()
281 const fiber = yield* storage.registerReplyHandler(
282 new Message.OutgoingRequest({
283 ...request,
284 respond: () => Effect.void
285 })
286 ).pipe(Effect.forkChild({ startImmediately: true }))
287 yield* TestClock.adjust(1)
288
289 yield* storage.unregisterShardReplyHandlers(request.envelope.address.shardId)
290 const exit = yield* Fiber.await(fiber)
291
292 assert(Exit.isFailure(exit))
293 const error = Cause.findErrorOption(exit.cause)
294 assert(Option.isSome(error))
295 assert(error.value instanceof ClusterError.EntityNotAssignedToRunner)
296 }).pipe(Effect.provide(MemoryLayer)))
297
298 it.effect("unregisterShardReplyHandlers annotates parked waiters during shutdown", () =>
299 Effect.gen(function*() {
300 const storage = yield* MessageStorage.MessageStorage
301 const request = yield* makeRequest()
302 const fiber = yield* storage.registerReplyHandler(
303 new Message.OutgoingRequest({
304 ...request,
305 respond: () => Effect.void
306 })
307 ).pipe(Effect.forkChild({ startImmediately: true }))
308 yield* TestClock.adjust(1)
309
310 yield* storage.unregisterShardReplyHandlers(request.envelope.address.shardId, { interrupt: true })
311 const exit = yield* Fiber.await(fiber)
312
313 assert(Exit.isFailure(exit))
314 assert(Cause.hasInterruptsOnly(exit.cause))
315 assert(exit.cause.reasons.some((reason) =>
316 reason._tag === "Interrupt" && reason.annotations.has(ClusterSchema.Abandon.key)
317 ))
318 }).pipe(Effect.provide(MemoryLayer)))
319
320 it.effect("reply handlers receive the persisted defect fallback for unencodable replies", () =>
321 Effect.gen(function*() {
322 const storage = yield* MessageStorage.MessageStorage
323 const snowflake = yield* Snowflake.Generator
324 const latch = yield* Latch.make()
325 const request = yield* makeRequest({ rpc: UnknownSuccessRpc })
326 yield* storage.saveRequest(request)
327 let received: Reply.Reply<any> | undefined
328 const fiber = yield* storage.registerReplyHandler(
329 new Message.OutgoingRequest({
330 ...request,
331 respond: (reply) => {
332 received = reply
333 return latch.open
334 }
335 })
336 ).pipe(Effect.forkChild)
337 yield* TestClock.adjust(1)
338 yield* storage.saveReply(
339 new Reply.ReplyWithContext({
340 reply: new Reply.WithExit({
341 id: snowflake.nextUnsafe(),
342 requestId: request.envelope.requestId,
343 exit: Exit.succeed(new Error("not json encodable")) as any
344 }),
345 context: request.context,
346 rpc: request.rpc
347 })
348 )
349 yield* latch.await
350 yield* Fiber.await(fiber)
351 expect(received?._tag).toEqual("WithExit")
352 const exit = (received as Reply.WithExit<any>).exit
353 expect(Exit.isFailure(exit)).toBe(true)
354 expect(JSON.stringify(exit)).toContain("MalformedMessage")
355 const stored = yield* storage.repliesForUnfiltered([request.envelope.requestId])
356 expect(JSON.stringify(stored)).toContain("MalformedMessage")
357 }).pipe(Effect.provide(MemoryLayer)))
358 })
359})
360
361export const GetUserRpc = Rpc.make("GetUser", {
362 payload: { id: Schema.Number }
363})
364
365const UnknownSuccessRpc = Rpc.make("UnknownSuccess", {
366 payload: { id: Schema.Number },
367 success: Schema.Unknown
368})
369
370export const makeRequest = Effect.fnUntraced(function*(options?: {
371 readonly rpc?: Rpc.AnyWithProps
372 readonly payload?: any
373 readonly entityId?: string
374}) {
375 const snowflake = yield* Snowflake.Generator
376 const rpc = options?.rpc ?? GetUserRpc
377 return new Message.OutgoingRequest({
378 envelope: Envelope.makeRequest<any>({
379 requestId: snowflake.nextUnsafe(),
380 address: EntityAddress.make({
381 shardId: ShardId.make("default", 1),
382 entityType: EntityType.make("test"),
383 entityId: EntityId.make(options?.entityId ?? "1")
384 }),
385 tag: rpc._tag,
386 payload: options?.payload ?? { id: 123 },
387 traceId: "noop",
388 spanId: "noop",
389 sampled: false,
390 headers: Headers.empty
391 }),
392 annotations: rpc.annotations,
393 context: Context.empty() as any,
394 rpc,
395 lastReceivedReply: Option.none(),
396 respond() {
397 return Effect.void
398 }
399 })
400})
401
402export class PrimaryKeyTest extends Rpc.make("PrimaryKeyTest", {
403 payload: {
404 id: Schema.Number
405 },
406 primaryKey: (value) => value.id.toString()
407}) {}
408
409export class StreamRpc extends Rpc.make("StreamTest", {
410 success: RpcSchema.Stream(Schema.Void, Schema.Never),
411 payload: {
412 id: Schema.Number
413 },
414 primaryKey: (value) => value.id.toString()
415}) {}
416
417export const makeReply = Effect.fnUntraced(function*(request: Message.OutgoingRequest<any>) {
418 const snowflake = yield* Snowflake.Generator
419 return new Reply.ReplyWithContext({
420 reply: new Reply.WithExit({
421 id: snowflake.nextUnsafe(),
422 requestId: request.envelope.requestId,
423 exit: Exit.void as any
424 }),
425 context: request.context,
426 rpc: request.rpc
427 })
428})
429
430export const makeAckChunk = Effect.fnUntraced(function*(
431 request: Message.OutgoingRequest<any>,
432 chunk: Reply.ReplyWithContext<any>
433) {
434 const snowflake = yield* Snowflake.Generator
435 return new Message.OutgoingEnvelope({
436 envelope: new Envelope.AckChunk({
437 id: snowflake.nextUnsafe(),
438 address: request.envelope.address,
439 requestId: chunk.reply.requestId,
440 replyId: chunk.reply.id
441 }),
442 rpc: request.rpc
443 })
444})
445
446export const makeChunkReply = Effect.fnUntraced(function*(request: Message.OutgoingRequest<any>, sequence = 0) {
447 const snowflake = yield* Snowflake.Generator
448 return new Reply.ReplyWithContext({
449 reply: new Reply.Chunk({
450 id: snowflake.nextUnsafe(),
451 requestId: request.envelope.requestId,
452 sequence,
453 values: [undefined]
454 }),
455 context: request.context,
456 rpc: request.rpc
457 })
458})
459
460export const makeEmptyReply = (request: Message.OutgoingRequest<any>) => {
461 return new Reply.ReplyWithContext({
462 reply: Reply.Chunk.emptyFrom(request.envelope.requestId),
463 context: request.context,
464 rpc: request.rpc
465 })
466}