com.feingto.cloud.cache.provider.RedisTokenProvider Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of feingto-core Show documentation
There is a newer version: 2.3.8.RELEASE
package com.feingto.cloud.cache.provider;

import com.feingto.cloud.core.json.JSON;
import com.feingto.cloud.domain.type.IntervalUnit;
import lombok.*;
import org.springframework.data.redis.core.BoundHashOperations;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.util.StringUtils;

import java.util.HashMap;
import java.util.Map;

/**
 * redis令牌桶算法
 *
 * @author longfei
 */
@Builder
public class RedisTokenProvider {
    public static final String CACHE_API_LIMIT = "REDIS_TOKEN_LIMIT";

    private final RedisTemplate redisTemplate;

    private String keyPrefix;//Redis key 前缀

    @Setter
    private String principal;// 身份标识

    private long frequency;// 频率

    private IntervalUnit intervalUnit;// 单位时间（分钟、小时等）

    @Setter
    private long limit;// 次数

    /**
     * 验证principal在单位时间内是否达到次数限制
     *
     * @return true: limit; false: access
     */
    public synchronized boolean isLimit() {
        long frequencyInMills;
        switch (intervalUnit) {
            case MILLISECONDS:
                frequencyInMills = frequency;
                break;
            case SECONDS:
                frequencyInMills = frequency * 1000;
                break;
            case MINUTES:
                frequencyInMills = frequency * 60000;
                break;
            case HOURS:
                frequencyInMills = frequency * 360000;
                break;
            case DAYS:
                frequencyInMills = frequency * 8640000;
                break;
            default:
                return false;// 不受限
        }
        double intervalPerPermit = frequencyInMills * 1.0 / limit;
        String key = this.generateKey();
        BoundHashOperations opts = redisTemplate.boundHashOps(CACHE_API_LIMIT);
        Object token = opts.get(key);
        if (token == null) {
            opts.put(key, TokenBucket.builder().lastRefillTime(System.currentTimeMillis()).tokensRemaining(limit - 1).build().toString());
            return false;
        } else {
            TokenBucket tokenBucket = TokenBucket.fromString((String) token);
            long lastRefillTime = tokenBucket.getLastRefillTime();
            long refillTime = System.currentTimeMillis();
            long intervalSinceLast = refillTime - lastRefillTime;
            long currentTokensRemaining;
            if (intervalSinceLast > frequencyInMills) {
                currentTokensRemaining = limit;
            } else {
                long grantedTokens = (long) (intervalSinceLast / intervalPerPermit);
                currentTokensRemaining = Math.min(grantedTokens + tokenBucket.getTokensRemaining(), limit);
            }
            if (currentTokensRemaining == 0) {
                tokenBucket.setTokensRemaining(currentTokensRemaining);
                opts.put(key, tokenBucket.toString());
                return true;
            } else {
                tokenBucket.setLastRefillTime(refillTime);
                tokenBucket.setTokensRemaining(currentTokensRemaining - 1);
                opts.put(key, tokenBucket.toString());
                return false;
            }
        }
    }

    /**
     * 重置令牌
     */
    public synchronized void clear() {
        redisTemplate.boundHashOps(CACHE_API_LIMIT).delete(this.generateKey());
    }

    private String generateKey() {
        if (StringUtils.isEmpty(keyPrefix))
            keyPrefix = "flow:limit:";
        return keyPrefix + frequency + ":" + intervalUnit + ":" + limit + ":" + principal;
    }

    @Builder
    @Data
    @NoArgsConstructor
    @AllArgsConstructor
    private static class TokenBucket {
        private long lastRefillTime;
        private long tokensRemaining;

        static TokenBucket fromHash(Map hash) {
            long lastRefillTime = Long.parseLong(hash.get("lastRefillTime"));
            int tokensRemaining = Integer.parseInt(hash.get("tokensRemaining"));
            return new TokenBucket(lastRefillTime, tokensRemaining);
        }

        static TokenBucket fromString(String json) {
            return JSON.build().json2pojo(json, TokenBucket.class);
        }

        Map toHash() {
            Map hash = new HashMap<>();
            hash.put("lastRefillTime", String.valueOf(lastRefillTime));
            hash.put("tokensRemaining", String.valueOf(tokensRemaining));
            return hash;
        }

        public String toString() {
            return JSON.build().obj2json(this);
        }
    }
}