使用redis实现一个令牌桶算法

import redis.clients.jedis.Jedis;

public class RedisTokenBucket {
    private final Jedis jedis;
    private final String key;
    private final long capacity; // 令牌桶容量
    private final double refillRate; // 令牌填充速率，单位tokens/秒

    public RedisTokenBucket(Jedis jedis, String key, long capacity, double refillRate) {
        this.jedis = jedis;
        this.key = key;
        this.capacity = capacity;
        this.refillRate = refillRate;
    }

    public boolean acquire(long tokens) {
        // 获取当前令牌数量
        String tokensStr = jedis.get(key);
        double currentTokens = tokensStr != null ? Double.parseDouble(tokensStr) : 0.0;

        // 计算令牌桶剩余容量
        double remainingCapacity = capacity - currentTokens;

        // 计算需要等待的时间（毫秒）
        long waitTime = (long) (tokens / refillRate * 1000);

        // 判断是否有足够的令牌可用
        if (tokens <= remainingCapacity) {
            // 更新令牌数量并设置过期时间
            jedis.incrByFloat(key, tokens);
            if (!jedis.exists(key)) {
                jedis.expire(key, (int) (capacity / refillRate));
            }
            return true;
        } else {
            // 没有足够的令牌可用，需要等待
            try {
                Thread.sleep(waitTime);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
            return false;
        }
    }

    public static void main(String[] args) {
        Jedis jedis = new Jedis("localhost", 6379);
        RedisTokenBucket tokenBucket = new RedisTokenBucket(jedis, "my_bucket", 10, 0.5); // 10个容量，每秒填充0.5个令牌

        // 测试获取令牌
        for (int i = 0; i < 20; i++) {
            boolean acquired = tokenBucket.acquire(1);
            System.out.println("Token acquired: " + acquired);
        }

        jedis.close();
    }
}
 

创建了一个名为 RedisTokenBucket 的类来表示Redis令牌桶。构造函数接受一个 Jedis 对象、一个键（用于在Redis中标识令牌桶）、令牌桶的容量和填充速率作为参数。

acquire 方法尝试获取指定数量的令牌。它首先从Redis中获取当前令牌数量，然后根据剩余容量和填充速率计算需要等待的时间。如果有足够的令牌可用，它将更新令牌数量并设置过期时间。如果没有足够的令牌可用，则会等待一段时间后再次尝试。

在 main 方法中，我们创建了一个Redis连接，并使用 RedisTokenBucket 类来模拟获取令牌的过程