From be29823ea26033a3d5a4818d5393433e79a11fe7 Mon Sep 17 00:00:00 2001
From: theanarkh <theratliter@gmail.com>
Date: Wed, 16 Oct 2024 09:36:16 +0800
Subject: [PATCH] lib: add UV_UDP_REUSEPORT for udp

---
 doc/api/dgram.md                              | 16 +++++++--
 lib/dgram.js                                  |  7 +++-
 src/udp_wrap.cc                               |  1 +
 test/common/udp.js                            | 25 +++++++++++++
 .../test-child-process-dgram-reuseport.js     | 35 +++++++++++++++++++
 test/parallel/test-cluster-dgram-reuseport.js | 31 ++++++++++++++++
 test/parallel/test-dgram-reuseport.js         | 19 ++++++++++
 7 files changed, 131 insertions(+), 3 deletions(-)
 create mode 100644 test/common/udp.js
 create mode 100644 test/parallel/test-child-process-dgram-reuseport.js
 create mode 100644 test/parallel/test-cluster-dgram-reuseport.js
 create mode 100644 test/parallel/test-dgram-reuseport.js

diff --git a/doc/api/dgram.md b/doc/api/dgram.md
index 26e01f100258f0..f12ad599acfafd 100644
--- a/doc/api/dgram.md
+++ b/doc/api/dgram.md
@@ -343,7 +343,8 @@ used when using `dgram.Socket` objects with the [`cluster`][] module. When
 `exclusive` is set to `false` (the default), cluster workers will use the same
 underlying socket handle allowing connection handling duties to be shared.
 When `exclusive` is `true`, however, the handle is not shared and attempted
-port sharing results in an error.
+port sharing results in an error. If creating a `dgram.Socket` with `reusePort`
+flag, `exclusive` will always be `true` when call `socket.bind()` with a port.
 
 A bound datagram socket keeps the Node.js process running to receive
 datagram messages.
@@ -916,6 +917,9 @@ chained.
 <!-- YAML
 added: v0.11.13
 changes:
+  - version: REPLACEME
+    pr-url: https://github.com/nodejs/node/pull/55403
+    description: The `reusePort` option is supported.
   - version: v15.8.0
     pr-url: https://github.com/nodejs/node/pull/37026
     description: AbortSignal support was added.
@@ -935,7 +939,15 @@ changes:
   * `type` {string} The family of socket. Must be either `'udp4'` or `'udp6'`.
     Required.
   * `reuseAddr` {boolean} When `true` [`socket.bind()`][] will reuse the
-    address, even if another process has already bound a socket on it.
+    address, even if another process has already bound a socket on it, but
+    only one socket can receive the data.
+    **Default:** `false`.
+  * `reusePort` {boolean} When `true` [`socket.bind()`][] will reuse the
+    address, even if another process has already bound a socket on it. Incoming
+    datagrams are distributed across the receiving sockets. The flag is available
+    only on some platforms, such as Linux 3.9+, DragonFlyBSD 3.6+, FreeBSD 12.0+,
+    Solaris 11.4, and AIX 7.2.5+. On unsupported platforms this function will
+    return an error.
     **Default:** `false`.
   * `ipv6Only` {boolean} Setting `ipv6Only` to `true` will
     disable dual-stack support, i.e., binding to address `::` won't make
