1import { assert, describe, it } from "@effect/vitest";
2import {
3 AuthEnvironmentMaintainScope,
4 AuthOrchestrationReadScope,
5 type AuthEnvironmentScope,
6 ScheduledTaskError,
7 ScheduledTaskId,
8 WS_METHODS,
9 WsRpcGroup,
10} from "@t3tools/contracts";
11import * as Deferred from "effect/Deferred";
12import * as Duration from "effect/Duration";
13import * as Effect from "effect/Effect";
14import * as Exit from "effect/Exit";
15import * as Fiber from "effect/Fiber";
16import * as Layer from "effect/Layer";
17import * as Metric from "effect/Metric";
18import * as Option from "effect/Option";
19import * as Stream from "effect/Stream";
20import * as Tracer from "effect/Tracer";
21import * as RpcTest from "effect/rpc/RpcTest";
22import * as TestClock from "effect/testing/TestClock";
23
24import { RPC_REQUIRED_SCOPES } from "../auth/RpcAuthorization.ts";
25import * as RpcAuthorization from "../auth/RpcAuthorization.ts";
26import { RpcInstrumentation, rpcInstrumentationLayer } from "./RpcInstrumentation.ts";
27
28type WsRpcMethod = keyof typeof RPC_REQUIRED_SCOPES;
29
30/** The server group narrowed to `tags`, so a test only implements the handlers it calls. */
31const groupOf = <const Tags extends ReadonlyArray<WsRpcMethod>>(...tags: Tags) =>
32 WsRpcGroup.omit(
33 ...[...WsRpcGroup.requests.keys()].filter(
34 (tag): tag is Exclude<WsRpcMethod, Tags[number]> =>
35 !(tags as ReadonlyArray<string>).includes(tag),
36 ),
37 ).middleware(RpcInstrumentation);
38
39/** The middleware ws.ts installs for a connection with `scopes`. */
40const connectionMiddleware = (scopes: ReadonlyArray<AuthEnvironmentScope>) =>
41 Layer.merge(RpcAuthorization.layer(scopes), rpcInstrumentationLayer);
42const readOnlyConnection = connectionMiddleware([AuthOrchestrationReadScope]);
43const taskId = ScheduledTaskId.make("scheduled-task:instrumented");
44const rpcSpanDefaults = { "rpc.transport": "websocket", "rpc.system": "effect-rpc" };
45
46/** Runs a test with a fresh metric registry and a tracer that keeps every span it ends. */
47const withTelemetry = <A, E, R>(
48 test: (ended: ReadonlyArray<Tracer.NativeSpan>) => Effect.Effect<A, E, R>,
49) => {
50 const ended: Array<Tracer.NativeSpan> = [];
51 const tracer = Tracer.make({
52 span: (options) => {
53 const span = new Tracer.NativeSpan(options);
54 const end = span.end.bind(span);
55 span.end = (endTime, exit) => {
56 end(endTime, exit);
57 ended.push(span);
58 };
59 return span;
60 },
61 });
62 return test(ended).pipe(
63 Effect.scoped,
64 Effect.withTracer(tracer),
65 Effect.provideService(Metric.MetricRegistry, new Map()),
66 );
67};
68
69const rpcSpans = (ended: ReadonlyArray<Tracer.NativeSpan>) =>
70 ended.filter((span) => span.name.startsWith("ws.rpc."));
71
72// RpcTest keeps Effect's own RpcServer/RpcClient spans, which ws.ts turns off with
73// `disableTracing: true`. Everything else comes from the middleware or the handlers.
74const appSpans = (ended: ReadonlyArray<Tracer.NativeSpan>) =>
75 ended.filter((span) => span.attributes.get("rpc.system.name") !== "effect_rpc");
76
77const exitTag = (span: Tracer.NativeSpan | undefined) =>
78 span?.status._tag === "Ended" ? span.status.exit._tag : undefined;
79
80const parentSpanId = (span: Tracer.NativeSpan | undefined) =>
81 span?.parent._tag === "Some" ? span.parent.value.spanId : undefined;
82
83const requestCount = (
84 snapshots: ReadonlyArray<Metric.Metric.Snapshot>,
85 method: string,
86 outcome: string,
87) =>
88 snapshots.find(
89 (snapshot): snapshot is Extract<Metric.Metric.Snapshot, { readonly type: "Counter" }> =>
90 snapshot.type === "Counter" &&
91 snapshot.id === "t3_rpc_requests_total" &&
92 snapshot.attributes?.["method"] === method &&
93 snapshot.attributes?.["outcome"] === outcome,
94 )?.state;
95
96const requestDuration = (snapshots: ReadonlyArray<Metric.Metric.Snapshot>, method: string) =>
97 snapshots.find(
98 (snapshot): snapshot is Extract<Metric.Metric.Snapshot, { readonly type: "Histogram" }> =>
99 snapshot.type === "Histogram" &&
100 snapshot.id === "t3_rpc_request_duration" &&
101 snapshot.attributes?.["method"] === method,
102 )?.state;
103
104describe("WS RPC instrumentation middleware", () => {
105 it.effect("records one span and request metric per call, including rejected calls", () =>
106 withTelemetry((ended) =>
107 Effect.gen(function* () {
108 const group = groupOf(
109 WS_METHODS.serverProbe,
110 WS_METHODS.scheduledTasksList,
111 WS_METHODS.serverRetryResourceTelemetry,
112 WS_METHODS.pullRequestsSubscribeRefreshes,
113 WS_METHODS.scheduledTasksSubscribe,
114 WS_METHODS.serverGetSettings,
115 );
116 const client = yield* RpcTest.makeClient(group).pipe(
117 Effect.provide(
118 Layer.mergeAll(
119 group.toLayerHandler(WS_METHODS.serverProbe, () =>
120 Effect.succeed({}).pipe(Effect.withSpan("serverProbe.child")),
121 ),
122 group.toLayerHandler(WS_METHODS.scheduledTasksList, () =>
123 Effect.annotateCurrentSpan({ "scheduled_task.id": taskId }).pipe(
124 Effect.andThen(new ScheduledTaskError({ message: "List failed." })),
125 ),
126 ),
127 group.toLayerHandler(WS_METHODS.serverRetryResourceTelemetry, () =>
128 Effect.die("authorization let a rejected call through"),
129 ),
130 group.toLayerHandler(WS_METHODS.pullRequestsSubscribeRefreshes, () =>
131 Stream.make(1, 2),
132 ),
133 group.toLayerHandler(WS_METHODS.scheduledTasksSubscribe, () =>
134 Stream.concat(
135 Stream.make({ tasks: [] }),
136 Stream.fail(new ScheduledTaskError({ message: "Subscription failed." })),
137 ),
138 ),
139 group.toLayerHandler(WS_METHODS.serverGetSettings, () => Effect.die("broken")),
140 readOnlyConnection,
141 ),
142 ),
143 );
144
145 assert.deepStrictEqual(yield* client[WS_METHODS.serverProbe]({}), {});
146 const listError = yield* client[WS_METHODS.scheduledTasksList]({}).pipe(Effect.flip);
147 assert.equal(listError._tag, "ScheduledTaskError");
148 // Retrying telemetry needs operate scope, so the handler must not run.
149 const rejection = yield* client[WS_METHODS.serverRetryResourceTelemetry]({}).pipe(
150 Effect.flip,
151 );
152 assert.equal(rejection._tag, "EnvironmentAuthorizationError");
153 const refreshes = yield* Stream.runCollect(
154 client[WS_METHODS.pullRequestsSubscribeRefreshes]({}),
155 );
156 assert.deepStrictEqual(Array.from(refreshes), [1, 2]);
157 const subscribeError = yield* Stream.runDrain(
158 client[WS_METHODS.scheduledTasksSubscribe]({}),
159 ).pipe(Effect.flip);
160 assert.equal(subscribeError._tag, "ScheduledTaskError");
161 const settingsExit = yield* Effect.exit(client[WS_METHODS.serverGetSettings]({}));
162 assert.isTrue(Exit.hasDies(settingsExit));
163
164 assert.deepStrictEqual(
165 rpcSpans(ended).map((span) => [span.name, Object.fromEntries(span.attributes)]),
166 [
167 [
168 "ws.rpc.server.probe",
169 {
170 ...rpcSpanDefaults,
171 "rpc.method": WS_METHODS.serverProbe,
172 "rpc.aggregate": "server",
173 },
174 ],
175 [
176 "ws.rpc.scheduledTasks.list",
177 {
178 ...rpcSpanDefaults,
179 "rpc.method": WS_METHODS.scheduledTasksList,
180 "rpc.aggregate": "scheduledTasks",
181 "scheduled_task.id": taskId,
182 },
183 ],
184 [
185 "ws.rpc.server.retryResourceTelemetry",
186 {
187 ...rpcSpanDefaults,
188 "rpc.method": WS_METHODS.serverRetryResourceTelemetry,
189 "rpc.aggregate": "server",
190 },
191 ],
192 [
193 "ws.rpc.pullRequests.subscribeRefreshes",
194 {
195 ...rpcSpanDefaults,
196 "rpc.method": WS_METHODS.pullRequestsSubscribeRefreshes,
197 "rpc.aggregate": "pull-requests",
198 },
199 ],
200 [
201 "ws.rpc.scheduledTasks.subscribe",
202 {
203 ...rpcSpanDefaults,
204 "rpc.method": WS_METHODS.scheduledTasksSubscribe,
205 "rpc.aggregate": "scheduledTasks",
206 },
207 ],
208 [
209 "ws.rpc.server.getSettings",
210 {
211 ...rpcSpanDefaults,
212 "rpc.method": WS_METHODS.serverGetSettings,
213 "rpc.aggregate": "server",
214 },
215 ],
216 ],
217 );
218 assert.deepStrictEqual(rpcSpans(ended).map(exitTag), [
219 "Success",
220 "Failure",
221 "Failure",
222 "Success",
223 "Failure",
224 "Failure",
225 ]);
226 const probeSpan = rpcSpans(ended)[0];
227 const child = ended.find((span) => span.name === "serverProbe.child");
228 assert.equal(parentSpanId(child), probeSpan?.spanId);
229
230 const snapshots = yield* Metric.snapshot;
231 assert.deepStrictEqual(
232 [
233 requestCount(snapshots, WS_METHODS.serverProbe, "success"),
234 requestCount(snapshots, WS_METHODS.scheduledTasksList, "failure"),
235 requestCount(snapshots, WS_METHODS.serverRetryResourceTelemetry, "failure"),
236 requestCount(snapshots, WS_METHODS.pullRequestsSubscribeRefreshes, "success"),
237 requestCount(snapshots, WS_METHODS.scheduledTasksSubscribe, "failure"),
238 requestCount(snapshots, WS_METHODS.serverGetSettings, "failure"),
239 ].map((state) => state?.count),
240 [1, 1, 1, 1, 1, 1],
241 );
242 for (const method of group.requests.keys()) {
243 assert.equal(requestDuration(snapshots, method)?.count, 1);
244 }
245 }),
246 ),
247 );
248
249 it.effect("keeps a stream's span and duration open until the subscription is interrupted", () =>
250 withTelemetry((ended) =>
251 Effect.gen(function* () {
252 const waiting = yield* Deferred.make<void>();
253 const group = groupOf(WS_METHODS.pullRequestsSubscribeRefreshes);
254 const client = yield* RpcTest.makeClient(group).pipe(
255 Effect.provide(
256 Layer.mergeAll(
257 group.toLayerHandler(WS_METHODS.pullRequestsSubscribeRefreshes, () =>
258 Stream.concat(
259 Stream.make(1),
260 Stream.fromEffect(
261 Deferred.succeed(waiting, undefined).pipe(
262 Effect.andThen(Effect.never),
263 Effect.withSpan("refreshes.wait"),
264 ),
265 ),
266 ),
267 ),
268 readOnlyConnection,
269 ),
270 ),
271 );
272
273 // Wall and monotonic time disagree before the call starts, so a duration that mixes
274 // the two clocks is caught.
275 yield* TestClock.adjust(Duration.seconds(1));
276 yield* TestClock.setTime(0);
277 const consumer = yield* Stream.runDrain(
278 client[WS_METHODS.pullRequestsSubscribeRefreshes]({}),
279 ).pipe(Effect.forkChild);
280 yield* Deferred.await(waiting);
281 yield* TestClock.adjust(Duration.millis(250));
282 // A backward wall-clock correction must not shorten the measured duration.
283 yield* TestClock.setTime(0);
284 assert.deepStrictEqual(appSpans(ended), []);
285
286 // The client waits for the server to stop the call, which ends the RPC span.
287 yield* Fiber.interrupt(consumer);
288
289 const [rpcSpan] = rpcSpans(ended);
290 assert.equal(rpcSpan?.name, "ws.rpc.pullRequests.subscribeRefreshes");
291 assert.equal(exitTag(rpcSpan), "Failure");
292 const waitSpan = ended.find((span) => span.name === "refreshes.wait");
293 assert.equal(parentSpanId(waitSpan), rpcSpan?.spanId);
294
295 const snapshots = yield* Metric.snapshot;
296 assert.equal(
297 requestCount(snapshots, WS_METHODS.pullRequestsSubscribeRefreshes, "interrupt")?.count,
298 1,
299 );
300 const duration = requestDuration(snapshots, WS_METHODS.pullRequestsSubscribeRefreshes);
301 assert.equal(duration?.count, 1);
302 assert.equal(duration?.sum, 250);
303 }),
304 ),
305 );
306
307 it.effect("records metrics but no spans for methods with tracing disabled", () =>
308 withTelemetry((ended) =>
309 Effect.gen(function* () {
310 const group = groupOf(WS_METHODS.serverSignalProcess);
311 const client = yield* RpcTest.makeClient(group).pipe(
312 Effect.provide(
313 Layer.mergeAll(
314 group.toLayerHandler(WS_METHODS.serverSignalProcess, (input) =>
315 Effect.succeed({
316 pid: input.pid,
317 signal: input.signal,
318 signaled: true,
319 message: Option.none(),
320 }).pipe(Effect.withSpan("signalProcess.child")),
321 ),
322 connectionMiddleware([AuthOrchestrationReadScope, AuthEnvironmentMaintainScope]),
323 ),
324 ),
325 );
326
327 const input = { pid: 4242, startTimeMs: 0, signal: "SIGINT" } as const;
328 assert.equal((yield* client[WS_METHODS.serverSignalProcess](input)).signaled, true);
329
330 assert.deepStrictEqual(appSpans(ended), []);
331 const snapshots = yield* Metric.snapshot;
332 assert.equal(requestCount(snapshots, WS_METHODS.serverSignalProcess, "success")?.count, 1);
333 }),
334 ),
335 );
336});