From 443d4fcff3cdf05032291ec85b1d9561d3798218 Mon Sep 17 00:00:00 2001
From: Chenyu Yang <chenyu.yang@intel.com>
Date: Thu, 7 Dec 2023 04:05:38 +0800
Subject: [PATCH] benchmark: add undici websocket benchmark

Refs: https://github.com/nodejs/performance/issues/114
PR-URL: https://github.com/nodejs/node/pull/50586
Reviewed-By: Luigi Pinca <luigipinca@gmail.com>
Reviewed-By: Matthew Aitken <maitken033380023@gmail.com>
---
 benchmark/websocket/simple.js              | 100 +++++++++++++++++++++
 test/benchmark/test-benchmark-websocket.js |   9 ++
 2 files changed, 109 insertions(+)
 create mode 100644 benchmark/websocket/simple.js
 create mode 100644 test/benchmark/test-benchmark-websocket.js

diff --git a/benchmark/websocket/simple.js b/benchmark/websocket/simple.js
new file mode 100644
index 00000000000000..1147602a7b1ea7
--- /dev/null
+++ b/benchmark/websocket/simple.js
@@ -0,0 +1,100 @@
+'use strict';
+
+const common = require('../common.js');
+const crypto = require('crypto');
+const http = require('http');
+const { WebSocket } = require('../../deps/undici/undici');
+
+const GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11';
+
+const configs = {
+  size: [64, 16 * 1024, 128 * 1024, 1024 * 1024],
+  useBinary: ['true', 'false'],
+  roundtrips: [5000, 1000, 100],
+};
+
+const bench = common.createBenchmark(main, configs);
+
+function createFrame(data, opcode) {
+  let infoLength = 2;
+  let payloadLength = data.length;
+
+  if (payloadLength >= 65536) {
+    infoLength += 8;
+    payloadLength = 127;
+  } else if (payloadLength > 125) {
+    infoLength += 2;
+    payloadLength = 126;
+  }
+
+  const info = Buffer.alloc(infoLength);
+
+  info[0] = opcode | 0x80;
+  info[1] = payloadLength;
+
+  if (payloadLength === 126) {
+    info.writeUInt16BE(data.length, 2);
+  } else if (payloadLength === 127) {
+    info[2] = info[3] = 0;
+    info.writeUIntBE(data.length, 4, 6);
+  }
+
+  return Buffer.concat([info, data]);
+}
+
+function main(conf) {
+  const frame = createFrame(Buffer.alloc(conf.size).fill('.'), conf.useBinary === 'true' ? 2 : 1);
+  const server = http.createServer();
+  server.on('upgrade', (req, socket) => {
+    const key = crypto
+      .createHash('sha1')
+      .update(req.headers['sec-websocket-key'] + GUID)
+      .digest('base64');
+
+    let bytesReceived = 0;
+    let roundtrip = 0;
+
+    socket.on('data', function onData(chunk) {
+      bytesReceived += chunk.length;
+
+      if (bytesReceived === frame.length + 4) { // +4 for the mask.
+        // Message completely received.
+        bytesReceived = 0;
+
+        if (++roundtrip === conf.roundtrips) {
+          socket.removeListener('data', onData);
+          socket.resume();
+          socket.end();
+          server.close();
+
+          bench.end(conf.roundtrips);
+        } else {
+          socket.write(frame);
+        }
+      }
+    });
+
+    socket.write(
+      [
+        'HTTP/1.1 101 Switching Protocols',
+        'Upgrade: websocket',
+        'Connection: Upgrade',
+        `Sec-WebSocket-Accept: ${key}`,
+        '\r\n',
+      ].join('\r\n'),
+    );
+
+    socket.write(frame);
+  });
+
+  server.listen(0, () => {
+    const ws = new WebSocket(`ws://localhost:${server.address().port}`);
+    ws.addEventListener('open', () => {
+      bench.start();
+    });
+
+    ws.addEventListener('message', (event) => {
+      ws.send(event.data);
+    });
+  });
+}
diff --git a/test/benchmark/test-benchmark-websocket.js b/test/benchmark/test-benchmark-websocket.js
new file mode 100644
index 00000000000000..d02e46a76b4e9e
--- /dev/null
+++ b/test/benchmark/test-benchmark-websocket.js
@@ -0,0 +1,9 @@
+'use strict';
+
+const common = require('../common');
+if (!common.enoughTestMem)
+  common.skip('Insufficient memory for Websocket benchmark test');
+
+const runBenchmark = require('../common/benchmark');
+
+runBenchmark('websocket', { NODEJS_BENCHMARK_ZERO_ALLOWED: 1 });