Go Channel Practice IV - Token Bucket

https://en.wikipedia.org/wiki/Token_bucket

先定义 TokenBucket

type TokenBucket struct {
	Depth  int
	Tokens int
	Delta  int
	ticker *time.Ticker
	mutex  sync.Mutex
	quit   chan bool
}

Depth 是 bucket 的总深度，Tokens 是当前可以用到的数目，Delta 是每秒向 bucket 增加的 token 数目。ticker, mutex 和 quit 类似 Leaky Bucket 的用法。

NewTokenBucket 创建一个 goroutine，不断在每个 tick 里向 token 里加 delta 个 token。

func NewTokenBucket(depth, delta int) *TokenBucket {
	ticker := time.NewTicker(time.Second)
	tb := &TokenBucket{
		Depth:  depth,
		Tokens: depth,
		Delta:  delta,
		ticker: ticker,
		mutex:  sync.Mutex{},
		quit:   make(chan bool),
	}

	go func() {
		for {
			select {
			case <-ticker.C:
				fmt.Println("tick...")
				tb.mutex.Lock()
				if tb.Tokens+tb.Delta > tb.Depth {
					tb.Tokens = tb.Depth
				} else {
					tb.Tokens += tb.Delta
				}
				tb.mutex.Unlock()
			case <-tb.quit:
				return
			}
		}
	}()

	return tb
}

ticker 和 goroutine 也需要 close

func (t *TokenBucket) Close() {
	t.ticker.Stop()
	t.quit <- true
}

创建一个 token bucket 之后调用，判断每次取的 token 是不是大于当前 bucket 内的 token，如果大于，则失败，否则成功。

func (t *TokenBucket) Withdraw(count int) bool {
	t.mutex.Lock()
	defer t.mutex.Unlock()
	if count > t.Tokens {
		fmt.Println("failed to withdraw")
		return false
	}
	t.Tokens -= count
	fmt.Println("withdraw succeeded")
	return true
}

完整测试如下

package main

import (
	"fmt"
	"sync"
	"time"
)

func main() {
	tb := NewTokenBucket(1000, 300)
	for i := 0; i < 300; i++ {
		if !tb.Withdraw(100) {
			time.Sleep(time.Millisecond * 100)
		}
	}
}

type TokenBucket struct {
	Depth  int
	Tokens int
	Delta  int
	ticker *time.Ticker
	mutex  sync.Mutex
	quit   chan bool
}

func NewTokenBucket(depth, delta int) *TokenBucket {
	ticker := time.NewTicker(time.Second)
	tb := &TokenBucket{
		Depth:  depth,
		Tokens: depth,
		Delta:  delta,
		ticker: ticker,
		mutex:  sync.Mutex{},
		quit:   make(chan bool),
	}

	go func() {
		for {
			select {
			case <-ticker.C:
				fmt.Println("tick...")
				tb.mutex.Lock()
				fmt.Println(tb.Tokens)
				if tb.Tokens+tb.Delta > tb.Depth {
					tb.Tokens = tb.Depth
				} else {
					tb.Tokens += tb.Delta
				}
				fmt.Println(tb.Tokens)
				tb.mutex.Unlock()
			case <-tb.quit:
				return
			}
		}
	}()

	return tb
}

func (t *TokenBucket) Close() {
	t.ticker.Stop()
	t.quit <- true
}

func (t *TokenBucket) Withdraw(count int) bool {
	t.mutex.Lock()
	defer t.mutex.Unlock()
	if count > t.Tokens {
		fmt.Println("failed to withdraw")
		return false
	}
	t.Tokens -= count
	fmt.Println("withdraw succeeded")
	return true
}

• 相关文章
• 全局相关文章