1import { assert, it } from "@effect/vitest"
2import { Deferred, Effect, Exit, Layer, Option, PubSub, Queue, Schema, Scope, Stream } from "effect"
3import {
4 ClusterSchema,
5 Entity,
6 EntityAddress,
7 EntityId,
8 EntityType,
9 Envelope,
10 Message,
11 MessageStorage,
12 RunnerHealth,
13 Runners,
14 RunnerServer,
15 RunnerStorage,
16 Sharding,
17 ShardingConfig,
18 Snowflake
19} from "effect/cluster"
20import { Headers } from "effect/http"
21import { Rpc, RpcMessage, type RpcSerialization, RpcServer, RpcTest } from "effect/rpc"
22import { TestClock } from "effect/testing"
23
24const ReproEntity = Entity.make("ReproRunnerServer", [
25 Rpc.make("ReproStream", { success: Schema.Int, payload: { id: Schema.Number }, stream: true })
26]).annotateRpcs(ClusterSchema.Persisted, false)
27
28const HoleCodecEntity = Entity.make("HoleCodecEntity", [
29 Rpc.make("Double", { success: Schema.Int, payload: { id: Schema.Number } })
30]).annotateRpcs(ClusterSchema.Persisted, false)
31
32const jsonCodec = Schema.toCodecJson as RpcSerialization.CodecFor
33
34// A hole codec that is observably different from `Schema.toCodecJson`: the
35// entity payload and the replies become JSON strings on the wire.
36const codecForJsonString =
37 (<S extends Schema.Top>(schema: S) =>
38 Schema.fromJsonString(Schema.toCodecJson(schema as any))) as RpcSerialization.CodecFor
39
40const layerProtocol = (codecFor: RpcSerialization.CodecFor) =>
41 Layer.effect(RpcServer.Protocol)(
42 Effect.map(Queue.unbounded<number>(), (disconnects) =>
43 RpcServer.Protocol.of({
44 run: () => Effect.never,
45 disconnects,
46 send: () => Effect.void,
47 end: () => Effect.void,
48 clientIds: Effect.succeed(new Set()),
49 initialMessage: Effect.succeedNone,
50 supportsAck: false,
51 supportsTransferables: false,
52 supportsSpanPropagation: false,
53 supportsNotifications: false,
54 codecFor
55 }))
56 )
57
58const makeHandlers = (
59 entities: Layer.Layer<never, never, any>,
60 codecFor: RpcSerialization.CodecFor = jsonCodec
61) =>
62 RunnerServer.layerHandlers.pipe(
63 Layer.provide(layerProtocol(codecFor)),
64 Layer.provideMerge(entities),
65 Layer.provideMerge(Sharding.layer),
66 Layer.provideMerge(Snowflake.layerGenerator),
67 Layer.provide(RunnerStorage.layerMemory),
68 Layer.provide(RunnerHealth.layerNoop),
69 Layer.provide(Runners.layerNoop),
70 Layer.provideMerge(MessageStorage.layerMemory),
71 Layer.provide(ShardingConfig.layer({
72 entityMailboxCapacity: 10,
73 entityTerminationTimeout: 0,
74 entityMessagePollInterval: 5000,
75 sendRetryInterval: 100,
76 refreshAssignmentsInterval: 0
77 }))
78 )
79
80const holeCodecHandlers = makeHandlers(
81 HoleCodecEntity.toLayer({ Double: ({ payload }) => Effect.succeed(payload.id * 2) }),
82 codecForJsonString
83)
84
85// A request for `entity`'s single RPC, addressed the way a sending runner would.
86const makeRequest = (
87 entity: Pick<Entity.Any, "type" | "getShardGroup">,
88 entityId: string,
89 options?: { readonly tag?: string; readonly payload?: unknown }
90) =>
91 Effect.gen(function*() {
92 const sharding = yield* Sharding.Sharding
93 const snowflake = yield* Snowflake.Generator
94 const id = EntityId.make(entityId)
95 const request: Envelope.PartialRequest = {
96 _tag: "Request",
97 requestId: snowflake.nextUnsafe(),
98 address: EntityAddress.make({
99 shardId: sharding.getShardId(id, entity.getShardGroup(id)),
100 entityType: EntityType.make(entity.type),
101 entityId: id
102 }),
103 tag: options?.tag ?? "ReproStream",
104 payload: options?.payload ?? { id: 1 },
105 headers: Headers.empty
106 }
107 return request
108 })
109
110// A runner RPC server whose replies are discarded.
111const makeServer = RpcServer.makeNoSerialization(Runners.Rpcs, { onFromServer: () => Effect.void })
112
113const writeRequest = (
114 server: RpcServer.RpcServer<any>,
115 tag: "Effect" | "Stream",
116 request: Envelope.PartialRequest,
117 options?: { readonly persisted?: boolean; readonly clientId?: number }
118) =>
119 server.write(options?.clientId ?? 0, {
120 _tag: "Request",
121 id: RpcMessage.RequestId(String(request.requestId)),
122 tag,
123 payload: { request, persisted: options?.persisted ?? false },
124 headers: Headers.empty
125 })
126
127const incomingRequest = (envelope: Envelope.PartialRequest, callerScope?: Scope.Scope) =>
128 new Message.IncomingRequest({
129 envelope,
130 lastSentReply: Option.none(),
131 respond: () => Effect.void,
132 codecFor: jsonCodec,
133 callerScope
134 })
135
136it.effect("completes a successful runner stream", () =>
137 Effect.gen(function*() {
138 yield* TestClock.adjust(1)
139 const snowflake = yield* Snowflake.Generator
140 const request = yield* makeRequest(ReproEntity, "one")
141 const client = yield* RpcTest.makeClient(Runners.Rpcs)
142 const queue = yield* client.Stream({ request, persisted: false }, { asQueue: true })
143 const first = yield* Queue.take(queue).pipe(Effect.timeout("1 second"), TestClock.withLive)
144 if (first._tag !== "Chunk") return assert.fail("expected the stream value before the terminal reply")
145 assert.deepStrictEqual(first.values, [1])
146 yield* client.Envelope({
147 envelope: new Envelope.AckChunk({
148 id: snowflake.nextUnsafe(),
149 address: request.address,
150 requestId: request.requestId,
151 replyId: Snowflake.Snowflake(first.id)
152 }),
153 persisted: false
154 })
155 const completion = yield* Effect.gen(function*() {
156 const last = yield* Queue.take(queue)
157 yield* Queue.take(queue).pipe(Effect.catchTag("Done", () => Effect.void))
158 return last
159 }).pipe(Effect.timeout("1 second"), TestClock.withLive)
160 assert.strictEqual(completion._tag, "WithExit")
161 }).pipe(Effect.provide(makeHandlers(
162 ReproEntity.toLayer({ ReproStream: () => Stream.make(1) })
163 ))))
164
165it.effect("releases a non-persisted Effect handler when the runner caller disconnects", () =>
166 Effect.gen(function*() {
167 const started = yield* Deferred.make<void>()
168 const stopped = yield* Deferred.make<void>()
169 yield* Effect.gen(function*() {
170 yield* TestClock.adjust(1)
171 const request = yield* makeRequest(HoleCodecEntity, "effect-disconnect", { tag: "Double" })
172 const server = yield* makeServer
173 yield* writeRequest(server, "Effect", request)
174 yield* TestClock.adjust(1)
175 assert.isTrue(yield* Deferred.isDone(started))
176 assert.isFalse(yield* Deferred.isDone(stopped))
177 yield* server.disconnect(0)
178 yield* TestClock.adjust(1)
179 assert.isTrue(yield* Deferred.isDone(stopped))
180 }).pipe(Effect.provide(makeHandlers(
181 HoleCodecEntity.toLayer({
182 Double: () =>
183 Rpc.fork(
184 Deferred.succeed(started, undefined).pipe(
185 Effect.andThen(Effect.never),
186 Effect.ensuring(Deferred.succeed(stopped, undefined))
187 )
188 )
189 })
190 )))
191 }))
192
193it.effect("forgets departed callers during rebuild without replaying or retaining mailbox slots", () =>
194 Effect.gen(function*() {
195 const rebuilding = yield* Deferred.make<void>()
196 const releaseBuild = yield* Deferred.make<void>()
197 const starts: Array<number> = []
198 let builds = 0
199 const entityLayer = ReproEntity.toLayer(
200 Effect.gen(function*() {
201 builds++
202 if (builds === 2) {
203 yield* Deferred.succeed(rebuilding, undefined)
204 yield* Deferred.await(releaseBuild)
205 }
206 return {
207 ReproStream: ({ payload }) => {
208 starts.push(payload.id)
209 return Rpc.fork(starts.length === 1 ? Stream.die("trigger entity rebuild") : Stream.never)
210 }
211 }
212 }),
213 { mailboxCapacity: 2 }
214 )
215 yield* Effect.gen(function*() {
216 yield* TestClock.adjust(1)
217 const sharding = yield* Sharding.Sharding
218 const server = yield* makeServer
219 const request = (id: number) => makeRequest(ReproEntity, "mailbox-during-rebuild", { payload: { id } })
220 const send = (clientId: number, id: number) =>
221 Effect.flatMap(request(id), (request) => writeRequest(server, "Stream", request, { clientId }))
222 yield* send(0, 1)
223 yield* TestClock.adjust("5 seconds")
224 assert.isTrue(yield* Deferred.isDone(rebuilding))
225 yield* send(1, 2)
226 yield* TestClock.adjust(1)
227 assert.deepStrictEqual(starts, [1])
228
229 const probe = (id: number) => Effect.flatMap(request(id), (request) => sharding.send(incomingRequest(request)))
230 // A and B fill both slots, proving B was admitted before its disconnect.
231 const full = yield* probe(3).pipe(Effect.flip)
232 assert.strictEqual(full._tag, "MailboxFull")
233 yield* server.disconnect(0).pipe(Effect.timeout("1 second"), TestClock.withLive)
234 yield* server.disconnect(1).pipe(Effect.timeout("1 second"), TestClock.withLive)
235 assert.isFalse(yield* Deferred.isDone(releaseBuild))
236 yield* Deferred.succeed(releaseBuild, undefined)
237 yield* TestClock.adjust(1)
238 assert.deepStrictEqual(starts, [1])
239
240 // Both the interrupted handler and the never-delivered request release
241 // their slots. Neither may run again when the rebuild finishes.
242 yield* probe(3)
243 yield* probe(4)
244 yield* TestClock.adjust(1)
245 assert.deepStrictEqual(starts, [1, 3, 4])
246 }).pipe(Effect.provide(makeHandlers(entityLayer)))
247 }))
248
249it.effect("releases a replayed handler when its caller disconnects during acquisition", () =>
250 Effect.gen(function*() {
251 const callerScope = yield* Effect.acquireRelease(Scope.make(), (scope) => Scope.close(scope, Exit.void))
252 const crash = yield* Deferred.make<void>()
253 const replaying = yield* Deferred.make<void>()
254 const disconnected = yield* Deferred.make<void>()
255 let starts = 0
256 let stops = 0
257 let crashes = 0
258 const entityLayer = ReproEntity.toLayer({
259 ReproStream: ({ payload }) => {
260 if (payload.id === 1) {
261 const invocation = ++starts
262 return Rpc.fork(
263 Stream.fromEffect(Effect.gen(function*() {
264 if (invocation === 2) yield* Deferred.succeed(replaying, undefined)
265 return yield* Effect.never
266 })).pipe(Stream.ensuring(Effect.sync(() => {
267 stops++
268 })))
269 )
270 }
271 return Rpc.fork(
272 ++crashes === 1
273 ? Stream.fromEffect(Deferred.await(crash).pipe(Effect.andThen(Effect.die("trigger entity rebuild"))))
274 : Stream.never
275 )
276 }
277 })
278 yield* Effect.gen(function*() {
279 yield* TestClock.adjust(1)
280 const sharding = yield* Sharding.Sharding
281 const request = (id: number) => makeRequest(ReproEntity, "disconnect-during-replay", { payload: { id } })
282 yield* sharding.send(incomingRequest(yield* request(1), callerScope))
283 yield* sharding.send(incomingRequest(yield* request(2)))
284 yield* TestClock.adjust(1)
285 assert.strictEqual(starts, 1)
286 assert.strictEqual(stops, 0)
287
288 // Close from another fiber as soon as the replacement handler starts,
289 // while the entity is still acquiring and replaying its requests.
290 yield* Deferred.await(replaying).pipe(
291 Effect.andThen(Scope.close(callerScope, Exit.void)),
292 Effect.andThen(Deferred.succeed(disconnected, undefined)),
293 Effect.forkChild({ startImmediately: true })
294 )
295 yield* Deferred.succeed(crash, undefined)
296 yield* TestClock.adjust("5 seconds")
297 assert.strictEqual(starts, 2)
298 assert.isTrue(yield* Deferred.isDone(disconnected))
299 // Both the original handler and its replacement must release resources.
300 assert.strictEqual(stops, 2)
301 }).pipe(Effect.provide(makeHandlers(entityLayer)))
302 }))
303
304it.effect("does not admit a request with an already-closed caller scope", () =>
305 Effect.gen(function*() {
306 let starts = 0
307 yield* Effect.gen(function*() {
308 yield* TestClock.adjust(1)
309 const sharding = yield* Sharding.Sharding
310 const callerScope = yield* Scope.make()
311 yield* Scope.close(callerScope, Exit.void)
312 const closed = incomingRequest(yield* makeRequest(ReproEntity, "closed-scope"), callerScope)
313 yield* Effect.exit(sharding.send(closed))
314 yield* TestClock.adjust(1)
315 assert.strictEqual(starts, 0)
316 // A live request must still fit in the only mailbox slot.
317 const live = incomingRequest(yield* makeRequest(ReproEntity, "closed-scope"))
318 const admitted = yield* Effect.exit(sharding.send(live))
319 assert.strictEqual(admitted._tag, "Success")
320 yield* TestClock.adjust(1)
321 assert.strictEqual(starts, 1)
322 }).pipe(Effect.provide(makeHandlers(
323 ReproEntity.toLayer({
324 ReproStream: () => {
325 starts++
326 return Rpc.fork(Stream.never)
327 }
328 }, { mailboxCapacity: 1 })
329 )))
330 }))
331
332for (const persisted of [false, true]) {
333 it.effect(`caller disconnect preserves ${persisted ? "persisted" : "Uninterruptible"} streams`, () =>
334 Effect.gen(function*() {
335 const pubsub = yield* Effect.acquireRelease(PubSub.unbounded<number>(), PubSub.shutdown)
336 const entity = persisted
337 ? ReproEntity.annotateRpcs(ClusterSchema.Persisted, true)
338 : ReproEntity.annotateRpcs(ClusterSchema.Uninterruptible, true)
339 yield* Effect.gen(function*() {
340 yield* TestClock.adjust(1)
341 const request = yield* makeRequest(entity, "protected-caller")
342 if (persisted) {
343 const driver = yield* MessageStorage.MemoryDriver
344 yield* driver.encoded.saveEnvelope({
345 envelope: yield* Schema.encodeEffect(Envelope.PartialJson)(request),
346 primaryKey: null,
347 deliverAt: null
348 })
349 }
350 const server = yield* makeServer
351 yield* writeRequest(server, "Stream", request, { persisted })
352 yield* TestClock.adjust(1)
353 assert.strictEqual(pubsub.subscribers.size, 1)
354 yield* server.disconnect(0)
355 yield* TestClock.adjust(1)
356 assert.strictEqual(pubsub.subscribers.size, 1)
357
358 if (!persisted) {
359 // An exemption from caller cleanup must not block explicit interrupts.
360 const sharding = yield* Sharding.Sharding
361 const snowflake = yield* Snowflake.Generator
362 yield* sharding.send(
363 new Message.IncomingEnvelope({
364 envelope: new Envelope.Interrupt({
365 id: snowflake.nextUnsafe(),
366 address: request.address,
367 requestId: request.requestId
368 })
369 })
370 )
371 yield* TestClock.adjust(1)
372 assert.strictEqual(pubsub.subscribers.size, 0)
373 }
374 }).pipe(Effect.provide(makeHandlers(
375 entity.toLayer({ ReproStream: () => Rpc.fork(Stream.fromPubSub(pubsub)) })
376 )))
377 }))
378}
379
380it.effect("fills the entity payload and reply holes with the serialization's codec", () =>
381 Effect.gen(function*() {
382 yield* TestClock.adjust(1)
383 // already encoded by the sender with the same hole codec
384 const request = yield* makeRequest(HoleCodecEntity, "hole", { tag: "Double", payload: JSON.stringify({ id: 21 }) })
385
386 const client = yield* RpcTest.makeClient(Runners.Rpcs)
387 const reply = yield* client.Effect({ request, persisted: false }).pipe(
388 Effect.timeout("1 second"),
389 TestClock.withLive
390 )
391
392 assert.strictEqual(typeof reply, "string", "the reply hole must carry the codec's output")
393 const decoded = JSON.parse(reply as any)
394 assert.strictEqual(decoded._tag, "WithExit")
395 assert.deepStrictEqual(decoded.exit, { _tag: "Success", value: 42 })
396 }).pipe(Effect.provide(holeCodecHandlers)))