Throttle validation NodeJS library implementation (using Redis, MySQL and Promises)

throttle.js

/**
 * Handle Throttling Values and restrict accessing when throttling
 * value exceeded.
 * 
 * This class implementation is inspired from Redis
 * Refer: http://redis.io/commands/incr#pattern-rate-limiter-1
 *
 * Gihan Karunarathne <[email protected]>
 */
'use strict';

var redis = require('redis');
var mysql = require('mysql');
var klass = require('klass');
var Promise = require('bluebird');
var debug = require('debug')('throttle');

var RateLimiter = module.exports = klass({
    /**
     * Initialize RateLimiter
     * USAGE:
     *   var RateLimiter = require('RateLimiter');
     *   var rateLimiter = new RateLimiter({
     *     redis: {
     *       host: '',
     *       port: ''
     *     },
     *     mysql: {
     *       host    : '',
     *       user    : '',
     *       password: '',
     *       database: ''
     *     }
     *   });
     */
    initialize: function(opts) {
        this.rateLimiterTag = 'rateLimiter';
        this.perSeconds = opts.perSeconds || 60; // Default 60 seconds
        debug('Redis Options: ', opts.redis);
        this.redisClient = redis.createClient(opts.redis);
        debug('MySQL Options: ', opts.mysql);
        this.connection = mysql.createConnection(opts.mysql);

        this.redisClient.on('ready', function() {
            debug('Rate Limiter: connected to redis server ...');
        });

        this.connection.connect(function(err) {
            if (err) {
                console.error(err);
                return;
            }
            debug('Connected to mysql server ', opts.mysql.host);
        });
    },

    /**
     * Check whether exceeds the rate limit
     * USAGE:
     *   rateLimiter.isExceed({
     *     serviceId: 101
     *   }).then(function(result) {
     *     // Handle validation here
     *     {
     *       continue: true/false
     *     }
     *   }).catch(function(err) {
     *     // Handle error here
     *   });
     *
     * Note: This method is inspired from Redis,
     * but deviate from original implementation.
     * Refer: http://redis.io/commands/incr#pattern-rate-limiter-1
     * According to rate limiter functionality, it's obvious that most of
     * time it's going to continue instead of blocking. Thus Redis algorithm
     * need to make two database calls. In this implementation,
     * it optimize database calls by only using decr database call
     * (trying to take advantage of returning value while decrease count.)
     */
    isExceed: function(args) {
        var self = this;
        debug('IsExeed: ', args);

        return new Promise(function(resolve, reject) {
            self.decrValue(args.serviceId).then(function(currentValue) {
                resolve({
                    continue: (currentValue > 0)
                });
            }).catch(function(err) {
                // Unable to decrease value, then set value
                self.getRateLimit(args).then(function(value) {
                    self.setThrottlingValue(args.serviceId, value.throttling).
                    then(function() {
                        resolve({
                            continue: true
                        });
                    }).catch(function(err) {
                        console.error(err);
                        reject(err);
                    });
                }).catch(function(err) {
                    console.error(err);
                    reject(err);
                });
            });
        });
    },

    /*****************************************************************
     * Redis - Rate Limit using Redis data structures
     ****************************************************************/
    /**
     * Decrease throttling value
     * USAGE:
     *   rateLimiter.decrValue(101).then(function(value) {
     *     // Handle current throttling value here
     *     {
     *       throttling: 999
     *     }
     *   }).catch(function(err) {
     *     // Handle error here
     *   });
     */
    decrValue: function(key) {
        var self = this;
        debug('DECRValue: ', key);

        return new Promise(function(resolve, reject) {
            var newKey = self.rateLimiterTag + key.toString();
            self.redisClient.decr(newKey, function(err, value) {
                debug('Value after decr: ', value);
                // If key is setting first time, it'll return -1
                if (value > 0) {
                    resolve(value);
                } else {
                    if (value > -1) {
                        /** Rationale: Increase by one, due to next time
                         validation happens after decrease the value */
                        self.redisClient.incr(newKey, function(err, newVal) {
                            if (err) {
                                reject(err);
                            } else {
                                resolve(value);
                            }
                        });
                    } else {
                        console.error(err);
                        reject(new Error('Throttling value is not set.'));
                    }
                }
            });
        });
    },

    /**
     * Cache throttling value using Redis.
     * Set expire time by using `perSeconds` value.
     * USAGE:
     *   rateLimiter.setThrottlingValue(101, 1000).then(function(){
     *     // Continue
     *   }).catch(function(err){
     *     // Handle error here
     *   });
     */
    setThrottlingValue: function(key, rateLimit) {
        var self = this;
        debug('SetThrottlingValue: ', key, rateLimit);

        return new Promise(function(resolve, reject) {
            self.redisClient.setex(self.rateLimiterTag + key.toString(),
                self.perSeconds, rateLimit,
                function(err, isSet) {
                    debug('SETEX: ', isSet, err);
                    if (err || !isSet) {
                        console.error(err);
                        reject(err);
                    } else {
                        resolve();
                    }
                });
        });
    },

    /**
     * MySQL - Get Rate Limit info
     * USAGE:
     *   rateLimiter.getRatelimit({
     *     serviceId: 101
     *   }).then(function(value) {
     *     // Handle value here
     *     {
     *       throttling: 1000
     *     }
     *   }).catch(function(err) {
     *     // Handler error here
     *   })
     */
    getRateLimit: function(args) {
        var self = this;
        debug('GetRateLimit: ', args);

        return new Promise(function(resolve, reject) {
            var columns = ['throttling'];
            var sql = 'SELECT ?? FROM ?? WHERE serviceId=?';
            // In order to avoid SQL injection
            var replace = [columns, 'tableName', args.serviceId];

            self.connection.query(sql, replace, function(err, rows) {
                if (err) {
                    console.error(err);
                    reject(err);
                } else {
                    debug('Rows: ', rows);
                    if (rows.length === 1) {
                        resolve({
                            throttling: rows[0].throttling
                        });
                    } else {
                        reject(new Error('Unable to find unique ' +
                            'throttling value'));
                    }
                }
            });
        });
    }

});

