001/**
002 * Copyright (c) 2015-2022, Michael Yang 杨福海 (fuhai999@gmail.com).
003 * <p>
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 * <p>
008 * http://www.apache.org/licenses/LICENSE-2.0
009 * <p>
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016package io.jboot.components.limiter.redis;
017
018import io.jboot.exception.JbootIllegalConfigException;
019import io.jboot.support.redis.JbootRedis;
020import io.jboot.support.redis.JbootRedisManager;
021
022/**
023 * 通过lua脚本来进行限次
024 */
025public class RedisRateLimitUtil {
026
027    private static final String RATE_LIMIT_SCRIPT = "local c" +
028            "\nc = redis.call('get',KEYS[1])" +
029            // 调用量已经超过最大值，直接返回
030            "\nif c and tonumber(c) > tonumber(ARGV[1]) then" +
031            "\nreturn tonumber(c);" +
032            "\nend" +
033            // 自增
034            "\nc = redis.call('incr',KEYS[1])" +
035            "\nif tonumber(c) == 1 then" +
036            // 从第一次调用开始限流，设置对应键值的过期
037            "\nredis.call('expire',KEYS[1],ARGV[2])" +
038            "\nend" +
039            "\nreturn c;";
040
041    private static JbootRedis redis;
042
043    /**
044     * 限制时长默认为1秒
045     */
046    public static boolean tryAcquire(String resource, int rate) {
047        return tryAcquire(resource, rate, 1);
048    }
049
050    /**
051     * 尝试是否能正常执行
052     *
053     * @param resource      资源名
054     * @param rate          限制次数
055     * @param periodSeconds 限制时长，单位为秒
056     * @return true 可以执行
057     * false 限次，禁止
058     */
059    public static boolean tryAcquire(String resource, int rate, int periodSeconds) {
060        if (redis == null) {
061            redis = JbootRedisManager.me().getRedis();
062            if (redis == null) {
063                throw new JbootIllegalConfigException("Redis config not well, can not use LimitScope.CLUSTER in @EnableLimit() ");
064            }
065        }
066        Long count = (Long) redis.eval(RATE_LIMIT_SCRIPT, 1, resource, String.valueOf(rate), String.valueOf(periodSeconds));
067        return count <= rate;
068    }
069}