dgram: support AbortSignal in createSocket

Nitzan Uziely · benjamingr · commit 21c8c7e5114c · 2021-01-25T20:36:00.000+02:00
PR-URL: #37026 Reviewed-By: Benjamin Gruenbaum <benjamingr@gmail.com> Reviewed-By: Rich Trott <rtrott@gmail.com>
diff --git a/doc/api/dgram.md b/doc/api/dgram.md
@@ -735,6 +735,9 @@ chained.
 <!-- YAML
 added: v0.11.13
 changes:
+  - version: REPLACEME
+    pr-url: https://github.com/nodejs/node/pull/37026
+    description: AbortSignal support was added.
   - version: v11.4.0
     pr-url: https://github.com/nodejs/node/pull/23798
     description: The `ipv6Only` option is supported.
@@ -759,6 +762,7 @@ changes:
   * `recvBufferSize` {number} Sets the `SO_RCVBUF` socket value.
   * `sendBufferSize` {number} Sets the `SO_SNDBUF` socket value.
   * `lookup` {Function} Custom lookup function. **Default:** [`dns.lookup()`][].
+  * `signal` {AbortSignal} An AbortSignal that may be used to close a socket.
 * `callback` {Function} Attached as a listener for `'message'` events. Optional.
 * Returns: {dgram.Socket}
 
@@ -770,6 +774,20 @@ method will bind the socket to the "all interfaces" address on a random port
 and port can be retrieved using [`socket.address().address`][] and
 [`socket.address().port`][].
 
+If the `signal` option is enabled, calling `.abort()` on the corresponding
+`AbortController` is similar to calling `.close()` on the socket:
+
+```js
+const controller = new AbortController();
+const { signal } = controller;
+const server = dgram.createSocket({ type: 'udp4', signal });
+server.on('message', (msg, rinfo) => {
+  console.log(`server got: ${msg} from ${rinfo.address}:${rinfo.port}`);
+});
+// Later, when you want to close the server.
+controller.abort();
+```
+
 ### `dgram.createSocket(type[, callback])`
 <!-- YAML
 added: v0.1.99
diff --git a/lib/dgram.js b/lib/dgram.js
@@ -52,6 +52,7 @@ const {
 } = errors.codes;
 const {
   isInt32,
+  validateAbortSignal,
   validateString,
   validateNumber,
   validatePort,
@@ -125,6 +126,20 @@ function Socket(type, listener) {
     recvBufferSize,
     sendBufferSize
   };
+
+  if (options?.signal !== undefined) {
+    const { signal } = options;
+    validateAbortSignal(signal, 'options.signal');
+    const onAborted = () => {
+      this.close();
+    };
+    if (signal.aborted) {
+      onAborted();
+    } else {
+      signal.addEventListener('abort', onAborted);
+      this.once('close', () => signal.removeEventListener('abort', onAborted));
+    }
+  }
 }
 ObjectSetPrototypeOf(Socket.prototype, EventEmitter.prototype);
 ObjectSetPrototypeOf(Socket, EventEmitter);
diff --git a/test/parallel/test-dgram-close-signal.js b/test/parallel/test-dgram-close-signal.js
@@ -0,0 +1,33 @@
+'use strict';
+
+const common = require('../common');
+const assert = require('assert');
+const dgram = require('dgram');
+
+{
+  // Test bad signal.
+  assert.throws(
+    () => dgram.createSocket({ type: 'udp4', signal: {} }),
+    {
+      code: 'ERR_INVALID_ARG_TYPE',
+      name: 'TypeError'
+    });
+}
+
+{
+  // Test close.
+  const controller = new AbortController();
+  const { signal } = controller;
+  const server = dgram.createSocket({ type: 'udp4', signal });
+  server.on('close', common.mustCall());
+  controller.abort();
+}
+
+{
+  // Test close with pre-aborted signal.
+  const controller = new AbortController();
+  controller.abort();
+  const { signal } = controller;
+  const server = dgram.createSocket({ type: 'udp4', signal });
+  server.on('close', common.mustCall());
+}
