Key weak collections by net sockets and servers, and carry the socket in undici client diagnostics (#1742)

packages/compiler/shims/undici/core/client.ts

@@ -370,7 +370,7 @@ export class Client extends DispatcherBase {
       localAddress: this.__localAddress,
     };
     if (channels.beforeConnect.hasSubscribers) {
-      channels.beforeConnect.publish({ connectParams });
+      channels.beforeConnect.publish({ connectParams, connector: this.__connector });
     }
     let settled = false;
     const onConnected = (err: Error | null, socket: Socket | null): void => {
@@ -428,7 +428,7 @@ export class Client extends DispatcherBase {
       if (this.__socket === socket) this.__onSocketClose(socket);
     });
     if (channels.connected.hasSubscribers) {
-      channels.connected.publish({ connectParams });
+      channels.connected.publish({ connectParams, connector: this.__connector, socket });
     }
     this.emit("connect", this.__url, [this as Dispatcher]);
     this.__resume(false);
@@ -447,7 +447,7 @@ export class Client extends DispatcherBase {
   ): void {
     this.__connecting = false;
     if (channels.connectError.hasSubscribers) {
-      channels.connectError.publish({ connectParams, error: err });
+      channels.connectError.publish({ connectParams, connector: this.__connector, error: err });
     }
     this.__onError(err);
     this.emit("connectionError", this.__url, [this as Dispatcher], err);
@@ -658,7 +658,7 @@ export class Client extends DispatcherBase {
     }
     header += request.headers;
     if (channels.sendHeaders.hasSubscribers) {
-      channels.sendHeaders.publish({ request: request.diagnosticView(), headers: header });
+      channels.sendHeaders.publish({ request: request.diagnosticView(), headers: header, socket });
     }
     if (body === null) {
       if (contentLength === 0) socket.write(`${header}content-length: 0\r\n\r\n`);

packages/compiler/src/ir/ir.ts

@@ -4348,7 +4348,9 @@ function canBoxDynCompositeAt(
       }
     }
     default:
-      return false;
+      // Runtime handles (a socket in a diagnostics message) box by
+      // reference wherever the walker meets them, as at the top level.
+      return DYN_HANDLE_KINDS.has(t.kind);
   }
 }
 

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/6620-weak-keys-net-sockets.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/6620-weak-keys-net-sockets.ts"
+ ],
+ "diags": []
+}

packages/runtime/src/scr_json.c

@@ -14065,7 +14065,8 @@ static int scr_weak_key(const ScrDyn *key, void **ptr, unsigned *kind) {
     if (key->v.handle.tag != SCR_DYNH_ARRAY_BUFFER && key->v.handle.tag != SCR_DYNH_SHARED_ARRAY_BUFFER && key->v.handle.tag != SCR_DYNH_WEAK_MAP &&
         key->v.handle.tag != SCR_DYNH_WEAK_SET && key->v.handle.tag != SCR_DYNH_STDIO &&
         key->v.handle.tag != SCR_DYNH_FETCH_REQUEST && key->v.handle.tag != SCR_DYNH_FETCH_RESPONSE &&
-        key->v.handle.tag != SCR_DYNH_MAP && key->v.handle.tag != SCR_DYNH_SET) return -1;
+        key->v.handle.tag != SCR_DYNH_MAP && key->v.handle.tag != SCR_DYNH_SET &&
+        key->v.handle.tag != SCR_DYNH_NET_SOCKET && key->v.handle.tag != SCR_DYNH_NET_SERVER) return -1;
     *ptr = key->v.handle.ptr; *kind = 4 + key->v.handle.tag; return 1;
   case SCR_DYN_TYPED_REF:
     if (!key->v.typed_ref.observed) return -1;

packages/runtime/src/scr_net.c

@@ -678,6 +678,7 @@ void scr_net_server_release(ScrNetServer *s) {
 #ifdef SCR_RC_AUDIT
     scr_net_live--;
 #endif
+    scr_weak_dispose_inline(s); /* weak collections keyed by this socket or server */
     free(s);
   }
 }
@@ -724,6 +725,7 @@ void scr_net_sock_release(ScrNetSocket *s) {
 #ifdef SCR_RC_AUDIT
     scr_net_live--;
 #endif
+    scr_weak_dispose_inline(s); /* weak collections keyed by this socket or server */
     free(s);
   }
 }

