Update RateLimiter to scale credits on update

Signed-off-by: Won Jun Jang <wjang@uber.com>

src/rate_limiter.js

@@ -25,14 +25,26 @@ export default class RateLimiter {
   }
 
   update(creditsPerSecond: number, maxBalance: number) {
+    this._updateBalance();
     this._creditsPerSecond = creditsPerSecond;
+    // The new balance should be proportional to the old balance.
+    this._balance = maxBalance * this._balance / this._maxBalance;
     this._maxBalance = maxBalance;
     if (this._balance > maxBalance) {
       this._balance = maxBalance;
     }
   }
 
   checkCredit(itemCost: number): boolean {
+    this._updateBalance();
+    if (this._balance >= itemCost) {
+      this._balance -= itemCost;
+      return true;
+    }
+    return false;
+  }
+
+  _updateBalance() {
     let currentTime: number = new Date().getTime();
     let elapsedTime: number = (currentTime - this._lastTick) / 1000;
     this._lastTick = currentTime;
@@ -41,10 +53,5 @@ export default class RateLimiter {
     if (this._balance > this._maxBalance) {
       this._balance = this._maxBalance;
     }
-    if (this._balance >= itemCost) {
-      this._balance -= itemCost;
-      return true;
-    }
-    return false;
   }
 }

test/init_tracer.js

@@ -168,7 +168,7 @@ describe('initTracer', () => {
     assert.equal(tracer._tags['x'], 'y');
   });
 
-it('should pass options to remote sampler and reporter', () => {
+  it('should pass options to remote sampler and reporter', () => {
     let logger = {
       info: function info(msg) {},
     };
@@ -189,7 +189,7 @@ it('should pass options to remote sampler and reporter', () => {
         sampler: {
           type: 'remote',
           param: 0,
-        }
+        },
       },
       {
         logger: logger,

test/rate_limiter.js

@@ -69,7 +69,11 @@ describe('leaky bucket ratelimiter should', () => {
     let initialDate = new Date(2011, 9, 1).getTime();
     let clock = sinon.useFakeTimers(initialDate);
     let limiter = new RateLimiter(0.1, 1);
-    assert.equal(true, limiter._balance <= 1 && limiter._balance >= 0);
+    assert.equal(
+      true,
+      limiter._balance <= 1 && limiter._balance >= 0,
+      'balance should be initialized to a random value between [0:1]'
+    );
     limiter._balance = 1.0;
     assert.equal(limiter.checkCredit(1), true, 'expected checkCredit to be true');
 
@@ -87,10 +91,11 @@ describe('leaky bucket ratelimiter should', () => {
     assert.equal(limiter.checkCredit(1), true, 'expected checkCredit to be true');
     assert.equal(limiter._balance, 2, 'balance should be at 2 after spending 1');
 
+    // move time 5s forward, not enough to accumulate credits for another message
+    clock = sinon.useFakeTimers(initialDate + 55000);
     // reduce the maxBalance so the limiter is capped at 1
     limiter.update(0.1, 1);
-    assert.equal(limiter._balance, 1, 'balance should be at 1 after update');
-    assert.equal(limiter.checkCredit(1), true, 'expected checkCredit to be true');
+    assert.equal(limiter._balance, 2.5 / 3.0, 'balance should be proportional to the original range');
     assert.equal(limiter.checkCredit(1), false, 'expected checkCredit to be false');
     clock.restore();
   });