websocket: add fast-path for string input (#3395)

nodejs · Jul 8, 2024 · 98dae4e · 98dae4e
1 parent 7436fee
commit 98dae4e
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 5 deletions.
diff --git a/lib/web/websocket/frame.js b/lib/web/websocket/frame.js
@@ -1,6 +1,6 @@
 'use strict'
 
-const { maxUnsigned16Bit } = require('./constants')
+const { maxUnsigned16Bit, opcodes } = require('./constants')
 
 const BUFFER_SIZE = 8 * 1024
 
@@ -89,6 +89,48 @@ class WebsocketFrameSend {
 
     return buffer
   }
+
+  /**
+   * @param {Uint8Array} buffer
+   */
+  static createFastTextFrame (buffer) {
+    const maskKey = generateMask()
+
+    const bodyLength = buffer.length
+
+    // mask body
+    for (let i = 0; i < bodyLength; ++i) {
+      buffer[i] ^= maskKey[i & 3]
+    }
+
+    let payloadLength = bodyLength
+    let offset = 6
+
+    if (bodyLength > maxUnsigned16Bit) {
+      offset += 8 // payload length is next 8 bytes
+      payloadLength = 127
+    } else if (bodyLength > 125) {
+      offset += 2 // payload length is next 2 bytes
+      payloadLength = 126
+    }
+    const head = Buffer.allocUnsafeSlow(offset)
+
+    head[0] = 0x80 /* FIN */ | opcodes.TEXT /* opcode TEXT */
+    head[1] = payloadLength | 0x80 /* MASK */
+    head[offset - 4] = maskKey[0]
+    head[offset - 3] = maskKey[1]
+    head[offset - 2] = maskKey[2]
+    head[offset - 1] = maskKey[3]
+
+    if (payloadLength === 126) {
+      head.writeUInt16BE(bodyLength, 2)
+    } else if (payloadLength === 127) {
+      head[2] = head[3] = 0
+      head.writeUIntBE(bodyLength, 4, 6)
+    }
+
+    return [head, buffer]
+  }
 }
 
 module.exports = {

diff --git a/lib/web/websocket/sender.js b/lib/web/websocket/sender.js
@@ -31,16 +31,25 @@ class SendQueue {
 
   add (item, cb, hint) {
     if (hint !== sendHints.blob) {
-      const frame = createFrame(item, hint)
       if (!this.#running) {
-        // fast-path
-        this.#socket.write(frame, cb)
+        // TODO(@tsctx): support fast-path for string on running
+        if (hint === sendHints.text) {
+          // special fast-path for string
+          const { 0: head, 1: body } = WebsocketFrameSend.createFastTextFrame(item)
+          this.#socket.cork()
+          this.#socket.write(head)
+          this.#socket.write(body, cb)
+          this.#socket.uncork()
+        } else {
+          // direct writing
+          this.#socket.write(createFrame(item, hint), cb)
+        }
       } else {
         /** @type {SendQueueNode} */
         const node = {
           promise: null,
           callback: cb,
-          frame
+          frame: createFrame(item, hint)
         }
         this.#queue.push(node)
       }