tests/corpus/6620-weak-keys-net-sockets.ts

@@ -0,0 +1,81 @@
+// Sockets and servers key WeakMaps and WeakSets by identity: the socket a
+// connect callback sees, the one a 'connection' listener sees, and the
+// same socket read back out of a diagnostics_channel message are one key.
+// Entries of closed sockets stop answering once the socket is gone.
+import * as net from "node:net";
+import type { Server, Socket } from "node:net";
+import { channel } from "node:diagnostics_channel";
+
+interface Conn {
+  label: string;
+  bytes: number;
+}
+
+const connections = new WeakMap<Socket, Conn>();
+const seen = new WeakSet<Socket>();
+const servers = new WeakMap<Server, string>();
+const opened = channel("example:socket:opened");
+
+opened.subscribe((message) => {
+  const { socket, label } = message as { socket: Socket; label: string };
+  console.log(`[diag] ${label} known ${connections.has(socket)}, seen ${seen.has(socket)}`);
+  const conn = connections.get(socket);
+  if (conn !== undefined) conn.label = `${conn.label}+diag`;
+});
+
+// Each dial waits for both ends to close, so the two sides log in order.
+let serverClosed: () => void = () => {};
+
+const server = net.createServer((socket) => {
+  seen.add(socket);
+  connections.set(socket, { label: "server side", bytes: 0 });
+  socket.on("data", (chunk) => {
+    const conn = connections.get(socket);
+    if (conn !== undefined) conn.bytes += chunk.length;
+    socket.end(`echo ${chunk.toString()}`);
+  });
+  socket.on("close", () => {
+    const conn = connections.get(socket);
+    console.log("server socket closed", conn?.label, conn?.bytes);
+    console.log("delete", connections.delete(socket), connections.delete(socket));
+    serverClosed();
+  });
+});
+servers.set(server, "echo server");
+
+function dial(port: number, label: string): Promise<void> {
+  const bothClosed = new Promise<void>((resolve) => {
+    serverClosed = resolve;
+  });
+  const clientClosed = new Promise<void>((resolve) => {
+    const client: Socket = net.connect(port, "127.0.0.1", () => {
+      connections.set(client, { label, bytes: 0 });
+      console.log(label, "connected; has", connections.has(client), "seen", seen.has(client));
+      opened.publish({ socket: client, label });
+      client.write(label);
+    });
+    let reply = "";
+    client.on("data", (chunk) => {
+      reply += chunk.toString();
+    });
+    client.on("close", () => {
+      console.log(label, "reply", JSON.stringify(reply), "entry", connections.get(client)?.label);
+      resolve();
+    });
+  });
+  return Promise.all([clientClosed, bothClosed]).then(() => undefined);
+}
+
+server.listen(0, "127.0.0.1", async () => {
+  console.log("server known", servers.has(server), servers.get(server));
+  const address = server.address();
+  const port = typeof address === "object" && address !== null ? address.port : 0;
+  await dial(port, "first");
+  await dial(port, "second");
+  try {
+    connections.set(1 as unknown as Socket, { label: "number", bytes: 0 });
+  } catch (err) {
+    console.log((err as Error).name, (err as Error).message);
+  }
+  server.close(() => console.log("closed; server still known", servers.has(server)));
+});

tests/fixtures/npm/cases/undici5/main.ts

@@ -81,6 +81,34 @@ channel("undici:request:trailers").subscribe((message) => {
   if (t !== undefined) console.log(`  [diag] ${t.label} done after ${t.dispatched} dispatch`);
 });
 
