Repository: websockets/ws
Pull request #2000 (2022-01-06 16:06:35)

Reduce buffer allocations
Do not convert strings to `Buffer`s if data does not need to be masked.

Refs: https://github.com/websockets/ws/pull/1998

lpinca: @ronag, this is what I had in mind. Anyway

1. The API is convoluted
1. It is harder to read and maintain
1. `Buffer.concat(Sender.frame(data, { opcode: 1 }))` no longer works if `data` is a string.

so I'm not ok with merging if the gain is marginal.
lpinca: Current implementation addresses my two main concerns:

1. The API is now simpler.
1. The `data` argument of the `Sender.frame()` function can now be a string but the function always returns a list of `Buffer`s when used externally. Behavior changes when an internal only option(`options[kByteLength]`) is used.

@ronag  what do you think?

Patch:
diff --git a/README.md b/README.md
index 82ca8db4..cab335b7 100644
--- a/README.md
+++ b/README.md
@@ -33,7 +33,7 @@ can use one of the many wrappers available on npm, like
   - [Multiple servers sharing a single HTTP/S server](#multiple-servers-sharing-a-single-https-server)
   - [Client authentication](#client-authentication)
   - [Server broadcast](#server-broadcast)
-  - [echo.websocket.org demo](#echowebsocketorg-demo)
+  - [Round-trip time](#round-trip-time)
   - [Use the Node.js streams API](#use-the-nodejs-streams-api)
   - [Other examples](#other-examples)
 - [FAQ](#faq)
@@ -324,14 +324,12 @@ wss.on('connection', function connection(ws) {
 });
 ```
 
-### echo.websocket.org demo
+### Round-trip time
 
 ```js
 import WebSocket from 'ws';
 
-const ws = new WebSocket('wss://echo.websocket.org/', {
-  origin: 'https://websocket.org'
-});
+const ws = new WebSocket('wss://websocket-echo.com/');
 
 ws.on('open', function open() {
   console.log('connected');
@@ -343,7 +341,7 @@ ws.on('close', function close() {
 });
 
 ws.on('message', function message(data) {
-  console.log(`Roundtrip time: ${Date.now() - data} ms`);
+  console.log(`Round-trip time: ${Date.now() - data} ms`);
 
   setTimeout(function timeout() {
     ws.send(Date.now());
@@ -356,9 +354,7 @@ ws.on('message', function message(data) {
 ```js
 import WebSocket, { createWebSocketStream } from 'ws';
 
-const ws = new WebSocket('wss://echo.websocket.org/', {
-  origin: 'https://websocket.org'
-});
+const ws = new WebSocket('wss://websocket-echo.com/');
 
 const duplex = createWebSocketStream(ws, { encoding: 'utf8' });
 
@@ -457,7 +453,7 @@ function heartbeat() {
   }, 30000 + 1000);
 }
 
-const client = new WebSocket('wss://echo.websocket.org/');
+const client = new WebSocket('wss://websocket-echo.com/');
 
 client.on('open', heartbeat);
 client.on('ping', heartbeat);
diff --git a/lib/permessage-deflate.js b/lib/permessage-deflate.js
index 50406971..94603c98 100644
--- a/lib/permessage-deflate.js
+++ b/lib/permessage-deflate.js
@@ -313,7 +313,7 @@ class PerMessageDeflate {
   /**
    * Compress data. Concurrency limited.
    *
-   * @param {Buffer} data Data to compress
+   * @param {(Buffer|String)} data Data to compress
    * @param {Boolean} fin Specifies whether or not this is the last fragment
    * @param {Function} callback Callback
    * @public
@@ -395,7 +395,7 @@ class PerMessageDeflate {
   /**
    * Compress data.
    *
-   * @param {Buffer} data Data to compress
+   * @param {(Buffer|String)} data Data to compress
    * @param {Boolean} fin Specifies whether or not this is the last fragment
    * @param {Function} callback Callback
    * @private
diff --git a/lib/sender.js b/lib/sender.js
index 2417656d..d331f653 100644
--- a/lib/sender.js
+++ b/lib/sender.js
@@ -11,6 +11,7 @@ const { EMPTY_BUFFER } = require('./constants');
 const { isValidStatusCode } = require('./validation');
 const { mask: applyMask, toBuffer } = require('./buffer-util');
 
+const kByteLength = Symbol('kByteLength');
 const maskBuffer = Buffer.alloc(4);
 
 /**
@@ -46,7 +47,7 @@ class Sender {
   /**
    * Frames a piece of data according to the HyBi WebSocket protocol.
    *
-   * @param {Buffer} data The data to frame
+   * @param {(Buffer|String)} data The data to frame
    * @param {Object} options Options object
    * @param {Boolean} [options.fin=false] Specifies whether or not to set the
    *     FIN bit
@@ -61,7 +62,7 @@ class Sender {
    *     modified
    * @param {Boolean} [options.rsv1=false] Specifies whether or not to set the
    *     RSV1 bit
-   * @return {Buffer[]} The framed data as a list of `Buffer` instances
+   * @return {(Buffer|String)[]} The framed data
    * @public
    */
   static frame(data, options) {
@@ -80,12 +81,27 @@ class Sender {
       }
 
       skipMasking = (mask[0] | mask[1] | mask[2] | mask[3]) === 0;
-      if (options.readOnly && !skipMasking) merge = true;
-
       offset = 6;
     }
 
-    let payloadLength = data.length;
+    let dataLength;
+
+    if (typeof data === 'string') {
+      if (
+        (!options.mask || skipMasking) &&
+        options[kByteLength] !== undefined
+      ) {
+        dataLength = options[kByteLength];
+      } else {
+        data = Buffer.from(data);
+        dataLength = data.length;
+      }
+    } else {
+      dataLength = data.length;
+      merge = options.mask && options.readOnly && !skipMasking;
+    }
+
+    let payloadLength = dataLength;
 
     if (data.length >= 65536) {
       offset += 8;
@@ -95,7 +111,7 @@ class Sender {
       payloadLength = 126;
     }
 
-    const target = Buffer.allocUnsafe(merge ? data.length + offset : offset);
+    const target = Buffer.allocUnsafe(merge ? dataLength + offset : offset);
 
     target[0] = options.fin ? options.opcode | 0x80 : options.opcode;
     if (options.rsv1) target[0] |= 0x40;
@@ -103,10 +119,10 @@ class Sender {
     target[1] = payloadLength;
 
     if (payloadLength === 126) {
-      target.writeUInt16BE(data.length, 2);
+      target.writeUInt16BE(dataLength, 2);
     } else if (payloadLength === 127) {
       target[2] = target[3] = 0;
-      target.writeUIntBE(data.length, 4, 6);
+      target.writeUIntBE(dataLength, 4, 6);
     }
 
     if (!options.mask) return [target, data];
@@ -164,36 +180,24 @@ class Sender {
       }
     }
 
+    const options = {
+      [kByteLength]: buf.length,
+      fin: true,
+      generateMask: this._generateMask,
+      mask,
+      maskBuffer: this._maskBuffer,
+      opcode: 0x08,
+      readOnly: false,
+      rsv1: false
+    };
+
     if (this._deflating) {
-      this.enqueue([this.doClose, buf, mask, cb]);
+      this.enqueue([this.dispatch, buf, false, options, cb]);
     } else {
-      this.doClose(buf, mask, cb);
+      this.sendFrame(Sender.frame(buf, options), cb);
     }
   }
 
-  /**
-   * Frames and sends a close message.
-   *
-   * @param {Buffer} data The message to send
-   * @param {Boolean} [mask=false] Specifies whether or not to mask `data`
-   * @param {Function} [cb] Callback
-   * @private
-   */
-  doClose(data, mask, cb) {
-    this.sendFrame(
-      Sender.frame(data, {
-        fin: true,
-        rsv1: false,
-        opcode: 0x08,
-        mask,
-        maskBuffer: this._maskBuffer,
-        generateMask: this._generateMask,
-        readOnly: false
-      }),
-      cb
-    );
-  }
-
   /**
    * Sends a ping message to the other peer.
    *
@@ -203,43 +207,40 @@ class Sender {
    * @public
    */
   ping(data, mask, cb) {
-    const buf = toBuffer(data);
+    let byteLength;
+    let readOnly;
 
-    if (buf.length > 125) {
+    if (typeof data === 'string') {
+      byteLength = Buffer.byteLength(data);
+      readOnly = false;
+    } else {
+      data = toBuffer(data);
+      byteLength = data.length;
+      readOnly = toBuffer.readOnly;
+    }
+
+    if (byteLength > 125) {
       throw new RangeError('The data size must not be greater than 125 bytes');
     }
 
+    const options = {
+      [kByteLength]: byteLength,
+      fin: true,
+      generateMask: this._generateMask,
+      mask,
+      maskBuffer: this._maskBuffer,
+      opcode: 0x09,
+      readOnly,
+      rsv1: false
+    };
+
     if (this._deflating) {
-      this.enqueue([this.doPing, buf, mask, toBuffer.readOnly, cb]);
+      this.enqueue([this.dispatch, data, false, options, cb]);
     } else {
-      this.doPing(buf, mask, toBuffer.readOnly, cb);
+      this.sendFrame(Sender.frame(data, options), cb);
     }
   }
 
-  /**
-   * Frames and sends a ping message.
-   *
-   * @param {Buffer} data The message to send
-   * @param {Boolean} [mask=false] Specifies whether or not to mask `data`
-   * @param {Boolean} [readOnly=false] Specifies whether `data` can be modified
-   * @param {Function} [cb] Callback
-   * @private
-   */
-  doPing(data, mask, readOnly, cb) {
-    this.sendFrame(
-      Sender.frame(data, {
-        fin: true,
-        rsv1: false,
-        opcode: 0x09,
-        mask,
-        maskBuffer: this._maskBuffer,
-        generateMask: this._generateMask,
-        readOnly
-      }),
-      cb
-    );
-  }
-
   /**
    * Sends a pong message to the other peer.
    *
@@ -249,43 +250,40 @@ class Sender {
    * @public
    */
   pong(data, mask, cb) {
-    const buf = toBuffer(data);
+    let byteLength;
+    let readOnly;
+
+    if (typeof data === 'string') {
+      byteLength = Buffer.byteLength(data);
+      readOnly = false;
+    } else {
+      data = toBuffer(data);
+      byteLength = data.length;
+      readOnly = toBuffer.readOnly;
+    }
 
-    if (buf.length > 125) {
+    if (byteLength > 125) {
       throw new RangeError('The data size must not be greater than 125 bytes');
     }
 
+    const options = {
+      [kByteLength]: byteLength,
+      fin: true,
+      generateMask: this._generateMask,
+      mask,
+      maskBuffer: this._maskBuffer,
+      opcode: 0x0a,
+      readOnly,
+      rsv1: false
+    };
+
     if (this._deflating) {
-      this.enqueue([this.doPong, buf, mask, toBuffer.readOnly, cb]);
+      this.enqueue([this.dispatch, data, false, options, cb]);
     } else {
-      this.doPong(buf, mask, toBuffer.readOnly, cb);
+      this.sendFrame(Sender.frame(data, options), cb);
     }
   }
 
-  /**
-   * Frames and sends a pong message.
-   *
-   * @param {Buffer} data The message to send
-   * @param {Boolean} [mask=false] Specifies whether or not to mask `data`
-   * @param {Boolean} [readOnly=false] Specifies whether `data` can be modified
-   * @param {Function} [cb] Callback
-   * @private
-   */
-  doPong(data, mask, readOnly, cb) {
-    this.sendFrame(
-      Sender.frame(data, {
-        fin: true,
-        rsv1: false,
-        opcode: 0x0a,
-        mask,
-        maskBuffer: this._maskBuffer,
-        generateMask: this._generateMask,
-        readOnly
-      }),
-      cb
-    );
-  }
-
   /**
    * Sends a data message to the other peer.
    *
@@ -303,11 +301,22 @@ class Sender {
    * @public
    */
   send(data, options, cb) {
-    const buf = toBuffer(data);
     const perMessageDeflate = this._extensions[PerMessageDeflate.extensionName];
     let opcode = options.binary ? 2 : 1;
     let rsv1 = options.compress;
 
+    let byteLength;
+    let readOnly;
+
+    if (typeof data === 'string') {
+      byteLength = Buffer.byteLength(data);
+      readOnly = false;
+    } else {
+      data = toBuffer(data);
+      byteLength = data.length;
+      readOnly = toBuffer.readOnly;
+    }
+
     if (this._firstFragment) {
       this._firstFragment = false;
       if (
@@ -319,7 +328,7 @@ class Sender {
             : 'client_no_context_takeover'
         ]
       ) {
-        rsv1 = buf.length >= perMessageDeflate._threshold;
+        rsv1 = byteLength >= perMessageDeflate._threshold;
       }
       this._compress = rsv1;
     } else {
@@ -331,30 +340,32 @@ class Sender {
 
     if (perMessageDeflate) {
       const opts = {
+        [kByteLength]: byteLength,
         fin: options.fin,
-        rsv1,
-        opcode,
+        generateMask: this._generateMask,
         mask: options.mask,
         maskBuffer: this._maskBuffer,
-        generateMask: this._generateMask,
-        readOnly: toBuffer.readOnly
+        opcode,
+        readOnly,
+        rsv1
       };
 
       if (this._deflating) {
-        this.enqueue([this.dispatch, buf, this._compress, opts, cb]);
+        this.enqueue([this.dispatch, data, this._compress, opts, cb]);
       } else {
-        this.dispatch(buf, this._compress, opts, cb);
+        this.dispatch(data, this._compress, opts, cb);
       }
     } else {
       this.sendFrame(
-        Sender.frame(buf, {
+        Sender.frame(data, {
+          [kByteLength]: byteLength,
           fin: options.fin,
-          rsv1: false,
-          opcode,
+          generateMask: this._generateMask,
           mask: options.mask,
           maskBuffer: this._maskBuffer,
-          generateMask: this._generateMask,
-          readOnly: toBuffer.readOnly
+          opcode,
+          readOnly,
+          rsv1: false
         }),
         cb
       );
@@ -362,13 +373,12 @@ class Sender {
   }
 
   /**
-   * Dispatches a data message.
+   * Dispatches a message.
    *
-   * @param {Buffer} data The message to send
+   * @param {(Buffer|String)} data The message to send
    * @param {Boolean} [compress=false] Specifies whether or not to compress
    *     `data`
    * @param {Object} options Options object
-   * @param {Number} options.opcode The opcode
    * @param {Boolean} [options.fin=false] Specifies whether or not to set the
    *     FIN bit
    * @param {Function} [options.generateMask] The function used to generate the
@@ -377,6 +387,7 @@ class Sender {
    *     `data`
    * @param {Buffer} [options.maskBuffer] The buffer used to store the masking
    *     key
+   * @param {Number} options.opcode The opcode
    * @param {Boolean} [options.readOnly=false] Specifies whether `data` can be
    *     modified
    * @param {Boolean} [options.rsv1=false] Specifies whether or not to set the
@@ -392,7 +403,7 @@ class Sender {
 
     const perMessageDeflate = this._extensions[PerMessageDeflate.extensionName];
 
-    this._bufferedBytes += data.length;
+    this._bufferedBytes += options[kByteLength];
     this._deflating = true;
     perMessageDeflate.compress(data, options.fin, (_, buf) => {
       if (this._socket.destroyed) {
@@ -403,7 +414,8 @@ class Sender {
         if (typeof cb === 'function') cb(err);
 
         for (let i = 0; i < this._queue.length; i++) {
-          const callback = this._queue[i][4];
+          const params = this._queue[i];
+          const callback = params[params.length - 1];
 
           if (typeof callback === 'function') callback(err);
         }
@@ -411,7 +423,7 @@ class Sender {
         return;
       }
 
-      this._bufferedBytes -= data.length;
+      this._bufferedBytes -= options[kByteLength];
       this._deflating = false;
       options.readOnly = false;
       this.sendFrame(Sender.frame(buf, options), cb);
@@ -428,7 +440,7 @@ class Sender {
     while (!this._deflating && this._queue.length) {
       const params = this._queue.shift();
 
-      this._bufferedBytes -= params[1].length;
+      this._bufferedBytes -= params[3][kByteLength];
       Reflect.apply(params[0], this, params.slice(1));
     }
   }
@@ -440,7 +452,7 @@ class Sender {
    * @private
    */
   enqueue(params) {
-    this._bufferedBytes += params[1].length;
+    this._bufferedBytes += params[3][kByteLength];
     this._queue.push(params);
   }
 


Test patch:
diff --git a/test/sender.test.js b/test/sender.test.js
index ef69d6ab..532239fa 100644
--- a/test/sender.test.js
+++ b/test/sender.test.js
@@ -47,6 +47,19 @@ describe('Sender', () => {
 
       assert.strictEqual(list[0][0] & 0x40, 0x40);
     });
+
+    it('accepts a string as first argument', () => {
+      const list = Sender.frame('€', {
+        readOnly: false,
+        rsv1: false,
+        mask: false,
+        opcode: 1,
+        fin: true
+      });
+
+      assert.deepStrictEqual(list[0], Buffer.from('8103', 'hex'));
+      assert.deepStrictEqual(list[1], Buffer.from('e282ac', 'hex'));
+    });
   });
 
   describe('#send', () => {
@@ -94,7 +107,7 @@ describe('Sender', () => {
 
             assert.strictEqual(chunks[0].length, 2);
             assert.notStrictEqual(chunk[0][0] & 0x40, 0x40);
-            assert.deepStrictEqual(chunks[1], Buffer.from('hi'));
+            assert.strictEqual(chunks[1], 'hi');
             done();
           }
         });
@@ -246,11 +259,12 @@ describe('Sender', () => {
 
           if (count % 2) {
             assert.ok(data.equals(Buffer.from([0x89, 0x02])));
-          } else {
+          } else if (count < 8) {
             assert.ok(data.equals(Buffer.from([0x68, 0x69])));
+          } else {
+            assert.strictEqual(data, 'hi');
+            done();
           }
-
-          if (count === 8) done();
         }
       });
       const sender = new Sender(mockSocket, {
@@ -278,11 +292,12 @@ describe('Sender', () => {
 
           if (count % 2) {
             assert.ok(data.equals(Buffer.from([0x8a, 0x02])));
-          } else {
+          } else if (count < 8) {
             assert.ok(data.equals(Buffer.from([0x68, 0x69])));
+          } else {
+            assert.strictEqual(data, 'hi');
+            done();
           }
-
-          if (count === 8) done();
         }
       });
       const sender = new Sender(mockSocket, {
diff --git a/test/websocket.integration.js b/test/websocket.integration.js
index 5ff87a64..abd96c61 100644
--- a/test/websocket.integration.js
+++ b/test/websocket.integration.js
@@ -6,43 +6,49 @@ const WebSocket = require('..');
 
 describe('WebSocket', () => {
   it('communicates successfully with echo service (ws)', (done) => {
-    const ws = new WebSocket('ws://echo.websocket.org/', {
-      origin: 'http://www.websocket.org',
+    const ws = new WebSocket('ws://websocket-echo.com/', {
       protocolVersion: 13
     });
-    const str = Date.now().toString();
 
     let dataReceived = false;
 
-    ws.on('open', () => ws.send(str));
+    ws.on('open', () => {
+      ws.send('hello');
+    });
+
     ws.on('close', () => {
       assert.ok(dataReceived);
       done();
     });
-    ws.on('message', (data) => {
+
+    ws.on('message', (message, isBinary) => {
       dataReceived = true;
-      assert.strictEqual(data, str);
+      assert.ok(!isBinary);
+      assert.strictEqual(message.toString(), 'hello');
       ws.close();
     });
   });
 
   it('communicates successfully with echo service (wss)', (done) => {
-    const ws = new WebSocket('wss://echo.websocket.org/', {
-      origin: 'https://www.websocket.org',
+    const ws = new WebSocket('wss://websocket-echo.com/', {
       protocolVersion: 13
     });
-    const str = Date.now().toString();
 
     let dataReceived = false;
 
-    ws.on('open', () => ws.send(str));
+    ws.on('open', () => {
+      ws.send('hello');
+    });
+
     ws.on('close', () => {
       assert.ok(dataReceived);
       done();
     });
-    ws.on('message', (data) => {
+
+    ws.on('message', (message, isBinary) => {
       dataReceived = true;
-      assert.strictEqual(data, str);
+      assert.ok(!isBinary);
+      assert.strictEqual(message.toString(), 'hello');
       ws.close();
     });
   });
diff --git a/test/websocket.test.js b/test/websocket.test.js
index 16cb33c5..b46f8928 100644
--- a/test/websocket.test.js
+++ b/test/websocket.test.js
@@ -33,7 +33,7 @@ describe('WebSocket', () => {
       );
 
       assert.throws(
-        () => new WebSocket('https://echo.websocket.org'),
+        () => new WebSocket('https://websocket-echo.com'),
         /^SyntaxError: The URL's protocol must be one of "ws:", "wss:", or "ws\+unix:"$/
       );
 
@@ -43,7 +43,7 @@ describe('WebSocket', () => {
       );
 
       assert.throws(
-        () => new WebSocket('wss://echo.websocket.org#foo'),
+        () => new WebSocket('wss://websocket-echo.com#foo'),
         /^SyntaxError: The URL contains a fragment identifier$/
       );
     });