diff --git a/lib/dgram.js b/lib/dgram.js
index baf747196293f5..74dd83889990ca 100644
--- a/lib/dgram.js
+++ b/lib/dgram.js
@@ -74,7 +74,7 @@ const {
 const { UV_UDP_REUSEADDR } = internalBinding('constants').os;
 
 const {
-  constants: { UV_UDP_IPV6ONLY },
+  constants: { UV_UDP_IPV6ONLY, UV_UDP_REUSEPORT },
   UDP,
   SendWrap,
 } = internalBinding('udp_wrap');
@@ -130,6 +130,7 @@ function Socket(type, listener) {
     connectState: CONNECT_STATE_DISCONNECTED,
     queue: undefined,
     reuseAddr: options?.reuseAddr, // Use UV_UDP_REUSEADDR if true.
+    reusePort: options?.reusePort,
     ipv6Only: options?.ipv6Only,
     recvBufferSize,
     sendBufferSize,
@@ -345,6 +346,10 @@ Socket.prototype.bind = function(port_, address_ /* , callback */) {
       flags |= UV_UDP_REUSEADDR;
     if (state.ipv6Only)
       flags |= UV_UDP_IPV6ONLY;
+    if (state.reusePort) {
+      exclusive = true;
+      flags |= UV_UDP_REUSEPORT;
+    }
 
     if (cluster.isWorker && !exclusive) {
       bindServerHandle(this, {
diff --git a/src/udp_wrap.cc b/src/udp_wrap.cc
index f51c47afef39cd..1c6652ad3172b6 100644
--- a/src/udp_wrap.cc
+++ b/src/udp_wrap.cc
@@ -231,6 +231,7 @@ void UDPWrap::Initialize(Local<Object> target,
   Local<Object> constants = Object::New(isolate);
   NODE_DEFINE_CONSTANT(constants, UV_UDP_IPV6ONLY);
   NODE_DEFINE_CONSTANT(constants, UV_UDP_REUSEADDR);
+  NODE_DEFINE_CONSTANT(constants, UV_UDP_REUSEPORT);
   target->Set(context,
               env->constants_string(),
               constants).Check();
diff --git a/test/common/udp.js b/test/common/udp.js
new file mode 100644
index 00000000000000..fd920173994a19
--- /dev/null
+++ b/test/common/udp.js
@@ -0,0 +1,25 @@
+'use strict';
+const dgram = require('dgram');
+
+const options = { type: 'udp4', reuseAddr: true };
+
+function checkSupportReusePort() {
+  return new Promise((resolve, reject) => {
+    const socket = dgram.createSocket(options);
+    socket.bind(0);
+    socket.on('listening', () => {
+      socket.close();
+      resolve();
+    });
+    socket.on('error', (err) => {
+      console.log('don not support reusePort flag', err);
+      socket.close();
+      reject();
+    });
+  });
+}
+
+module.exports = {
+  checkSupportReusePort,
+  options,
+};
diff --git a/test/parallel/test-child-process-dgram-reuseport.js b/test/parallel/test-child-process-dgram-reuseport.js
new file mode 100644
index 00000000000000..bb148d420dea6e
--- /dev/null
+++ b/test/parallel/test-child-process-dgram-reuseport.js
@@ -0,0 +1,35 @@
+'use strict';
+const common = require('../common');
+const { checkSupportReusePort, options } = require('../common/udp');
+const assert = require('assert');
+const child_process = require('child_process');
+const dgram = require('dgram');
+
+if (!process.env.isWorker) {
+  checkSupportReusePort().then(() => {
+    const socket = dgram.createSocket(options);
+    socket.bind(0, common.mustCall(() => {
+      const port = socket.address().port;
+      const workerOptions = { env: { 'isWorker': 1, port } };
+      let count = 2;
+      for (let i = 0; i < 2; i++) {
+        const worker = child_process.fork(__filename, workerOptions);
+        worker.on('exit', common.mustCall((code) => {
+          assert.strictEqual(code, 0);
+          if (--count === 0) {
+            socket.close();
+          }
+        }));
+      }
+    }));
+  }, process.exit);
+  return;
+}
+
+const socket = dgram.createSocket(options);
+
+socket.bind(+process.env.port, common.mustCall(() => {
+  socket.close(common.mustCall(() => {
+    process.exit(0);
+  }));
+})).on('error', common.mustNotCall());
diff --git a/test/parallel/test-cluster-dgram-reuseport.js b/test/parallel/test-cluster-dgram-reuseport.js
new file mode 100644
index 00000000000000..ea439e2877c3b0
--- /dev/null
+++ b/test/parallel/test-cluster-dgram-reuseport.js
@@ -0,0 +1,31 @@
+'use strict';
+const common = require('../common');
+const { checkSupportReusePort, options } = require('../common/udp');
+const assert = require('assert');
+const cluster = require('cluster');
+const dgram = require('dgram');
+
+if (cluster.isPrimary) {
+  checkSupportReusePort().then(() => {
+    cluster.fork().on('exit', common.mustCall((code) => {
+      assert.strictEqual(code, 0);
+    }));
+  }, process.exit);
+  return;
+}
+
+let waiting = 2;
+function close() {
+  if (--waiting === 0)
+    cluster.worker.disconnect();
+}
+
+const socket1 = dgram.createSocket(options);
+const socket2 = dgram.createSocket(options);
+
+socket1.bind(0, () => {
+  socket2.bind(socket1.address().port, () => {
+    socket1.close(close);
+    socket2.close(close);
+  }).on('error', common.mustNotCall());
+}).on('error', common.mustNotCall());
diff --git a/test/parallel/test-dgram-reuseport.js b/test/parallel/test-dgram-reuseport.js
new file mode 100644
index 00000000000000..f1bb9d124c4c06
--- /dev/null
+++ b/test/parallel/test-dgram-reuseport.js
@@ -0,0 +1,19 @@
+'use strict';
+const common = require('../common');
+const { checkSupportReusePort, options } = require('../common/udp');
+const dgram = require('dgram');
+
+function test() {
+  const socket1 = dgram.createSocket(options);
+  const socket2 = dgram.createSocket(options);
+  socket1.bind(0, common.mustCall(() => {
+    socket2.bind(socket1.address().port, common.mustCall(() => {
+      socket1.close();
+      socket2.close();
+    }));
+  }));
+  socket1.on('error', common.mustNotCall());
+  socket2.on('error', common.mustNotCall());
+}
+
+checkSupportReusePort().then(test, process.exit);