qps-limit.js 2.75 KB
'use strict';

const logger = global.yoho.logger;
const cache = global.yoho.cache.master;
const config = global.yoho.config;
const ONE_DAY = 60 * 60 * 24;
const MAX_QPS = config.maxQps;
const MAX_QPS_10m = config.maxQps10m;
const _ = require('lodash');

const PAGES = {
    '/product/^\\/(\\d+)\\.html/': 5,
    '/product/list/index': 5,
    '/product/index/index': 5,
    '/product/search/list': 5
};

function urlJoin(a, b) {
    if (_.endsWith(a, '/') && _.startsWith(b, '/')) {
        return a + b.substring(1, b.length);
    } else if (!_.endsWith(a, '/') && !_.startsWith(b, '/')) {
        return a + '/' + b;
    } else {
        return a + b;
    }
}

module.exports = (limiter, policy) => {
    const req = limiter.req,
        res = limiter.res;

    const key = `pc:limiter:${limiter.remoteIp}`;
    const key10m = `pc:limiter:10m:${limiter.remoteIp}`;

    res.on('render', function() {
        let route = req.route ? req.route.path : '';
        let appPath = req.app.mountpath;

        if (_.isArray(route) && route.length > 0) {
            route = route[0];
        }

        let pageKey = urlJoin(appPath, route.toString()); // route may be a regexp
        let pageIncr = PAGES[pageKey] || 0;

        if (pageIncr > 0) {
            cache.incrAsync(key, pageIncr);
            cache.incrAsync(key10m, pageIncr);
        }
    });

    return cache.getMultiAsync([key, key10m]).then((results) => {
        let result = results[key];
        let result10m = results[key10m];

        logger.debug('qps limiter: ' + key + '@' + result + ' max: ' + MAX_QPS);
        logger.debug('qps limiter 10m: ' + key10m + '@' + result10m + ' max: ' + MAX_QPS_10m);

        // 默认数据设置
        if (!result && !_.isNumber(result)) {
            cache.setAsync(key, 1, 60);  // 设置key,1m失效
        }

        if (!result10m && !_.isNumber(result10m)) {
            cache.setAsync(key10m, 1, 600);  // 设置key,10m失效
        }

        // 第一次访问,都没计数,直接过
        if (!result && !_.isNumber(result) && !result10m && !_.isNumber(result10m)) {
            return Promise.resolve(true);
        }

        if (result === -1 || result10m === -1) {
            return Promise.resolve(true);
        }

        // 判断 qps 10分钟
        if (result10m > MAX_QPS_10m) {
            cache.touch(key10m, ONE_DAY);
            logger.debug('req limit', key10m);

            return Promise.resolve(policy);
        }

        // 判断 qps 1分钟
        if (result > MAX_QPS) {
            cache.touch(key, ONE_DAY);
            logger.debug('req limit', key);

            return Promise.resolve(policy);
        }

        cache.incrAsync(key, 1); // qps + 1
        cache.incrAsync(key10m, 1); // qps + 1
        return Promise.resolve(true);
    });
};