basic client with simple unit test

Author: tmacro

lib/Config.js

@@ -1839,6 +1839,11 @@ class Config extends EventEmitter {
 
 
         if (config.rateLimiting?.enabled) {
+            // rate limiting uses the same localCache config defined for S3 to avoid
+            // config duplication.
+            assert(config.localCache, 'missing required property of rateLimit ' +
+                'configuration: localCache');
+
             this.rateLimiting.enabled = true;
 
             assert.strictEqual(typeof config.rateLimiting.serviceUserArn, 'string');

lib/api/apiUtils/rateLimit/client.js

@@ -0,0 +1,56 @@
+const fs = require('fs');
+
+const Redis = require('ioredis');
+
+const { config } = require('../../../Config');
+
+
+const updateCounterScript = fs.readFileSync(`${__dirname  }/updateCounter.lua`).toString();
+
+const SCRIPTS = {
+    updateCounter: {
+        numberOfKeys: 1,
+        lua: updateCounterScript,
+    },
+};
+
+class RateLimitClient {
+    constructor(redisConfig) {
+        this.redis = new Redis({
+            ...redisConfig,
+            scripts: SCRIPTS,
+            lazyConnect: true,
+        });
+    }
+
+    getLocalCounter(key, cb) {
+        this.redis.get(key, cb);
+    }
+
+    updateLocalCounters(batch, cb) {
+        const pipeline = this.redis.pipeline();
+        for (const { key, ts, counter } of batch) {
+            pipeline.updateCounter(key, ts, counter);
+        }
+
+
+        pipeline.exec((err, results) => {
+            if (err) {
+                cb(err);
+                return;
+            }
+
+            cb(null, results.map(([_, counter], i) => [batch[i].key, counter]));
+        });
+    }
+}
+
+let counterClient;
+if (config.rateLimiting.enabled) {
+    counterClient = new RateLimitClient(config.localCache);
+}
+
+module.exports = {
+    counterClient,
+    RateLimitClient
+}

lib/api/apiUtils/rateLimit/updateCounter.lua

@@ -0,0 +1,14 @@
+local counterExists = redis.call('EXISTS', KEYS[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
+end

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

@@ -0,0 +1,58 @@
+const assert = require('assert');
+
+const { RateLimitClient } = require('../../../../../lib/api/apiUtils/rateLimit/client');
+
+class RedisStub {
+    constructor() {
+        this.data = {}
+    }
+
+    pipeline() {
+        return new PipelineStub();
+    }
+}
+
+class PipelineStub {
+    constructor() {
+        this.ops = [];
+    }
+
+    updateCounter(key, ts, value) {
+        this.ops.push([key, ts, value]);
+    }
+
+    exec(cb) {
+        cb(null, this.ops.map(v => [1, v[2]]));
+    }
+}
+
+describe('test RateLimitClient', () => {
+    let client;
+
+    before(() => {
+        client = new RateLimitClient({});
+    });
+
+    beforeEach(() => {
+        client.redis = new RedisStub();
+    });
+
+    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 },
+        ]
+
+        client.updateLocalCounters(batch, (err, results) => {
+            assert.ifError(err);
+            assert.deepStrictEqual(results, [
+                ['foo', ts + 100],
+                ['bar', ts + 200],
+                ['qux', ts + 300],
+            ]);
+            done();
+        });
+    });
+});