websocket: 200x faster generate mask (#3204)

tsctx · Uzlopak · web-flow · commit c0dc3dd37a41 · 2024-05-07T19:27:54.000+09:00
* websocket: improve performance of generate mask

* apply comments

* Apply suggestions from code review

* Apply suggestions from code review

* Apply suggestions from code review

Co-authored-by: Aras Abbasi &lt;aras.abbasi@googlemail.com&gt;

* Apply suggestions from code review

* Update lib/web/websocket/frame.js

* Update lib/web/websocket/frame.js

* Update lib/web/websocket/frame.js

Co-authored-by: Aras Abbasi &lt;aras.abbasi@googlemail.com&gt;

---------

Co-authored-by: Aras Abbasi &lt;aras.abbasi@googlemail.com&gt;
diff --git a/benchmarks/websocket/generate-mask.mjs b/benchmarks/websocket/generate-mask.mjs
@@ -0,0 +1,22 @@
+import { randomFillSync, randomBytes } from 'node:crypto'
+import { bench, group, run } from 'mitata'
+
+const BUFFER_SIZE = 16384
+
+const buf = Buffer.allocUnsafe(BUFFER_SIZE)
+let bufIdx = BUFFER_SIZE
+
+function generateMask () {
+  if (bufIdx === BUFFER_SIZE) {
+    bufIdx = 0
+    randomFillSync(buf, 0, BUFFER_SIZE)
+  }
+  return [buf[bufIdx++], buf[bufIdx++], buf[bufIdx++], buf[bufIdx++]]
+}
+
+group('generate', () => {
+  bench('generateMask', () => generateMask())
+  bench('crypto.randomBytes(4)', () => randomBytes(4))
+})
+
+await run()
diff --git a/lib/web/websocket/frame.js b/lib/web/websocket/frame.js
@@ -2,13 +2,34 @@
 
 const { maxUnsigned16Bit } = require('./constants')
 
+const BUFFER_SIZE = 16386
+
 /** @type {import('crypto')} */
 let crypto
+let buffer = null
+let bufIdx = BUFFER_SIZE
+
 try {
   crypto = require('node:crypto')
 /* c8 ignore next 3 */
 } catch {
+  crypto = {
+    // not full compatibility, but minimum.
+    randomFillSync: function randomFillSync (buffer, _offset, _size) {
+      for (let i = 0; i < buffer.length; ++i) {
+        buffer[i] = Math.random() * 255 | 0
+      }
+      return buffer
+    }
+  }
+}
 
+function generateMask () {
+  if (bufIdx === BUFFER_SIZE) {
+    bufIdx = 0
+    crypto.randomFillSync((buffer ??= Buffer.allocUnsafe(BUFFER_SIZE)), 0, BUFFER_SIZE)
+  }
+  return [buffer[bufIdx++], buffer[bufIdx++], buffer[bufIdx++], buffer[bufIdx++]]
 }
 
 class WebsocketFrameSend {
@@ -17,11 +38,12 @@ class WebsocketFrameSend {
    */
   constructor (data) {
     this.frameData = data
-    this.maskKey = crypto.randomBytes(4)
   }
 
   createFrame (opcode) {
-    const bodyLength = this.frameData?.byteLength ?? 0
+    const frameData = this.frameData
+    const maskKey = generateMask()
+    const bodyLength = frameData?.byteLength ?? 0
 
     /** @type {number} */
     let payloadLength = bodyLength // 0-125
@@ -43,10 +65,10 @@ class WebsocketFrameSend {
     buffer[0] = (buffer[0] & 0xF0) + opcode // opcode
 
     /*! ws. MIT License. Einar Otto Stangvik <einaros@gmail.com> */
-    buffer[offset - 4] = this.maskKey[0]
-    buffer[offset - 3] = this.maskKey[1]
-    buffer[offset - 2] = this.maskKey[2]
-    buffer[offset - 1] = this.maskKey[3]
+    buffer[offset - 4] = maskKey[0]
+    buffer[offset - 3] = maskKey[1]
+    buffer[offset - 2] = maskKey[2]
+    buffer[offset - 1] = maskKey[3]
 
     buffer[1] = payloadLength
 
@@ -61,8 +83,8 @@ class WebsocketFrameSend {
     buffer[1] |= 0x80 // MASK
 
     // mask body
-    for (let i = 0; i < bodyLength; i++) {
-      buffer[offset + i] = this.frameData[i] ^ this.maskKey[i % 4]
+    for (let i = 0; i < bodyLength; ++i) {
+      buffer[offset + i] = frameData[i] ^ maskKey[i & 3]
     }
 
     return buffer
