Go限制请求次数——令牌池

抓数据的网站限定1秒只能有10次请求,因此设计了一个令牌管理机制来控制请求数量。

设计思路如下:

  • 发请求前需要先获取令牌
  • 限定某时间段内的发放的令牌数量
  • 任务执行完成后不能归还令牌,只能使用定时器不断重置令牌
  • 如果当前goroutine数量过多时也不重置令牌
package main

import (
    "errors"
    "fmt"
    "sync"
    "sync/atomic"
    "time"
)

//节流器

type throttle struct {
    D      time.Duration //周期是D,
    C      int64         //限制一个周期最多操作C次
    Mu     sync.Mutex
    Token  chan bool //令牌池
    num    int64     //当前的goroutine数量
    maxNum int64     //允许工作goroutine最大数量
}

//如果两个周期后还没有申请到令牌,就报错超时
//目前用不到,如果限制routine最大数量需要靠这来监控
var ErrApplyTimeout = errors.New("apply token time out")

func NewThrottle(D time.Duration, C, maxNum int64) *throttle {
    instance := &throttle{
        D:      D,
        C:      C,
        Token:  make(chan bool, C),
        maxNum: maxNum,
    }
    go instance.reset()
    return instance
}

//每周期重新填充一次令牌池
func (t *throttle) reset() {
    ticker := time.NewTicker(t.D)
    for _ = range ticker.C {
        //goroutine数量不超过最大数量时再填充令牌池
        if t.num >= t.maxNum {
            continue
        }
        t.Mu.Lock()
        supply := t.C - int64(len(t.Token))
        fmt.Printf("reset token:%d\n", supply)
        for supply > 0 {
            t.Token <- true
            supply--
        }
        t.Mu.Unlock()
    }
}

//申请令牌,如果过两个周期还没申请到就报超时退出
func (t *throttle) ApplyToken() (bool, error) {
    select {
    case <-t.Token:
        return true, nil
    case <-time.After(t.D * 2):
        return false, ErrApplyTimeout
    }
}

func (t *throttle) Work(job func()) {
    if ok, err := t.ApplyToken(); !ok {
        fmt.Println(err)
        return
    }
    go func() {
        atomic.AddInt64(&t.num, 1)
        defer atomic.AddInt64(&t.num, -1)
        job()
    }()
}
func main() {
    t := NewThrottle(time.Second, 10, 20) //每秒10次,同时最多20个routine存在
    for {
        t.Work(doWork)
    }
}

//真正的工作函数 假设每个需要执行5秒
func doWork() {
    fmt.Println(time.Now())
    <-time.After(5 * time.Second)
}

以下是最初设计的,后来实际生产的时候并没有用上面的代码。而是用了另一种简陋实现

阅读原文

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
平台声明:文章内容(如有图片或视频亦包括在内)由作者上传并发布,文章内容仅代表作者本人观点,简书系信息发布平台,仅提供信息存储服务。

推荐阅读更多精彩内容

  • Spring Cloud为开发人员提供了快速构建分布式系统中一些常见模式的工具(例如配置管理,服务发现,断路器,智...
    卡卡罗2017阅读 134,993评论 19 139
  • 一、写在最前 轰轰烈烈的双十二已经过去小半个月了,程序猿的我坐在办公桌上思考,双十二这么大的访问量,这群电商是怎么...
    爱情小傻蛋阅读 8,012评论 0 13
  • Android 自定义View的各种姿势1 Activity的显示之ViewRootImpl详解 Activity...
    passiontim阅读 173,606评论 25 708
  • 今日头条旗下团队西瓜视频近日播出的直播节目《百万答题》,你参加了吗? 1 我大中国从不差人儿 近日,网络上很多直播...
    兰思阅读 280评论 0 0
  • #日记迎春20170321日行一善 感谢妈妈在这儿,我早起就有饭吃,孩子们也都收拾妥帖,只需送他俩出门就好,给了我...
    静俭阅读 196评论 0 0