Add cancelAll method to clear the task queue

src/util/token-bucket.js

@@ -14,10 +14,11 @@ class TokenBucket {
     constructor (maxTokens, refillRate, startingTokens = maxTokens) {
         this._maxTokens = maxTokens;
         this._refillRate = refillRate;
-        this._pendingTasks = [];
+        this._pendingTaskRecords = [];
         this._tokenCount = startingTokens;
         this._timer = new Timer();
         this._timer.start();
+        this._timeout = null;
         this._lastUpdateTime = this._timer.timeElapsed();
     }
 
@@ -30,13 +31,13 @@ class TokenBucket {
      * @memberof TokenBucket
      */
     do (task, cost = 1) {
-        let wrappedTask;
+        const newRecord = {};
         const promise = new Promise((resolve, reject) => {
-            wrappedTask = () => {
+            newRecord.wrappedTask = () => {
                 const canRun = this._refillAndSpend(cost);
                 if (canRun) {
                     // Remove this task from the queue and run it
-                    this._pendingTasks.shift();
+                    this._pendingTaskRecords.shift();
                     try {
                         resolve(task());
                     } catch (e) {
@@ -44,25 +45,38 @@ class TokenBucket {
                     }
 
                     // Tell the next wrapper to start trying to run its task
-                    if (this._pendingTasks.length > 0) {
-                        const nextWrappedTask = this._pendingTasks[0];
-                        nextWrappedTask();
+                    if (this._pendingTaskRecords.length > 0) {
+                        const nextRecord = this._pendingTaskRecords[0];
+                        nextRecord.wrappedTask();
                     }
                 } else {
                     // This task can't run yet. Estimate when it will be able to, then try again.
-                    this._waitUntilAffordable(cost).then(() => wrappedTask());
+                    newRecord.reject = reject;
+                    this._waitUntilAffordable(cost).then(() => newRecord.wrappedTask());
                 }
             };
         });
-        this._pendingTasks.push(wrappedTask);
+        this._pendingTaskRecords.push(newRecord);
 
-        if (this._pendingTasks.length === 1) {
-            wrappedTask();
+        if (this._pendingTaskRecords.length === 1) {
+            newRecord.wrappedTask();
         }
 
         return promise;
     }
 
+    /**
+     * Cancel all pending tasks, rejecting all their promises.
+     */
+    cancelAll () {
+        if (this._timeout !== null) {
+            clearTimeout(this._timeout);
+            this._timeout = null;
+        }
+        this._pendingTaskRecords.forEach(r => r.reject());
+        this._pendingTaskRecords = [];
+    }
+
     /**
      * Shorthand for calling @ _refill() then _spend(cost).
      * @see {@link TokenBucket#_refill}
@@ -121,9 +135,17 @@ class TokenBucket {
             return Promise.reject(new Error(`Task cost ${cost} is greater than bucket limit ${this._maxTokens}`));
         }
         return new Promise(resolve => {
-            const tokensNeeded = this._tokenCount - cost;
+            const tokensNeeded = Math.max(cost - this._tokenCount, 0);
             const estimatedWait = Math.ceil(1000 * tokensNeeded / this._refillRate);
-            setTimeout(resolve, estimatedWait);
+
+            let timeout = null;
+            const onTimeout = () => {
+                if (this._timeout === timeout) {
+                    this._timeout = null;
+                }
+                resolve();
+            };
+            this._timeout = timeout = setTimeout(onTimeout, estimatedWait);
         });
     }
 }

test/unit/util_token-bucket.js

@@ -14,18 +14,28 @@ test('constructor', t => {
 
     const taskResults = [];
     const promises = [];
+    const goodCancelMessage = 'Task was canceled correctly';
+    bukkit.do(() => taskResults.push('nope'), 999).then(
+        () => {
+            t.fail('Task should have been canceled');
+        },
+        () => {
+            taskResults.push(goodCancelMessage);
+        }
+    );
+    bukkit.cancelAll();
     promises.push(
-        bukkit.do(() => taskResults.push('a'), 100).then(() =>
-            testCompare(t, timer.timeElapsed(), '>=', 100, 'Costly task must wait')
+        bukkit.do(() => taskResults.push('a'), 50).then(() =>
+            testCompare(t, timer.timeElapsed(), '>=', 50, 'Costly task must wait')
         ),
-        bukkit.do(() => taskResults.push('b'), 0).then(() =>
-            testCompare(t, timer.timeElapsed(), '<', 150, 'Cheap task should run soon')
+        bukkit.do(() => taskResults.push('b'), 10).then(() =>
+            testCompare(t, timer.timeElapsed(), '>=', 60, 'Tasks must run in serial')
         ),
-        bukkit.do(() => taskResults.push('c'), 101).then(() =>
-            testCompare(t, timer.timeElapsed(), '>=', 200, 'Tasks must run in serial')
+        bukkit.do(() => taskResults.push('c'), 1).then(() =>
+            testCompare(t, timer.timeElapsed(), '<', 80, 'Cheap task should run soon')
         )
     );
     return Promise.all(promises).then(() => {
-        t.deepEqual(taskResults, ['a', 'b', 'c'], 'All tasks must run in correct order');
+        t.deepEqual(taskResults, [goodCancelMessage, 'a', 'b', 'c'], 'All tasks must run in correct order');
     });
 });