+// A transport's connection registry: sockets the wrapped connector dialed,
+// keyed weakly, and the client diagnostics messages that carry them.
+const dialed = new WeakMap<Socket, number>();
+const clientDiag = { beforeConnect: 0, connector: 0, connected: 0, sendHeaders: 0, known: 0, unknown: 0 };
+const noteSocket = (socket: Socket): void => {
+  const n = dialed.get(socket);
+  if (n === undefined) clientDiag.unknown++;
+  else clientDiag.known++;
+};
+channel("undici:client:beforeConnect").subscribe((message) => {
+  const { connector } = message as { connector: unknown };
+  clientDiag.beforeConnect++;
+  if (typeof connector === "function") clientDiag.connector++;
+});
+channel("undici:client:connected").subscribe((message) => {
+  const { socket, connectParams } = message as {
+    socket: Socket;
+    connectParams: { protocol: string; hostname: string };
+  };
+  clientDiag.connected++;
+  if (connectParams.protocol === "http:" && connectParams.hostname === "127.0.0.1") noteSocket(socket);
+});
+channel("undici:client:sendHeaders").subscribe((message) => {
+  const { socket, headers } = message as { socket: Socket; headers: string };
+  clientDiag.sendHeaders++;
+  if (headers.length > 0) noteSocket(socket);
+});
+
 /** A transport-style Agent: one Pool per origin built by the factory, a
  * wrapped connector, and a Client interceptor that marks dispatches. */
 class TracingAgent extends Agent {
@@ -102,7 +130,9 @@ class TracingAgent extends Agent {
           pipelining: 1,
           connect(options, callback) {
             opened++;
-            return connect(options, callback);
+            const socket = connect(options, callback);
+            dialed.set(socket, opened);
+            return socket;
           },
           interceptors: {
             Client: [
@@ -198,6 +228,9 @@ server.listen(0, "127.0.0.1", async () => {
     await r.body.text();
   }
   console.log("connections opened:", agent.connections);
+  console.log(
+    `client diagnostics: ${clientDiag.beforeConnect} before connect (${clientDiag.connector} with a connector), ${clientDiag.connected} connected, ${clientDiag.sendHeaders} send headers; ${clientDiag.known} carry a dialed socket, ${clientDiag.unknown} another`,
+  );
   console.log("stats:", agent.statsFor(origin));
 
   console.log("== Pool with a connection limit");

tests/fixtures/npm/cases/undici7/main.ts

@@ -115,6 +115,34 @@ channel("undici:request:trailers").subscribe((message) => {
   if (t !== undefined) console.log(`  [diag] ${t.label} done after ${t.dispatched} dispatch`);
 });
 
+// A transport's connection registry: sockets the wrapped connector dialed,
+// keyed weakly, and the client diagnostics messages that carry them.
+const dialed = new WeakMap<Socket, number>();
+const clientDiag = { beforeConnect: 0, connector: 0, connected: 0, sendHeaders: 0, known: 0, unknown: 0 };
+const noteSocket = (socket: Socket): void => {
+  const n = dialed.get(socket);
+  if (n === undefined) clientDiag.unknown++;
+  else clientDiag.known++;
+};
+channel("undici:client:beforeConnect").subscribe((message) => {
+  const { connector } = message as { connector: unknown };
+  clientDiag.beforeConnect++;
+  if (typeof connector === "function") clientDiag.connector++;
+});
+channel("undici:client:connected").subscribe((message) => {
+  const { socket, connectParams } = message as {
+    socket: Socket;
+    connectParams: { protocol: string; hostname: string };
+  };
+  clientDiag.connected++;
+  if (connectParams.protocol === "http:" && connectParams.hostname === "127.0.0.1") noteSocket(socket);
+});
+channel("undici:client:sendHeaders").subscribe((message) => {
+  const { socket, headers } = message as { socket: Socket; headers: string };
+  clientDiag.sendHeaders++;
+  if (headers.length > 0) noteSocket(socket);
+});
+
 /** A transport-style Agent: one Pool per origin built by the factory, a
  * wrapped connector, and a Client interceptor that marks dispatches. */
 class TracingAgent extends Agent {
@@ -136,7 +164,9 @@ class TracingAgent extends Agent {
           pipelining: 1,
           connect(options, callback) {
             opened++;
-            return connect(options, callback);
+            const socket = connect(options, callback);
+            dialed.set(socket, opened);
+            return socket;
           },
           interceptors: {
             Client: [
@@ -232,6 +262,9 @@ server.listen(0, "127.0.0.1", async () => {
     await r.body.text();
   }
   console.log("connections opened:", agent.connections);
+  console.log(
+    `client diagnostics: ${clientDiag.beforeConnect} before connect (${clientDiag.connector} with a connector), ${clientDiag.connected} connected, ${clientDiag.sendHeaders} send headers; ${clientDiag.known} carry a dialed socket, ${clientDiag.unknown} another`,
+  );
   console.log("stats:", agent.statsFor(origin));
 
   console.log("== Pool with a connection limit");