【Java】RedisTemplate来控制某段时间内执行最大次数

一、前言

RedisTemplate来控制某段时间内执行最大次数的Java示例代码。

二、使用场景

假设我们希望限制某个操作（如发送邮件、访问特定API等）在1小时内最多执行n次。

三、思路

1. 使用redis的zset数据结构，将当前时间戳作为分值。
2. 每次发送短信时，先删除1个小时以前的元素(即分值从0，一个小时前的时间戳)。
3. 删除过期的元素后再统计，key中的元素个数(1个小时内的发送次数)，看是否大于6次。
4. 发送成功后，将当前时间戳做为值和分值保存在zset中。

四、代码示例

一个小时内同一用户最多发送6次短信

import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ZSetOperations.TypedTuple;
import org.springframework.stereotype.Component;

import java.util.Date;
import java.util.Set;

import static java.util.stream.Collectors.toList;

@Component
public class RateLimiter {

    private static final String KEY_PREFIX = "rate_limiter:send_msg:";
    private static final long INTERVAL_MS = TimeUnit.HOURS.toMillis(1);
    private static final int MAX_COUNT = 6;

    private final RedisTemplate<String, Object> redisTemplate;

    public RateLimiter(RedisTemplate<String, Object> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    public boolean canPerformOperation(String userId) {
        String key = KEY_PREFIX + userId;

        // 获取当前时间戳（毫秒）
        long now = System.currentTimeMillis();

        // 删除已过期的计数
        long startOfInterval = now - INTERVAL_MS;
        Set<TypedTuple<Object>> expiredCounts = redisTemplate.opsForZSet().rangeByScoreWithScores(key, 0, startOfInterval);
        if (!expiredCounts.isEmpty()) {
            List<Object> expiredKeys = expiredCounts.stream().map(TypedTuple::getValue).collect(toList());
            redisTemplate.opsForZSet().removeRangeByScore(key, 0, startOfInterval);
            System.out.println("Removed " + expiredCounts.size() + " expired counts");
        }

        // 获取剩余次数
        long currentCount = redisTemplate.opsForZSet().zCard(key);
        long remainingCount = MAX_COUNT - currentCount;

        // 如果剩余次数大于0，则添加新的计数项并返回true
        if (remainingCount > 0) {
            redisTemplate.opsForZSet().add(key, now, now);
            return true;
        } else {
            return false;
        }
    }

    // 示例用法：
    public static void main(String[] args) {
        RateLimiter rateLimiter = ...; // 初始化RateLimiter实例
        String userId = "100000";
        if (rateLimiter.canPerformOperation( userId )) {
            System.out.println("Sending msg...");
        } else {
            System.out.println("Rate limit exceeded. Unable to send msg.");
        }
    }
}