Gitlib

Enjoy Yourself, Enjoy Life.

0%

基于Redis使用令牌桶算法实现流量控制

发表于 分类于

令牌桶

令牌桶算法是常见的限流算法，用来控制发送到网络上的数据的数目，并允许突发数据的发送，其原理也很简单：

  1. 首先设有一个令牌桶，桶内存放令牌，一开始令牌桶内的令牌是满的（桶内令牌的数量可根据服务器情况设定）;
  2. 每次访问从桶内取走一个令牌，当桶内令牌为0，则不允许再访问;
  3. 每隔一段时间，再放入令牌，最多使桶内令牌满额。（可以根据实际情况，每隔一段时间放入若干个令牌，或直接补满令牌桶）;

我们可以使用redis的队列作为令牌桶容器使用，使用lPush（入队），rPop（出队），实现令牌加入与消耗的操作， 代码如下：

TokenBucket.class.php

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
<?php

class TokenBucket{ 

    private $_config; // redis设定
    private $_redis;  // redis对象
    private $_queue;  // 令牌桶
    private $_max;    // 最大令牌数

    /**
     * 初始化
     * @param Array $config redis连接设定
     */
    public function __construct($config, $queue, $max){
        $this->_config = $config;
        $this->_queue = $queue;
        $this->_max = $max;
        $this->_redis = $this->connect();
    }

    /**
     * 加入令牌
     * @param  Int $num 加入的令牌数量
     * @return Int 加入的数量
     */
    public function add($num=0){

        // 当前剩余令牌数
        $curnum = intval($this->_redis->lSize($this->_queue));

        // 最大令牌数
        $maxnum = intval($this->_max);

        // 计算最大可加入的令牌数量，不能超过最大令牌数
        $num = $maxnum>=$curnum+$num? $num : $maxnum-$curnum;

        // 加入令牌
        if($num>0){
            $token = array_fill(0, $num, 1);
            $this->_redis->lPush($this->_queue, ...$token);
            return $num;
        }

        return 0;

    }

    /**
     * 获取令牌
     * @return Boolean
     */
    public function get(){
        return $this->_redis->rPop($this->_queue)? true : false;
    }

    /**
     * 重设令牌桶，填满令牌
     */
    public function reset(){
        $this->_redis->delete($this->_queue);
        $this->add($this->_max);
    }

    /**
     * 创建redis连接
     * @return Link
     */
    private function connect(){
        try{
            $redis = new Redis();
            $redis->connect($this->_config['host'],$this->_config['port']);
        }catch(RedisException $e){
            throw new Exception($e->getMessage());
            return false;
        }
        return $redis;
    }


}
?>

traffic.limit.php

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
<?php


require '../common.php';
require 'TokenBucket.class.php';


// redis连接设定
$config = array(
    'host' => '127.0.0.1',
    'port' => 6379
);

// 令牌桶容器
$queue = 'traffic';

// 最大令牌数
$max = 5;

// 创建TrafficShaper对象
$tokenBucket = new TokenBucket($config, $queue, $max);

// 重设令牌桶，填满令牌
$tokenBucket->reset();

// 循环获取令牌，令牌桶内只有5个令牌，因此最后3次获取失败
for ($i=0; $i<8; $i++) {
    if($tokenBucket->get()) {
    	print_log('请求执行');
    } else {
    	print_log('请求Discard');
    }
}

?>

输出：

1
2
3
4
5
6
7
8
9
gitlib@devops:/www/websites/www/gitlib/practise/php/redis$ php traffic.limit.test.php 
2019-09-08 14:42:28 请求执行
2019-09-08 14:42:28 请求执行
2019-09-08 14:42:28 请求执行
2019-09-08 14:42:28 请求执行
2019-09-08 14:42:28 请求执行
2019-09-08 14:42:28 请求Discard
2019-09-08 14:42:28 请求Discard
2019-09-08 14:42:28 请求Discard

至于定期加入令牌，可以使用crontab实现，定时调用add方法加入若干令牌即可。