node:dgram: run 'message' in the async context of createSocket()

src/js/node/dgram.ts

@@ -66,6 +66,9 @@ const SymbolDispose = Symbol.dispose;
 const SymbolAsyncDispose = Symbol.asyncDispose;
 const ObjectDefineProperty = Object.defineProperty;
 const FunctionPrototypeBind = Function.prototype.bind;
+const WeakMapPrototypeGet = WeakMap.prototype.get;
+const WeakMapPrototypeSet = WeakMap.prototype.set;
+const WeakMapPrototypeDelete = WeakMap.prototype.delete;
 
 // Mirrors Node's SystemError shape for ERR_SOCKET_BUFFER_SIZE: name
 // "SystemError", an enumerable `info` object, errno/syscall accessors that
@@ -309,6 +312,8 @@ interface QueuedSend {
   port: number | undefined;
   address: string | null | undefined;
   req: SendRequest;
+  // The async context of a send() that has a callback, where node makes the request that it completes: https://github.com/nodejs/node/blob/v26.3.0/lib/dgram.js#L738
+  frame: unknown;
 }
 
 interface UDPHandle {
@@ -327,6 +332,22 @@ interface UDPHandle {
   [kOwnerSymbol]?: $Node.DgramSocket;
 }
 
+// The async context of createSocket(), for a handle whose native socket is not open yet. node makes the UDP handle there, and what that handle delivers ('message', a receive error) runs in it: https://github.com/nodejs/node/blob/v26.3.0/lib/dgram.js#L143
+const handleFrames = new WeakMap<UDPHandle, object>();
+
+// Bun.udpSocket() takes the socket's Bun.ModuleGraph from the async context, so a socket that another graph binds opens in the context of bind(): it stays the binder's.
+function openInFrameOfHandle(frame: unknown, udpOptions: any) {
+  const AsyncContextFrame = require("internal/async_context_frame");
+  if (AsyncContextFrame.graphOf(frame) !== AsyncContextFrame.currentGraph()) return Bun.udpSocket(udpOptions);
+  // Not AsyncContextFrame.run(): handed `Bun` or `process` as a value, an ahead-of-time build keeps every member of it, and the apply() in run() hides the type of what the call returns.
+  const previous = AsyncContextFrame.exchange(frame);
+  try {
+    return Bun.udpSocket(udpOptions);
+  } finally {
+    AsyncContextFrame.exchange(previous);
+  }
+}
+
 function newHandle(type: string, lookup: LookupFunction | undefined) {
   if (lookup === undefined) {
     lookup = require("node:dns").lookup as LookupFunction;
@@ -363,6 +384,9 @@ function newHandle(type: string, lookup: LookupFunction | undefined) {
 
   handle.onmessage = onMessage;
 
+  const frame = $getInternalField($asyncContext, 0);
+  if (frame !== undefined) WeakMapPrototypeSet.$call(handleFrames, handle, frame);
+
   return handle;
 }
 
@@ -453,7 +477,14 @@ function handleQueueSend(
   // same way.
   const length = data.byteLength;
   handle.sendQueue ??= [];
-  handle.sendQueue.push({ data, length, port, address, req });
+  handle.sendQueue.push({
+    data,
+    length,
+    port,
+    address,
+    req,
+    frame: req.callback ? $getInternalField($asyncContext, 0) : undefined,
+  });
   handle.queueCount += 1;
   handle.queueSize += length;
   return 0;
@@ -509,7 +540,20 @@ function completeQueuedSend(handle: UDPHandle, entry: QueuedSend, err: unknown)
   // callback cannot skip a sibling entry or run inside close()/drain.
   const req = entry.req;
   if (req && typeof req.oncomplete === "function") {
-    process.nextTick(runQueuedSendComplete, req, err, err === undefined ? entry.length : 0);
+    const sent = err === undefined ? entry.length : 0;
+    const frame = entry.frame;
+    // Queued inside the frame of the send(), so that what the callback throws is reported in it too. The frames are equal in every program without AsyncLocalStorage and Bun.ModuleGraph, which never loads the helper.
+    if (!req.callback || frame === $getInternalField($asyncContext, 0))
+      process.nextTick(runQueuedSendComplete, req, err, sent);
+    else {
+      const AsyncContextFrame = require("internal/async_context_frame");
+      const previous = AsyncContextFrame.exchange(frame);
+      try {
+        process.nextTick(runQueuedSendComplete, req, err, sent);
+      } finally {
+        AsyncContextFrame.exchange(previous);
+      }
+    }
   }
 }
 
@@ -911,8 +955,17 @@ function startBunSocket(
     };
     // Private name: a cluster-shared descriptor is read one datagram at a time so workers share the load.
     if (sharedHandle) $putByIdDirectPrivate(udpOptions, "sharedFd", true);
-    (Bun.udpSocket(udpOptions) as Promise<$Node.UDPSocketHandle>).$then(
+    // Not for a shared bind: its handle is the one the cluster primary sent (node's replaceHandle(): https://github.com/nodejs/node/blob/v26.3.0/lib/dgram.js#L213).
+    const handle = state.handle!;
+    const frame = WeakMapPrototypeGet.$call(handleFrames, handle);
+    const opened =
+      sharedHandle !== undefined || frame === $getInternalField($asyncContext, 0)
+        ? Bun.udpSocket(udpOptions)
+        : openInFrameOfHandle(frame, udpOptions);
+    (opened as Promise<$Node.UDPSocketHandle>).$then(
       (socket: $Node.UDPSocketHandle) => {
+        // The native socket holds the frame now. Until a bind succeeds, the handle keeps it.
+        WeakMapPrototypeDelete.$call(handleFrames, handle);
         if (!state.handle) {
           // Closed while the bind was in flight.
           socket.close();

test/js/bun/module-graph/module-graph-gc.test.ts

@@ -5,6 +5,8 @@ import { afterAll, describe, expect, jest, test } from "bun:test";
 import { rmSync, writeFileSync } from "fs";
 import { bunEnv, bunExe, tempDir, tls } from "harness";
 import { AsyncLocalStorage } from "node:async_hooks";
+import dgram from "node:dgram";
+import { once } from "node:events";
 import { join } from "path";
 import { collect, heapFromTimer, Lifetimes } from "./module-graph-lifetimes";
 
@@ -39,8 +41,14 @@ const dir = String(
     // What the graph opens reports to, and is stopped through, `control` (a host object in
     // `globals`), so the host can stop it without holding anything of the graph.
     "io.mjs": `
+      import dgram from "node:dgram";
       import http from "node:http";
       import https from "node:https";
+      // A node:dgram socket that nothing has bound yet.
+      export function dgramSocket() { return dgram.createSocket("udp4"); }
+      // Calls the listeners of the host's target while the graph's script is on the stack.
+      export function dispatch(target) { target.dispatchEvent(new Event("call")); }
+      export function dgramSend(socket) { socket.send("from the graph"); }
       export function serve() {
         const server = Bun.serve({ port: 0, fetch: () => new Response("graph " + TAG) });
         control.httpPort = server.port;
@@ -398,6 +406,89 @@ describe("ModuleGraph GC: what the graph's context owns", () => {
     expect(await accepts(state.httpPort)).toBe(false);
   });
 
+  // node:dgram opens the socket in bind(), inside the async context of createSocket(). A socket that the host bound
+  // does not keep the graph that made it, and not a graph whose script was on the stack when the host bound it. A
+  // socket the graph bound is the graph's: its handlers hold the graph until the host closes the socket.
+  test("of the node:dgram sockets the host holds, one that the graph bound keeps the graph and one that the host bound does not", async () => {
+    const lifetimes = new Lifetimes();
+    const storage = new AsyncLocalStorage<string>();
+    const sockets: Record<string, dgram.Socket> = {};
+    const bind = (socket: dgram.Socket) =>
+      new Promise<string>(resolve => {
+        socket.once("error", (error: NodeJS.ErrnoException) => resolve(String(error.code)));
+        socket.bind(0, "127.0.0.1", () => resolve("listening"));
+      });
+    const rows = [
+      "made in the graph, bound by the host",
+      "made in the graph, bound by the graph",
+      "made by the host before, bound by the host in a listener the graph's script calls",
+    ];
+    try {
+      await (async () => {
+        for (const row of rows) {
+          const graph = lifetimes.track(row, new ModuleGraph({ globals: { TAG: row, control: control() } }));
+          const io = await graph.import(file("io.mjs"));
+          if (row.startsWith("made in the graph")) {
+            const socket = (sockets[row] = graph.run(() => storage.run("createSocket()", () => io.dgramSocket())));
+            const bound = row.endsWith("bound by the host") ? bind(socket) : graph.run(() => bind(socket));
+            expect(await bound).toBe("listening");
+          } else {
+            // A listener of the host's that the graph's script calls runs as the host, in an async context that holds the graph.
+            const socket = (sockets[row] = dgram.createSocket("udp4"));
+            const target = new EventTarget();
+            let ranAs: string | undefined;
+            let bound: Promise<string> | undefined;
+            target.addEventListener("call", () => {
+              ranAs = ModuleGraph.current === undefined ? "the host" : "a graph";
+              bound = storage.run("bind()", () => bind(socket));
+            });
+            graph.run(() => io.dispatch(target));
+            expect(ranAs).toBe("the host");
+            expect(await bound).toBe("listening");
+          }
+          graph.dispose();
+        }
+      })();
+      const boundByTheGraph = "made in the graph, bound by the graph";
+      expect(await lifetimes.stillAlive(...rows.filter(row => row !== boundByTheGraph))).toEqual([]);
+      expect(await lifetimes.survives(boundByTheGraph)).toBe(true);
+    } finally {
+      for (const socket of Object.values(sockets)) socket.close();
+    }
+    expect(await lifetimes.stillAlive(...rows)).toEqual([]);
+  });
+
+  // A send() that waits behind backpressure keeps the async context of its call for its callback. With no callback
+  // it keeps nothing: the queue of a socket of the host's does not keep the graph whose script sent.
+  test("a queued node:dgram send() with no callback does not keep the graph that sent it", async () => {
+    const { kStateSymbol } = require("bun:internal-for-testing").exposedInternals["internal/dgram"];
+    const lifetimes = new Lifetimes();
+    const socket = dgram.createSocket("udp4");
+    socket.bind(0, "127.0.0.1");
+    await once(socket, "listening");
+    socket.connect(socket.address().port, "127.0.0.1");
+    await once(socket, "connect");
+    const handle = socket[kStateSymbol].handle;
+    const realSocket = handle.socket;
+    const refusesEverySend = { send: () => false };
+    try {
+      await (async () => {
+        const graph = lifetimes.track("graph", new ModuleGraph({ globals: { TAG: "sends", control: control() } }));
+        const io = await graph.import(file("io.mjs"));
+        // With a queue, every send() waits behind it, and a drain that is refused leaves it there: backpressure.
+        handle.socket = refusesEverySend;
+        handle.sendQueue = [];
+        graph.run(() => io.dgramSend(socket));
+        graph.dispose();
+      })();
+      expect(await lifetimes.stillAlive("graph")).toEqual([]);
+      expect(handle.sendQueue).toHaveLength(1);
+    } finally {
+      handle.socket = realSocket;
+      socket.close();
+    }
+  });
+
   test("a repeating timer keeps its graph alive; clearing it lets the graph go", async () => {
     const lifetimes = new Lifetimes();
     const state = control();