Golang token bucket-frequency limit

Token bucket algorithm

　　The token bucket algorithm is generally used as frequency limit, flow limit, etc., may have a single-speed two-color, single-speed three-color, two-speed three-color method.

Our specific requirement is to limit the frequency of calls to the API, so we achieve a single-speed two-color.

Package Mainimport ("Errors" "FMT" "StrConv" "Sync" "Time") const Token_granularity = 1 000type MAP struct {lock sync.        Rwmutex Bucket map[string]*tokenbucket}//can implement a lock-free algorithm Todofunc (M *map) Set (k string, times, interval int) {                M.lock.lock () defer m.lock.unlock () If _, OK: = M.bucket[k];!ok {tb: = new (Tokenbucket) Tb. Init (times, interval) m.bucket[k] = tb} return}func Atoi (s string) int {I, _: = Strc Onv. Atoi (s) return I}func Newmap () *map {return &map{bucket:make (map[string]*tokenbucket)}}type Tokenbucke t struct {lastquery time. Time tokens int burst int step int add int mu sync.        Mutex}func (TB *tokenbucket) Init (quota int, interval int) {Tb.burst = quota * token_granularity//maximum reserved tokens Tb.tokens = quota * token_granularity//current reserved Token tb.step =Quota * Token_granularity/quota//per consumption of tokens Tb.add = quota * token_granularity * interval/interval//per second tokens added Tb.lastquery = time. Now ()}func (TB *tokenbucket) Tokenbucketquery () error {now: = time. Now () diff: = now. Sub (tb.lastquery) Token: = Int (diff. nanoseconds ()/1000000000) * Tb.add Tb.mu.Lock () defer tb.mu.Unlock () if token! = 0 {T                B.lastquery = Now Tb.tokens + = token//Add token for this time} if Tb.tokens > Tb.burst {//exceeds maximum token number reset Tb.tokens = Tb.burst} if Tb.tokens >= tb.step {//vs. per consumption tb.tokens-= t B.step return nil} return errors. New ("Not Enough")}//@testfunc Main () {var TB tokenbucket TB. Init (5, 1) calls 5 cnt in//1s: = 0 for {err: = tb. Tokenbucketquery () if err! = Nil {fmt.                   PRINTLN (Err)} else {     Fmt. Println ("Take")} cnt + = 1 FMT. PRINTLN (CNT) time. Sleep (100000 * time. microsecond)//time. Sleep (5* time. Second)}}

Test results:

　　　　

Now the implementation is the need to lock to ensure multi-threaded security, do not know whether there is a lock-free implementation, pending research

Golang token bucket-frequency limit