Rate-limiting with a token bucket for the WeDo extension (#1540)

* Initial version of token bucket

* Add rateLimiter util

* Remove check for motor.isOn in stopAllMotors

* Fix unit test

* Fix unit test with stubbed timer, and cleanup

* Add comment

* Reduce WeDo rate limit to 20 sends/sec

* Move rate limit into a constant

* Stop button stops motors and tone even if rate limit exceeded

src/extensions/scratch3_wedo2/index.js

@@ -7,6 +7,7 @@ const log = require('../../util/log');
 const BLESession = require('../../io/bleSession');
 const Base64Util = require('../../util/base64-util');
 const MathUtil = require('../../util/math-util');
+const RateLimiter = require('../../util/rateLimiter.js');
 
 /**
  * Icon svg to be displayed at the left edge of each extension block, encoded as a data URI.
@@ -30,6 +31,12 @@ const UUID = {
  */
 const BLESendInterval = 100;
 
+/**
+ * A maximum number of BLE message sends per second, to be enforced by the rate limiter.
+ * @type {number}
+ */
+const BLESendRateMax = 20;
+
 /**
  * Enum for WeDo2 sensor and output types.
  * @readonly
@@ -260,15 +267,16 @@ class WeDo2Motor {
 
     /**
      * Turn this motor off.
+     * @param {boolean} [useLimiter=true] - if true, use the rate limiter
      */
-    setMotorOff () {
+    setMotorOff (useLimiter = true) {
         const cmd = new Uint8Array(4);
         cmd[0] = this._index + 1; // connect id
         cmd[1] = WeDo2Commands.MOTOR_POWER; // command
         cmd[2] = 1; // 1 byte to follow
         cmd[3] = 0; // power in range 0-100
 
-        this._parent._send(UUID.OUTPUT_COMMAND, Base64Util.uint8ArrayToBase64(cmd));
+        this._parent._send(UUID.OUTPUT_COMMAND, Base64Util.uint8ArrayToBase64(cmd), useLimiter);
 
         this._isOn = false;
     }
@@ -346,19 +354,6 @@ class WeDo2 {
             distance: 0
         };
 
-        /**
-         * A flag that is true while we are busy sendng data to the BLE session.
-         * @type {boolean}
-         * @private
-         */
-        this._sending = false;
-
-        /**
-         * ID for a timeout which is used to clear the sending flag if it has been
-         * true for a long time.
-         */
-        this._sendingTimeoutID = null;
-
         /**
          * The Bluetooth connection session for reading/writing device data.
          * @type {BLESession}
@@ -369,6 +364,14 @@ class WeDo2 {
 
         this._onConnect = this._onConnect.bind(this);
         this._onMessage = this._onMessage.bind(this);
+
+        /**
+         * A rate limiter utility, to help limit the rate at which we send BLE messages
+         * over the socket to Scratch Link to a maximum number of sends per second.
+         * @type {RateLimiter}
+         * @private
+         */
+        this._rateLimiter = new RateLimiter(BLESendRateMax);
     }
 
     /**
@@ -406,8 +409,11 @@ class WeDo2 {
      */
     stopAllMotors () {
         this._motors.forEach(motor => {
-            if (motor && motor.isOn) {
-                motor.setMotorOff();
+            if (motor) {
+                // Send the motor off command without using the rate limiter.
+                // This allows the stop button to stop motors even if we are
+                // otherwise flooded with commands.
+                motor.setMotorOff(false);
             }
         });
     }
@@ -473,7 +479,9 @@ class WeDo2 {
         cmd[0] = WeDo2ConnectIDs.PIEZO; // connect id
         cmd[1] = WeDo2Commands.STOP_TONE; // command
 
-        return this._send(UUID.OUTPUT_COMMAND, Base64Util.uint8ArrayToBase64(cmd));
+        // Send this command without using the rate limiter, because it is only triggered
+        // by the stop button.
+        return this._send(UUID.OUTPUT_COMMAND, Base64Util.uint8ArrayToBase64(cmd), false);
     }
 
     /**
@@ -545,24 +553,18 @@ class WeDo2 {
      * Write a message to the device BLE session.
      * @param {number} uuid - the UUID of the characteristic to write to
      * @param {Uint8Array} message - the message to write.
+     * @param {boolean} [useLimiter=true] - if true, use the rate limiter
      * @return {Promise} - a promise result of the write operation
      * @private
      */
-    _send (uuid, message) {
-        if (!this.getPeripheralIsConnected()) return;
-        if (this._sending) return;
+    _send (uuid, message, useLimiter = true) {
+        if (!this.getPeripheralIsConnected()) return Promise.resolve();
 
-        this._sending = true;
+        if (useLimiter) {
+            if (!this._rateLimiter.okayToSend()) return Promise.resolve();
+        }
 
-        this._sendingTimeoutID = window.setTimeout(() => {
-            this._sending = false;
-        }, 5000);
-
-        return this._ble.write(UUID.IO_SERVICE, uuid, message, 'base64')
-            .then(() => {
-                this._sending = false;
-                window.clearTimeout(this._sendingTimeoutID);
-            });
+        return this._ble.write(UUID.IO_SERVICE, uuid, message, 'base64');
     }
 
     /**

src/util/rateLimiter.js

@@ -0,0 +1,73 @@
+const Timer = require('../util/timer');
+
+class RateLimiter {
+    /**
+     * A utility for limiting the rate of repetitive send operations, such as
+     * bluetooth messages being sent to hardware devices. It uses the token bucket
+     * strategy: a counter accumulates tokens at a steady rate, and each send costs
+     * a token. If no tokens remain, it's not okay to send.
+     * @param {number} maxRate the maximum number of sends allowed per second
+     * @constructor
+     */
+    constructor (maxRate) {
+        /**
+         * The maximum number of tokens.
+         * @type {number}
+         */
+        this._maxTokens = maxRate;
+
+        /**
+         * The interval in milliseconds for refilling one token. It is calculated
+         * so that the tokens will be filled to maximum in one second.
+         * @type {number}
+         */
+        this._refillInterval = 1000 / maxRate;
+
+        /**
+         * The current number of tokens in the bucket.
+         * @type {number}
+         */
+        this._count = this._maxTokens;
+
+        this._timer = new Timer();
+        this._timer.start();
+
+        /**
+         * The last time in milliseconds when the token count was updated.
+         * @type {number}
+         */
+        this._lastUpdateTime = this._timer.timeElapsed();
+    }
+
+    /**
+     * Check if it is okay to send a message, by updating the token count,
+     * taking a token and then checking if we are still under the rate limit.
+     * @return {boolean} true if we are under the rate limit
+     */
+    okayToSend () {
+        // Calculate the number of tokens to refill the bucket with, based on the
+        // amount of time since the last refill.
+        const now = this._timer.timeElapsed();
+        const timeSinceRefill = now - this._lastUpdateTime;
+        const refillCount = Math.floor(timeSinceRefill / this._refillInterval);
+
+        // If we're adding at least one token, reset _lastUpdateTime to now.
+        // Otherwise, don't reset it so that we can continue measuring time until
+        // the next refill.
+        if (refillCount > 0) {
+            this._lastUpdateTime = now;
+        }
+
+        // Refill the tokens up to the maximum
+        this._count = Math.min(this._maxTokens, this._count + refillCount);
+
+        // If we have at least one token, use one, and it's okay to send.
+        if (this._count > 0) {
+            this._count--;
+            return true;
+        }
+        return false;
+    }
+}
+
+module.exports = RateLimiter;

test/unit/util_rateLimiter.js

@@ -0,0 +1,32 @@
+const test = require('tap').test;
+const RateLimiter = require('../../src/util/rateLimiter.js');
+
+test('rate limiter', t => {
+    // Create a rate limiter with maximum of 20 sends per second
+    const rate = 20;
+    const limiter = new RateLimiter(rate);
+
+    // Simulate time passing with a stubbed timer
+    let simulatedTime = Date.now();
+    limiter._timer = {timeElapsed: () => simulatedTime};
+
+    // The rate limiter starts with a number of tokens equal to the max rate
+    t.equal(limiter._count, rate);
+
+    // Running okayToSend a number of times equal to the max rate
+    // uses up all of the tokens
+    for (let i = 0; i < rate; i++) {
+        t.true(limiter.okayToSend());
+        // Tokens are counting down
+        t.equal(limiter._count, rate - (i + 1));
+    }
+    t.false(limiter.okayToSend());
+
+    // Advance the timer enough so we get exactly one more token
+    // One extra millisecond is required to get over the threshold
+    simulatedTime += (1000 / rate) + 1;
+    t.true(limiter.okayToSend());
+    t.false(limiter.okayToSend());
+
+    t.end();
+});