一、概述

令牌桶是一种常用的流量控制技术。令牌桶本身没有丢弃和优先级策略。

原理

1．令牌以一定的速率放入桶中。

2．每个令牌允许源发送一定数量的比特。

3．发送一个包，流量调节器就要从桶中删除与包大小相等的令牌数。

4．如果没有足够的令牌发送包，这个包就会等待直到有足够的令牌（在整形器的情况下）或者包被丢弃，也有可能被标记更低的DSCP（在策略者的情况下）。

5．桶有特定的容量，如果桶已经满了，新加入的令牌就会被丢弃。因此，在任何时候，源发送到网络上的最大突发数据量与桶的大小成比例。令牌桶允许突发，但是不能超过限制。

二、原理图

在这里插入图片描述

三、令牌桶工具类

简单模拟；

TreeSet 桶

package com.tuwer.util;

import java.time.Instant;
import java.util.*;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

/**
 * <p>令牌桶</p>
 *
 * @author 土味儿
 * @version 1.0
 * @Date 2023/6/14
 */
public class TokenBucket {
    /**
     * 速率（每秒）
     */
    private int rate;
    /**
     * 令牌桶
     */
    private TreeSet<String> bucket;

    public TokenBucket(int rate) {
        this.rate = Math.max(rate, 1);
        this.bucket = new TreeSet<>();
        this.put();
    }

    /**
     * 获取令牌
     *
     * @return
     */
    public boolean getToken() {
        // 从桶中弹出第一个令牌；如果弹出成功为true
        return Objects.nonNull(this.bucket.pollFirst());
    }

    /**
     * 补充令牌
     * 使桶内令牌保持与速率一致
     */
    private void put() {
        // 开启新线程
        Executors.newSingleThreadExecutor().execute(() -> {
            int num;
            // 匀速补充令牌
            while (true) {
                // 需要补充的令牌数量
                num = this.rate - this.bucket.size();
                this.bucket.addAll(createToken(num));

                // 休眼1秒
                try {
                    TimeUnit.SECONDS.sleep(1);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        });
    }

    /**
     * 生成num个令牌
     *
     * @param num
     * @return
     */
    private Set<String> createToken(int num) {
        Set<String> res = new HashSet<>();
        if (num < 1) {
            return res;
        }

        // 当前时间戳
        long second = Instant.now().getEpochSecond();
        for (int i = 1; i < num + 1; i++) {
            res.add(second + "-" + i);
        }
        return res;
    }
}

List 桶

package com.tuwer.util;

import java.time.Instant;
import java.util.*;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

/**
 * <p>令牌桶</p>
 *
 * @author 土味儿
 * @version 1.0
 * @Date 2023/6/14
 */
public class TokenBucketOfList {
    /**
     * 速率（每秒）
     */
    private int rate;
    /**
     * 令牌桶
     */
    private List<String> bucket;

    public TokenBucketOfList(int rate) {
        this.rate = Math.max(rate, 1);
        this.bucket = new ArrayList<>(this.rate);
        this.put();
    }

    /**
     * 获取令牌
     *
     * @return
     */
    public boolean getToken() {
        if(bucket.isEmpty()){
            return false;
        }
        // 从桶中弹出第一个令牌
        bucket.remove(0);
        return true;
    }

    /**
     * 补充令牌
     * 使桶内令牌保持与速率一致
     */
    private void put() {
        // 开启新线程
        Executors.newSingleThreadExecutor().execute(() -> {
            int num;
            // 匀速补充令牌
            while (true) {
                // 需要补充的令牌数量
                num = this.rate - this.bucket.size();
                this.bucket.addAll(createToken(num));

                // 休眼1秒
                try {
                    TimeUnit.SECONDS.sleep(1);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        });
    }

    /**
     * 生成num个令牌
     *
     * @param num
     * @return
     */
    private Set<String> createToken(int num) {
        Set<String> res = new HashSet<>();
        if (num < 1) {
            return res;
        }

        // 当前时间戳
        long second = Instant.now().getEpochSecond();
        for (int i = 1; i < num + 1; i++) {
            res.add(second + "-" + i);
        }
        return res;
    }
}

四、模拟API接口

1、服务1

每秒可以处理5个请求

package com.tuwer.service;

import com.tuwer.util.TokenBucket;

import java.util.ArrayList;
import java.util.List;

/**
 * <p>服务类</p>
 *
 * @author 土味儿
 * @version 1.0
 * @Date 2023/6/14
 */
public class MyService1 {
    private static TokenBucket tokenBucket = new TokenBucket(5);

    public List<String> userList(){
        if(tokenBucket.getToken()){
            List<String> list = new ArrayList<>();
            list.add("张三");
            list.add("李四");
            list.add("王五");
            list.add("赵六");
            System.out.println(list);
            return list;
        }else{
            System.out.println("\nMyService1 服务忙，请稍候再试！\n");
            return null;
        }
    }
}

2、服务2

每秒可以处理3个请求

package com.tuwer.service;

import com.tuwer.util.TokenBucket;
import com.tuwer.util.TokenBucketOfList;

import java.util.ArrayList;
import java.util.List;

/**
 * <p>服务类</p>
 *
 * @author 土味儿
 * @version 1.0
 * @Date 2023/6/14
 */
public class MyService2 {
    private static TokenBucket tokenBucket = new TokenBucket(3);

    public List<String> userList(){
        if(tokenBucket.getToken()){
            List<String> list = new ArrayList<>();
            list.add("AAA");
            list.add("BBB");
            list.add("CCC");
            System.out.println(list);
            return list;
        }else{
            System.out.println("\nMyService2 服务忙，请稍候再试！\n");
            return null;
        }
    }
}

五、测试

向两个服务分别发送30个请求，每个请求间隔150毫秒

public class MyTest {
    public static void main(String[] args) {
        MyService1 myService1 = new MyService1();
        MyService2 myService2 = new MyService2();
        int num = 30;
        for (int i = 0; i < num; i++) {
            try {
                TimeUnit.MILLISECONDS.sleep(150);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }

            myService1.userList();
            myService2.userList();
        }
    }
}