1import { describe, expect, it } from "@effect/vitest"
2import { Context, Effect, Exit, Fiber, Latch, Layer, Option, Schema } from "effect"
3import {
4 EntityAddress,
5 EntityId,
6 EntityType,
7 Envelope,
8 Message,
9 MessageStorage,
10 Reply,
11 ShardId,
12 ShardingConfig,
13 Snowflake
14} from "effect/cluster"
15import { Headers } from "effect/http"
16import { Rpc, RpcSchema } from "effect/rpc"
17import { TestClock } from "effect/testing"
18
19const MemoryLayer = MessageStorage.layerMemory.pipe(
20 Layer.provideMerge(Snowflake.layerGenerator),
21 Layer.provide(ShardingConfig.layerDefaults)
22)
23
24describe("MessageStorage", () => {
25 describe("memory", () => {
26 it.effect("saves a request", () =>
27 Effect.gen(function*() {
28 const storage = yield* MessageStorage.MessageStorage
29 const request = yield* makeRequest()
30 const result = yield* storage.saveRequest(request)
31 expect(result._tag).toEqual("Success")
32 const messages = yield* storage.unprocessedMessages([request.envelope.address.shardId])
33 expect(messages).toHaveLength(1)
34 }).pipe(Effect.provide(MemoryLayer)))
35
36 it.effect("detects duplicates", () =>
37 Effect.gen(function*() {
38 const storage = yield* MessageStorage.MessageStorage
39 yield* storage.saveRequest(
40 yield* makeRequest({
41 rpc: PrimaryKeyTest,
42 payload: PrimaryKeyTest.payloadSchema.make({ id: 123 })
43 })
44 )
45 const result = yield* storage.saveRequest(
46 yield* makeRequest({
47 rpc: PrimaryKeyTest,
48 payload: PrimaryKeyTest.payloadSchema.make({ id: 123 })
49 })
50 )
51 expect(result._tag).toEqual("Duplicate")
52 }).pipe(Effect.provide(MemoryLayer)))
53
54 it.effect("unprocessedMessages excludes complete requests", () =>
55 Effect.gen(function*() {
56 const storage = yield* MessageStorage.MessageStorage
57 const request = yield* makeRequest()
58 yield* storage.saveRequest(request)
59 yield* storage.saveReply(yield* makeReply(request))
60 const messages = yield* storage.unprocessedMessages([request.envelope.address.shardId])
61 expect(messages).toHaveLength(0)
62 }).pipe(Effect.provide(MemoryLayer)))
63
64 it.effect("repliesFor", () =>
65 Effect.gen(function*() {
66 const storage = yield* MessageStorage.MessageStorage
67 const request = yield* makeRequest()
68 yield* storage.saveRequest(request)
69 let replies = yield* storage.repliesFor([request])
70 expect(replies).toHaveLength(0)
71 yield* storage.saveReply(yield* makeReply(request))
72 replies = yield* storage.repliesFor([request])
73 expect(replies).toHaveLength(1)
74 expect(replies[0].requestId).toEqual(request.envelope.requestId)
75 }).pipe(Effect.provide(MemoryLayer)))
76
77 it.effect("registerReplyHandler", () =>
78 Effect.gen(function*() {
79 const storage = yield* MessageStorage.MessageStorage
80 const latch = yield* Latch.make()
81 const request = yield* makeRequest()
82 yield* storage.saveRequest(request)
83 const fiber = yield* storage.registerReplyHandler(
84 new Message.OutgoingRequest({
85 ...request,
86 respond: () => latch.open
87 })
88 ).pipe(Effect.forkChild)
89 yield* TestClock.adjust(1)
90 yield* storage.saveReply(yield* makeReply(request))
91 yield* latch.await
92 yield* Fiber.await(fiber)
93 }).pipe(Effect.provide(MemoryLayer)))
94 })
95})
96
97export const GetUserRpc = Rpc.make("GetUser", {
98 payload: { id: Schema.Number }
99})
100
101export const makeRequest = Effect.fnUntraced(function*(options?: {
102 readonly rpc?: Rpc.AnyWithProps
103 readonly payload?: any
104 readonly entityId?: string
105 readonly shardId?: ShardId.ShardId
106}) {
107 const snowflake = yield* Snowflake.Generator
108 const rpc = options?.rpc ?? GetUserRpc
109 return new Message.OutgoingRequest({
110 envelope: Envelope.makeRequest<any>({
111 requestId: snowflake.nextUnsafe(),
112 address: EntityAddress.make({
113 shardId: options?.shardId ?? ShardId.make("default", 1),
114 entityType: EntityType.make("test"),
115 entityId: EntityId.make(options?.entityId ?? "1")
116 }),
117 tag: rpc._tag,
118 payload: options?.payload ?? { id: 123 },
119 traceId: "noop",
120 spanId: "noop",
121 sampled: false,
122 headers: Headers.empty
123 }),
124 annotations: rpc.annotations,
125 context: Context.empty() as any,
126 rpc,
127 lastReceivedReply: Option.none(),
128 respond() {
129 return Effect.void
130 }
131 })
132})
133
134export class PrimaryKeyTest extends Rpc.make("PrimaryKeyTest", {
135 payload: {
136 id: Schema.Number
137 },
138 primaryKey: (value) => value.id.toString()
139}) {}
140
141export class LongKeyRpc extends Rpc.make("LongKeyRpc", {
142 payload: {
143 id: Schema.String
144 },
145 primaryKey: (value) => value.id
146}) {}
147
148export class StreamRpc extends Rpc.make("StreamTest", {
149 success: RpcSchema.Stream(Schema.Void, Schema.Never),
150 payload: {
151 id: Schema.Number
152 },
153 primaryKey: (value) => value.id.toString()
154}) {}
155
156export const makeReply = Effect.fnUntraced(function*(request: Message.OutgoingRequest<any>) {
157 const snowflake = yield* Snowflake.Generator
158 return new Reply.ReplyWithContext({
159 reply: new Reply.WithExit({
160 id: snowflake.nextUnsafe(),
161 requestId: request.envelope.requestId,
162 exit: Exit.void as any
163 }),
164 context: request.context,
165 rpc: request.rpc
166 })
167})
168
169export const makeAckChunk = Effect.fnUntraced(function*(
170 request: Message.OutgoingRequest<any>,
171 chunk: Reply.ReplyWithContext<any>
172) {
173 const snowflake = yield* Snowflake.Generator
174 return new Message.OutgoingEnvelope({
175 envelope: new Envelope.AckChunk({
176 id: snowflake.nextUnsafe(),
177 address: request.envelope.address,
178 requestId: chunk.reply.requestId,
179 replyId: chunk.reply.id
180 }),
181 rpc: request.rpc
182 })
183})
184
185export const makeChunkReply = Effect.fnUntraced(function*(request: Message.OutgoingRequest<any>, sequence = 0) {
186 const snowflake = yield* Snowflake.Generator
187 return new Reply.ReplyWithContext({
188 reply: new Reply.Chunk({
189 id: snowflake.nextUnsafe(),
190 requestId: request.envelope.requestId,
191 sequence,
192 values: [undefined]
193 }),
194 context: request.context,
195 rpc: request.rpc
196 })
197})
198
199export const makeEmptyReply = (request: Message.OutgoingRequest<any>) => {
200 return new Reply.ReplyWithContext({
201 reply: Reply.Chunk.emptyFrom(request.envelope.requestId),
202 context: request.context,
203 rpc: request.rpc
204 })
205}