1// @effect-diagnostics nodeBuiltinImport:off - Measures the real Node HTTP and WebSocket transports.
2import * as NodeHttp from "node:http";
3import * as NodeZlib from "node:zlib";
4
5import * as NodeSocket from "@effect/platform-node/NodeSocket";
6import { WsRpcGroup } from "@t3tools/contracts";
7import * as Effect from "effect/Effect";
8import * as Exit from "effect/Exit";
9import * as Layer from "effect/Layer";
10import * as Schema from "effect/Schema";
11import * as Scope from "effect/Scope";
12import { RpcClient, RpcSerialization } from "effect/rpc";
13import * as Socket from "effect/socket/Socket";
14
15export class TransferHttpRequestError extends Schema.TaggedError<TransferHttpRequestError>()(
16 "TransferHttpRequestError",
17 {
18 url: Schema.String,
19 cause: Schema.Defect(),
20 },
21) {}
22
23export interface HttpTransferMeasurement {
24 readonly status: number;
25 readonly contentEncoding: string | null;
26 readonly encodedBody: Uint8Array;
27 readonly encodedBodyBytes: number;
28 readonly decodedBody: Uint8Array;
29 readonly decodedBodyBytes: number;
30 /** HTTP response bytes read from the socket, including status line and headers. */
31 readonly wireBytes: number;
32}
33
34export const measureHttpGet = Effect.fn("TransferBudget.measureHttpGet")(function* (input: {
35 readonly url: string;
36 readonly headers?: Readonly<Record<string, string>>;
37}) {
38 return yield* Effect.tryPromise({
39 try: () =>
40 new Promise<HttpTransferMeasurement>((resolve, reject) => {
41 let socketBytesBeforeResponse = 0;
42 const request = NodeHttp.get(
43 input.url,
44 {
45 agent: false,
46 headers: {
47 "accept-encoding": "gzip",
48 connection: "close",
49 ...input.headers,
50 },
51 },
52 (response) => {
53 const chunks: Buffer[] = [];
54 response.on("data", (chunk: Buffer) => chunks.push(chunk));
55 response.once("error", reject);
56 response.once("end", () => {
57 try {
58 const encodedBody = Buffer.concat(chunks);
59 const header = response.headers["content-encoding"];
60 const contentEncoding = Array.isArray(header)
61 ? (header[0] ?? null)
62 : (header ?? null);
63 const decodedBody =
64 contentEncoding === "gzip" ? NodeZlib.gunzipSync(encodedBody) : encodedBody;
65 resolve({
66 status: response.statusCode ?? 0,
67 contentEncoding,
68 encodedBody,
69 encodedBodyBytes: encodedBody.byteLength,
70 decodedBody,
71 decodedBodyBytes: decodedBody.byteLength,
72 wireBytes: Math.max(0, response.socket.bytesRead - socketBytesBeforeResponse),
73 });
74 } catch (cause) {
75 reject(cause);
76 }
77 });
78 },
79 );
80 request.once("socket", (socket) => {
81 socketBytesBeforeResponse = socket.bytesRead;
82 });
83 request.once("error", reject);
84 request.setTimeout(10_000, () => {
85 request.destroy(new Error(`Timed out reading ${input.url}`));
86 });
87 }),
88 catch: (cause) => new TransferHttpRequestError({ url: input.url, cause }),
89 });
90});
91
92export interface WebSocketTransferTotals {
93 readonly wireBytes: number;
94 readonly decodedBytes: number;
95 readonly messages: number;
96}
97
98export interface WebSocketTransferRecorder {
99 readonly connect: (
100 url: string,
101 protocols: string | string[] | undefined,
102 cookie: string,
103 ) => globalThis.WebSocket;
104 readonly totals: () => WebSocketTransferTotals;
105 readonly negotiatedExtensions: () => string;
106 /** Resolves once the upgrade completes, so totals taken after it exclude the upgrade response. */
107 readonly awaitOpen: Effect.Effect<void>;
108}
109
110interface NodeWebSocketWithTransport extends NodeSocket.NodeWS.WebSocket {
111 readonly _socket?: {
112 readonly bytesRead: number;
113 };
114}
115
116function rawDataBytes(data: NodeSocket.NodeWS.RawData): number {
117 if (Array.isArray(data)) {
118 return data.reduce((total, chunk) => total + chunk.byteLength, 0);
119 }
120 return data.byteLength;
121}
122
123function makeWebSocketTransferRecorder(): WebSocketTransferRecorder {
124 let socket: NodeWebSocketWithTransport | null = null;
125 // Held separately from the WebSocket so wire totals survive a close, which
126 // is when a reconnect measurement reads them.
127 let transport: NodeWebSocketWithTransport["_socket"] | null = null;
128 let decodedBytes = 0;
129 let messages = 0;
130 let resolveOpen: () => void = () => {};
131 const opened = new Promise<void>((resolve) => {
132 resolveOpen = resolve;
133 });
134
135 return {
136 connect: (url, protocols, cookie) => {
137 const nextSocket = new NodeSocket.NodeWS.WebSocket(url, protocols, {
138 headers: { cookie },
139 perMessageDeflate: true,
140 }) as NodeWebSocketWithTransport;
141 socket = nextSocket;
142 nextSocket.once("open", () => {
143 transport = nextSocket._socket ?? null;
144 resolveOpen();
145 });
146 nextSocket.on("message", (data) => {
147 const bytes = rawDataBytes(data);
148 decodedBytes += bytes;
149 messages += 1;
150 });
151 return nextSocket as unknown as globalThis.WebSocket;
152 },
153 totals: () => ({
154 wireBytes: transport?.bytesRead ?? socket?._socket?.bytesRead ?? 0,
155 decodedBytes,
156 messages,
157 }),
158 negotiatedExtensions: () => socket?.extensions ?? "",
159 awaitOpen: Effect.promise(() => opened).pipe(
160 Effect.timeoutOrElse({
161 duration: "10 seconds",
162 orElse: () => Effect.die(new Error("Timed out waiting for the WebSocket to open")),
163 }),
164 ),
165 };
166}
167
168function layerCountingWsRpcProtocol(input: {
169 readonly url: string;
170 readonly cookie: string;
171 readonly recorder: WebSocketTransferRecorder;
172}) {
173 // Socket.makeWebSocket only ever passes its `protocols` option here.
174 const layerWebSocketConstructor = Layer.succeed(Socket.WebSocketConstructor, (url, protocols) =>
175 input.recorder.connect(url, protocols as string | string[] | undefined, input.cookie),
176 );
177 return RpcClient.layerProtocolSocket().pipe(
178 Layer.provide(
179 Socket.layerWebSocket(input.url, { openTimeout: "10 seconds" }).pipe(
180 Layer.provide(layerWebSocketConstructor),
181 ),
182 ),
183 Layer.provide(RpcSerialization.layerJson),
184 );
185}
186
187const makeCountingWsRpcClient = RpcClient.make(WsRpcGroup);
188export type CountingWsRpcClient = Effect.Success<typeof makeCountingWsRpcClient>;
189
190export interface MeasuredWsClient {
191 readonly client: CountingWsRpcClient;
192 readonly recorder: WebSocketTransferRecorder;
193 /** Fork subscription consumers here so they stop before the socket closes. */
194 readonly scope: Scope.Scope;
195 /** Closes the socket now. The enclosing scope closes it otherwise. */
196 readonly close: Effect.Effect<void>;
197}
198
199/**
200 * Opens one WebSocket RPC client on a child of the current scope. Several
201 * clients can share one test scope and still disconnect independently, which
202 * a reconnect measurement needs.
203 */
204export const openMeasuredWsClient = Effect.fn("TransferBudget.openMeasuredWsClient")(
205 function* (input: { readonly url: string; readonly cookie: string }) {
206 const recorder = makeWebSocketTransferRecorder();
207 const parent = yield* Effect.scope;
208 const scope = yield* Scope.fork(parent);
209 const protocol = yield* Layer.buildWithScope(
210 layerCountingWsRpcProtocol({ url: input.url, cookie: input.cookie, recorder }),
211 scope,
212 );
213 const client = yield* makeCountingWsRpcClient.pipe(
214 Effect.provide(protocol),
215 Scope.provide(scope),
216 );
217 yield* recorder.awaitOpen;
218 return {
219 client,
220 recorder,
221 scope,
222 close: Scope.close(scope, Exit.void),
223 } satisfies MeasuredWsClient;
224 },
225);