package cn.reghao.tnb.sb.ratelimiter;

import lombok.Getter;
import lombok.Setter;

import java.util.Date;

/**
 * 令牌桶算法
 *
 * @author reghao
 * @date 2025-07-15 14:01:04
 */
public class TokenBucketRateLimiter {
    // 令牌桶容量
    private int capacity;
    // 令牌产生速率
    private int rate;
    // 令牌数量
    private int tokenAmount;

    public TokenBucketRateLimiter(int capacity, int rate) {
        this.capacity = capacity;
        this.rate = rate;
        tokenAmount = capacity;
        startThread();
    }

    private void startThread() {
        new Thread(new Runnable() {
            @Override
            public void run() {
                // 以恒定速率发放令牌
                while (true) {
                    synchronized (this) {
                        tokenAmount++;
                        if (tokenAmount > capacity) {
                            tokenAmount = capacity;
                        }
                    }
                    try {
                        Thread.sleep(1000/rate);
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                }
            }
        }).start();
    }

    public synchronized boolean tryAcquire(Request request) {
        if (tokenAmount > 0) {
            tokenAmount--;
            handleRequest(request);
            return true;
        } else {
            return false;
        }
    }

    // 处理请求
    private void handleRequest(Request request) {
        request.setHandleTime(new Date());
        System.out.println(request.getCode() + "号请求被处理，请求发起时间:"
                + request.getLaunchTime() + ",请求处理时间: " + request.getHandleTime() +
                ", 处理耗时:" + (request.getHandleTime().getTime() -
                request.getLaunchTime().getTime())+"ms");
    }

    @Setter
    @Getter
    static class Request {
        private int code;
        private Date launchTime;
        private Date handleTime;

        public Request(int code, Date launchTime) {
            this.launchTime = launchTime;
            this.code = code;
        }
    }

    public static void main(String[] args) throws InterruptedException {
        TokenBucketRateLimiter tokenBucketLimiter = new TokenBucketRateLimiter(5, 2);
        for (int i = 1; i <= 10; i++) {
            Request request = new Request(i, new Date());
            if (tokenBucketLimiter.tryAcquire(request)) {
                System.out.println(i + "号请求被接受");
            } else {
                System.out.println(i + "号请求被拒绝");
            }
        }
    }
}
