Skip to content
Source
Sign inSign up
Open History to see the last commit

NetworkTransferMeasurement.integration.ts

215 lines7.0 KB

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 NodeWS from "@effect/platform-node/NodeWS";
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 readonly wireBytes: number;
32​
33export const measureHttpGet = Effect.fn("TransferBudget.measureHttpGet")(function* (input: {
34 readonly url: string;
35 readonly headers?: Readonly<Record<string, string>>;
36}) {
37 return yield* Effect.tryPromise({
38 try: () =>
39 new Promise<HttpTransferMeasurement>((resolve, reject) => {
40 let socketBytesBeforeResponse = 0;
41 const request = NodeHttp.get(
42 input.url,
43 {
44 agent: false,
45 headers: {
46 "accept-encoding": "gzip",
47 connection: "close",
48 ...input.headers,
49 },
50 },
51 (response) => {
52 const chunks: Buffer[] = [];
53 response.on("data", (chunk: Buffer) => chunks.push(chunk));
54 response.once("error", reject);
55 response.once("end", () => {
56 try {
57 const encodedBody = Buffer.concat(chunks);
58 const header = response.headers["content-encoding"];
59 const contentEncoding = Array.isArray(header)
60 ? (header[0] ?? null)
61 : (header ?? null);
62 const decodedBody =
63 contentEncoding === "gzip" ? NodeZlib.gunzipSync(encodedBody) : encodedBody;
64 resolve({
65 status: response.statusCode ?? 0,
66 contentEncoding,
67 encodedBody,
68 encodedBodyBytes: encodedBody.byteLength,
69 decodedBody,
70 decodedBodyBytes: decodedBody.byteLength,
71 wireBytes: Math.max(0, response.socket.bytesRead - socketBytesBeforeResponse),
72 });
73 } catch (cause) {
74 reject(cause);
75 }
76 });
77 },
78 );
79 request.once("socket", (socket) => {
80 socketBytesBeforeResponse = socket.bytesRead;
81 });
82 request.once("error", reject);
83 request.setTimeout(10_000, () => {
84 request.destroy(new Error(`Timed out reading ${input.url}`));
85 });
86 }),
87 catch: (cause) => new TransferHttpRequestError({ url: input.url, cause }),
88 });
89});
90​
91export interface WebSocketTransferTotals {
92 readonly wireBytes: number;
93 readonly decodedBytes: number;
94 readonly messages: number;
96​
97export interface WebSocketTransferRecorder {
98 readonly connect: (
99 url: string,
100 protocols: string | string[] | undefined,
101 cookie: string,
102 ) => globalThis.WebSocket;
103 readonly totals: () => WebSocketTransferTotals;
104 readonly negotiatedExtensions: () => string;
105 /** Resolves after upgrade; later totals exclude the upgrade response. */
106 readonly awaitOpen: Effect.Effect<void>;
108​
109interface NodeWebSocketWithTransport extends NodeWS.WebSocket {
110 readonly _socket?: {
111 readonly bytesRead: number;
112 };
114​
115function rawDataBytes(data: NodeWS.RawData): number {
116 if (Array.isArray(data)) {
117 return data.reduce((total, chunk) => total + chunk.byteLength, 0);
118 }
119 return data.byteLength;
121​
122function makeWebSocketTransferRecorder(): WebSocketTransferRecorder {
123 let socket: NodeWebSocketWithTransport | null = null;
124 // Keep wire totals after WebSocket close for reconnect measurements.
125 let transport: NodeWebSocketWithTransport["_socket"] | null = null;
126 let decodedBytes = 0;
127 let messages = 0;
128 let resolveOpen: () => void = () => {};
129 const opened = new Promise<void>((resolve) => {
130 resolveOpen = resolve;
131 });
132​
133 return {
134 connect: (url, protocols, cookie) => {
135 const nextSocket = new NodeWS.WebSocket(url, protocols, {
136 headers: { cookie },
137 perMessageDeflate: true,
138 }) as NodeWebSocketWithTransport;
139 socket = nextSocket;
140 nextSocket.once("open", () => {
141 transport = nextSocket._socket ?? null;
142 resolveOpen();
143 });
144 nextSocket.on("message", (data) => {
145 const bytes = rawDataBytes(data);
146 decodedBytes += bytes;
147 messages += 1;
148 });
149 return nextSocket as unknown as globalThis.WebSocket;
150 },
151 totals: () => ({
152 wireBytes: transport?.bytesRead ?? socket?._socket?.bytesRead ?? 0,
153 decodedBytes,
154 messages,
155 }),
156 negotiatedExtensions: () => socket?.extensions ?? "",
157 awaitOpen: Effect.promise(() => opened).pipe(
158 Effect.timeoutOrElse({
159 duration: "10 seconds",
160 orElse: () => Effect.die(new Error("Timed out waiting for the WebSocket to open")),
161 }),
162 ),
163 };
165​
166function layerCountingWsRpcProtocol(input: {
167 readonly url: string;
168 readonly cookie: string;
169 readonly recorder: WebSocketTransferRecorder;
170}) {
171 const layerWebSocketConstructor = Layer.succeed(Socket.WebSocketConstructor, (url, protocols) =>
172 input.recorder.connect(url, protocols as string | string[] | undefined, input.cookie),
173 );
174 return RpcClient.layerProtocolSocket().pipe(
175 Layer.provide(
176 Socket.layerWebSocket(input.url, { openTimeout: "10 seconds" }).pipe(
177 Layer.provide(layerWebSocketConstructor),
178 ),
179 ),
180 Layer.provide(RpcSerialization.layerJson),
181 );
183​
184const makeCountingWsRpcClient = RpcClient.make(WsRpcGroup);
185export type CountingWsRpcClient = Effect.Success<typeof makeCountingWsRpcClient>;
186​
187export interface MeasuredWsClient {
188 readonly client: CountingWsRpcClient;
189 readonly recorder: WebSocketTransferRecorder;
190 readonly scope: Scope.Scope;
191 readonly close: Effect.Effect<void>;
193​
194export const openMeasuredWsClient = Effect.fn("TransferBudget.openMeasuredWsClient")(
195 function* (input: { readonly url: string; readonly cookie: string }) {
196 const recorder = makeWebSocketTransferRecorder();
197 const parent = yield* Effect.scope;
198 const scope = yield* Scope.fork(parent);
199 const protocol = yield* Layer.buildWithScope(
200 layerCountingWsRpcProtocol({ url: input.url, cookie: input.cookie, recorder }),
201 scope,
202 );
203 const client = yield* makeCountingWsRpcClient.pipe(
204 Effect.provide(protocol),
205 Scope.provide(scope),
206 );
207 yield* recorder.awaitOpen;
208 return {
209 client,
210 recorder,
211 scope,
212 close: Scope.close(scope, Exit.void),
213 } satisfies MeasuredWsClient;
214 },
215);