Java常见限流用法介绍和实现

在现代分布式系统中，限流（Rate Limiting）是一种重要的流量控制手段，可以有效地保护后端服务不被过载。限流的应用场景非常广泛，例如API接口的调用、用户请求的限制等。在Java中，常见的限流实现方式主要包括：令牌桶算法、漏桶算法、计数器法、Redis的限流等。接下来，我们将详细介绍这些限流方式，并给出相应的代码示例。

1. 令牌桶算法

令牌桶算法是一种基于令牌的流量控制机制。系统以固定的速率向桶中放入令牌，当请求到达时，需从桶中获取令牌，若桶中有令牌，则允许处理请求，否则拒绝请求。

以下是一个简单的Java实现示例：

import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

public class TokenBucket {
    private final long capacity; // 桶的容量
    private final long rate; // 生成令牌的速率
    private long tokens; // 当前令牌数
    private long lastRefillTime; // 上次填充令牌的时间

    public TokenBucket(long capacity, long rate) {
        this.capacity = capacity;
        this.rate = rate;
        this.tokens = 0;
        this.lastRefillTime = System.currentTimeMillis();
    }

    public synchronized boolean getToken() {
        refillTokens();
        if (tokens > 0) {
            tokens--;
            return true; // 成功获取到令牌
        }
        return false; // 没有令牌可用
    }

    private void refillTokens() {
        long now = System.currentTimeMillis();
        long delta = now - lastRefillTime;
        long tokensToAdd = delta / 1000 * rate; // 每秒生成rate个令牌

        if (tokens + tokensToAdd > capacity) {
            tokens = capacity; // 桶满
        } else {
            tokens += tokensToAdd;
        }
        lastRefillTime = now;
    }

    public static void main(String[] args) throws InterruptedException {
        TokenBucket tokenBucket = new TokenBucket(10, 2); // 桶容量为10，速率为每秒2个令牌

        for (int i = 0; i < 20; i++) {
            if (tokenBucket.getToken()) {
                System.out.println("Request " + (i + 1) + " is allowed.");
            } else {
                System.out.println("Request " + (i + 1) + " is denied.");
            }
            TimeUnit.MILLISECONDS.sleep(500); // 模拟请求间隔
        }
    }
}

2. 漏桶算法

漏桶算法是一种相对平稳的流量控制机制。请求以不规则的速度进入桶中，桶以固定的速率流出请求。当桶满时，新请求将被拒绝。

import java.util.concurrent.TimeUnit;

public class LeakyBucket {
    private final double rate; // 固定的流出速率
    private double currentWater; // 当前水量
    private long lastCheckTime; // 上次检查时间

    public LeakyBucket(double rate) {
        this.rate = rate;
        this.currentWater = 0.0;
        this.lastCheckTime = System.currentTimeMillis();
    }

    public synchronized boolean addWater(double water) {
        leak();
        if (currentWater + water > 1.0) {
            return false; // 桶满，拒绝请求
        }
        currentWater += water;
        return true; // 成功加入水
    }

    private void leak() {
        long now = System.currentTimeMillis();
        double delta = (now - lastCheckTime) / 1000.0; // 转换为秒
        currentWater -= delta * rate; // 漏出水
        if (currentWater < 0) {
            currentWater = 0; // 水量不可为负
        }
        lastCheckTime = now;
    }

    public static void main(String[] args) throws InterruptedException {
        LeakyBucket leakyBucket = new LeakyBucket(0.5); // 每秒漏出0.5个单位的水

        for (int i = 0; i < 10; i++) {
            if (leakyBucket.addWater(0.2)) {
                System.out.println("Request " + (i + 1) + " is allowed.");
            } else {
                System.out.println("Request " + (i + 1) + " is denied.");
            }
            TimeUnit.MILLISECONDS.sleep(300); // 模拟请求间隔
        }
    }
}

3. 使用Redis实现限流

利用Redis的特性，我们也可以在分布式环境中实现限流。Redis的INCR命令可以用于计数，并且配合过期时间，可以轻松实现限流功能。

以下是基于Spring Boot的Redis限流实现示例：

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Service;

import java.util.concurrent.TimeUnit;

@Service
public class RateLimiterService {
    @Autowired
    private StringRedisTemplate redisTemplate;

    public boolean tryAcquire(String key, int limit, int periodInSeconds) {
        Long currentCount = redisTemplate.opsForValue().increment(key, 1);
        if (currentCount == 1) {
            redisTemplate.expire(key, periodInSeconds, TimeUnit.SECONDS); // 设置过期时间
        }
        return currentCount <= limit; // 判断是否超出限制
    }
}

在Controller中使用：

@RestController
public class ApiController {
    @Autowired
    private RateLimiterService rateLimiterService;

    @GetMapping("/api/request")
    public ResponseEntity<String> request() {
        String key = "api_request_count";
        int limit = 10; // 每10秒允许10个请求
        int periodInSeconds = 10;

        if (rateLimiterService.tryAcquire(key, limit, periodInSeconds)) {
            return ResponseEntity.ok("Request processed.");
        } else {
            return ResponseEntity.status(HttpStatus.TOO_MANY_REQUESTS).body("Too many requests.");
        }
    }
}

结论

限流是一种保护后端服务的有效手段，在Java中可以通过多种方式进行实现。无论是通过令牌桶算法、漏桶算法，还是使用Redis等工具，我们都可以根据具体的需求选择合适的限流策略。希望本文能为您理解限流的常用方式乃至实现提供一定的帮助。