rework lua script and add basic functional test

Author: tmacro

lib/api/apiUtils/rateLimit/client.js

@@ -23,24 +23,47 @@ class RateLimitClient {
         });
     }
 
-    getLocalCounter(key, cb) {
-        this.redis.get(key, cb);
-    }
+    /**
+     * @typedef {Object} CounterUpdateBatch
+     * @property {string} key - counter key
+     * @property {number} cost - cost to add to counter
+     */
+
+    /**
+     * @typedef {Object} CounterUpdateBatchResult
+     * @property {string} key - counter key
+     * @property {number} value - current value of counter
+     */
 
+    /**
+     * @callback RateLimitClient~batchUpdate
+     * @param {Error} err
+     * @param {undefined|CounterUpdateBatchResult[]}
+     */
+
+    /**
+     * Add cost to the counter at key.
+     * Returns the new value for  the counter
+     *
+     * @param {CounterUpdateBatch[]} batch - batch of counter updates
+     * @param {RateLimitClient~batchUpdate} cb
+     */
     updateLocalCounters(batch, cb) {
         const pipeline = this.redis.pipeline();
-        for (const { key, ts, counter } of batch) {
-            pipeline.updateCounter(key, ts, counter);
+        for (const { key, cost } of batch) {
+            pipeline.updateCounter(key, cost);
         }
 
-
         pipeline.exec((err, results) => {
             if (err) {
                 cb(err);
                 return;
             }
 
-            cb(null, results.map(([_, counter], i) => [batch[i].key, counter]));
+            cb(null, results.map(([_, value], i) => ({
+                key: batch[i].key,
+                value,
+            })));
         });
     }
 }

lib/api/apiUtils/rateLimit/updateCounter.lua

@@ -1,14 +1,21 @@
-local counterExists = redis.call('EXISTS', KEYS[1])
+local ts = redis.call('TIME')
+local currentTime = ts[1] * 1000
+currentTime = currentTime + math.floor(ts[2] / 1000)
+
+local newValue = currentTime + tonumber(ARGV[1])
 
-if counterExists ~= 1 then
-    redis.call('SET', KEYS[1], ARGV[2])
-    return ARGV[2]
-else
-    local currentValue = redis.call('GET', KEYS[1])
-    local localDelay = currentValue - ARGV[1]
-    local remoteDelay = ARGV[2] - ARGV[1]
-    local newCounter = localDelay + remoteDelay + ARGV[1]
-    redis.call('SET', KEYS[1], newCounter)
-    redis.call('EXPIREAT', KEYS[1], newCounter)
-    return newCounter
+local counterExists = redis.call('EXISTS', KEYS[1])
+if counterExists == 1 then
+    local currentValue = tonumber(redis.call('GET', KEYS[1]))
+    if currentValue > currentTime then
+        local currentCost = currentValue - currentTime
+        newValue = newValue + currentCost
+    end
 end
+
+redis.call('SET', KEYS[1], newValue)
+
+local expiry = math.ceil(newValue / 1000)
+redis.call('EXPIREAT', KEYS[1], expiry)
+
+return newValue

tests/functional/aws-node-sdk/test/rateLimit/client.js

@@ -0,0 +1,46 @@
+const assert = require('assert');
+
+const { config } = require('../../../../../lib/Config');
+const { RateLimitClient } = require('../../../../../lib/api/apiUtils/rateLimit/client');
+
+
+const counterKey = 'foo';
+
+describe('Test RateLimitClient', () => {
+    let client;
+
+    before(done => {
+        client = new RateLimitClient(config.localCache);
+        client.redis.connect(done);
+    });
+
+    beforeEach(done => {
+        client.redis.del(counterKey, err => done(err));
+    });
+
+    it('should set the value of an empty counter', done => {
+        const batch = [{ key: counterKey, cost: 10000 }];
+        client.updateLocalCounters(batch, (err, res) => {
+            assert.ifError(err);
+            assert.strictEqual(res.length, 1);
+            assert.strictEqual(res[0].key, counterKey);
+            done();
+        });
+    });
+
+    it('should increment the value of an existing counter', done => {
+        const batch = [{ key: counterKey, cost: 10000 }];
+        client.updateLocalCounters(batch, (err, res) => {
+            assert.ifError(err);
+            console.log({ res });
+            const { value: existingValue } = res[0];
+            client.updateLocalCounters(batch, (err, res) => {
+                assert.ifError(err);
+            console.log({ res });
+                const { value: newValue } = res[0];
+                assert(newValue > existingValue, `${newValue} is not greater than ${existingValue}`);
+                done();
+            });
+        });
+    });
+});

tests/unit/api/apiUtils/rateLimit/client.js

@@ -17,12 +17,12 @@ class PipelineStub {
         this.ops = [];
     }
 
-    updateCounter(key, ts, value) {
-        this.ops.push([key, ts, value]);
+    updateCounter(key, cost) {
+        this.ops.push([key, cost]);
     }
 
     exec(cb) {
-        cb(null, this.ops.map(v => [1, v[2]]));
+        cb(null, this.ops.map(v => [1, v[1]]));
     }
 }
 
@@ -38,19 +38,18 @@ describe('test RateLimitClient', () => {
     });
 
     it('should update a batch of counters', done => {
-        const ts = Date.now();
         const batch = [
-            { key: 'foo', ts, counter: ts + 100 },
-            { key: 'bar', ts, counter: ts + 200 },
-            { key: 'qux', ts, counter: ts + 300 },
+            { key: 'foo', cost: 100 },
+            { key: 'bar', cost: 200 },
+            { key: 'qux', cost: 300 },
         ]
 
         client.updateLocalCounters(batch, (err, results) => {
             assert.ifError(err);
             assert.deepStrictEqual(results, [
-                ['foo', ts + 100],
-                ['bar', ts + 200],
-                ['qux', ts + 300],
+                { key: 'foo', value: 100 },
+                { key: 'bar', value: 200 },
+                { key: 'qux', value: 300 },
             ]);
             done();
         });