throttle_test.js

/**
 *  Test cases for Throttle Class
 *
 * Gihan Karunarathne <[email protected]>
 *
 * USAGE:
 *   1. Install dependencies using 'npm install'
 *   2. Run test cases using 'npm test'
 */
'use strict';

var RateLimiter = require('./throttle');

var assert = require('assert');
var async = require('async');
var testConf = require('./wtest.config.json');
var debug = require('debug')('test:throttle');

describe('Throttle', function() {
    var rateLimiter;

    before(function(done) {
        rateLimiter = new RateLimiter({
            redis: testConf.redis,
            mysql: testConf.mysql,
            perSeconds: 5
        });
        done();
    });

    after(function(done) {
        done();
    });

    describe('Throttle', function() {

        it('should not block first time request', function(done) {
            rateLimiter.isExceed({
                serviceId: 1
            }).then(function(result) {
                assert.ok(result.continue);
                done();
            }).catch(function(err) {
                assert.ifError(err);
            });
        });

        it('should not block second time request', function(done) {
            var serviceId = 1;
            rateLimiter.isExceed({
                serviceId: serviceId
            }).then(function(result) {
                assert.ok(result.continue);
                rateLimiter.isExceed({
                    serviceId: serviceId
                }).then(function(result) {
                    assert.ok(result.continue);
                    done();
                }).catch(function(err) {
                    assert.ifError(err);
                });
            }).catch(function(err) {
                assert.ifError(err);
            });
        });

        it('should block after rate limit exceed', function(done) {
            var serviceId = 101;
            rateLimiter.getRateLimit({
                serviceId: serviceId
            }).then(function(value) {
                debug(value);
                var arr = Array.apply(null, {
                    length: value.throttling
                }).map(function(num) {
                    return serviceId;
                }); // http://stackoverflow.com/a/20066663/1461060

                // Should not block until throttling value reached
                async.eachSeries(arr, function(serviceIdTmp, callback) {
                    rateLimiter.isExceed({
                        serviceId: serviceIdTmp
                    }).then(function(result) {
                        assert.ok(result.continue);
                        callback();
                    }).catch(function(err) {
                        callback(err);
                    });
                }, function(err) {
                    assert.ifError(err);
                    var arr2 = Array.apply(null, {
                        length: 5
                    }).map(function(num) {
                        return serviceId;
                    }); // http://stackoverflow.com/a/20066663/1461060

                    // Should block after throttling value exceeded
                    async.eachSeries(arr2, function(serviceIdTmp2, callback) {
                        rateLimiter.isExceed({
                            serviceId: serviceIdTmp2
                        }).then(function(result2) {
                            debug('Should not validate exceed ', result2);
                            assert.ok(!result2.continue);
                            callback();
                        }).catch(function(err) {
                            assert.ifError(err);
                        });
                    }, function(err) {
                        assert.ifError(err);
                        done();
                    });
                }); // END - async
            }).catch(function(err) {
                assert.ifError(err);
            }); // END - GetRateLimit
        }); // END - it

    });

});

wtest.config.json

{
    "mysql": {
        "host": "localhost",
        "user": "root",
        "password": "abc123",
        "database": "database"
    },
    "redis": {
        "host": "localhost",
        "port": "6379",
        "options": {}
    }
}