漏桶算法与令牌桶算法

1
2
3
4
5
6
7
8

type Limiter struct {
	limit Limit // 向桶中放置令牌的速率
	burst int // 桶的容量
	mu     sync.Mutex
	tokens float64 // 可用令牌容量
	last time.Time // 上次放入token的时间
	lastEvent time.Time
}

1
2
3

func (lim *Limiter) AllowN(now time.Time, n int) bool {
	return lim.reserveN(now, n, 0).ok // 由于仅需要一个合规否，顾合规的通过，不合规的丢弃
}

 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

func (lim *Limiter) reserveN(now time.Time, n int, maxFutureReserve time.Duration) Reservation {
	lim.mu.Lock()

	if lim.limit == Inf {
		lim.mu.Unlock()
		return Reservation{
			ok:        true,
			lim:       lim,
			tokens:    n,
			timeToAct: now,
		}
	}
	// 这里拿到的是now，上次更新token时间和桶内token数量
	now, last, tokens := lim.advance(now)
	// 计算剩余的token
	tokens -= float64(n)

	// Calculate the wait duration
	var waitDuration time.Duration
	if tokens < 0 {
		waitDuration = lim.limit.durationFromTokens(-tokens)
	}

	// 确定是否合规，n是token
    // token 的数量要小于桶的容量，并且 等待时间小于最大等待时间
	ok := n <= lim.burst && waitDuration <= maxFutureReserve

	// Prepare reservation
	r := Reservation{
		ok:    ok,
		lim:   lim,
		limit: lim.limit,
	}
	if ok {
		r.tokens = n
		r.timeToAct = now.Add(waitDuration)
	}

	// Update state
	if ok {
		lim.last = now
		lim.tokens = tokens
		lim.lastEvent = r.timeToAct
	} else {
		lim.last = last
	}

	lim.mu.Unlock()
	return r
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

func (lim *Limiter) advance(now time.Time) (newNow time.Time, newLast time.Time, newTokens float64) {
   last := lim.last
   if now.Before(last) { // 计算上次放入token是否在传入now之前
      last = now
   }

   // 当 last 很旧时，避免在下面进行 delta 溢出。
   // maxElapsed 计算装满需要多少时间
   maxElapsed := lim.limit.durationFromTokens(float64(lim.burst) - lim.tokens)
   elapsed := now.Sub(last) // 上次装入到现在的时差
   if elapsed > maxElapsed { // 上次如果放入token时间超长，就让他与装满时间相等
      elapsed = maxElapsed // 即，让桶为满的
   }

   // 装桶的动作，下面函数表示，elapsed时间内可以生成多少个token
   delta := lim.limit.tokensFromDuration(elapsed)
   tokens := lim.tokens + delta // 当前的token
   if burst := float64(lim.burst); tokens > burst {
      tokens = burst // 这里表示token溢出，让他装满就好
   }

   return now, last, tokens
}

 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

func (lim *Limiter) WaitN(ctx context.Context, n int) (err error) {
	if n > lim.burst && lim.limit != Inf {
		return fmt.Errorf("rate: Wait(n=%d) exceeds limiter's burst %d", n, lim.burst)
	}
	// 外部已取消
	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
	}
	// Determine wait limit
	now := time.Now()
	waitLimit := InfDuration
	if deadline, ok := ctx.Deadline(); ok {
		waitLimit = deadline.Sub(now)
	}
	// 三个方法的核心，这里给定了deatline
	r := lim.reserveN(now, n, waitLimit)
	if !r.ok {
		return fmt.Errorf("rate: Wait(n=%d) would exceed context deadline", n)
	}
	// Wait if necessary
	delay := r.DelayFrom(now)
	if delay == 0 {
		return nil
	}
	t := time.NewTimer(delay)
	defer t.Stop()
	select {
	case <-t.C:
		// We can proceed.
		return nil
	case <-ctx.Done():
		// Context was canceled before we could proceed.  Cancel the
		// reservation, which may permit other events to proceed sooner.
		r.Cancel()
		return ctx.Err()
	}
}

 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

package main

import (
	"log"
	"strconv"
	"time"

	"golang.org/x/time/rate"
)

func main() {
	timeLayout := "2006-01-02:15:04:05.0000"
	limiter := rate.NewLimiter(1, 5) // BT(1,5)
	log.Println("bucket current capacity: " + strconv.Itoa(limiter.Burst()))
	length := 20 // 一共请求20次
	chs := make([]chan string, length)
	for i := 0; i < length; i++ {
		chs[i] = make(chan string, 1)
		go func(taskId string, ch chan string, r *rate.Limiter) {
			err := limiter.Allow()
			if !err {
				ch <- "Task-" + taskId + " unallow " + time.Now().Format(timeLayout)
			}

			time.Sleep(time.Duration(5) * time.Millisecond)
			ch <- "Task-" + taskId + " run success  " + time.Now().Format(timeLayout)
			return

		}(strconv.FormatInt(int64(i), 10), chs[i], limiter)
	}
	for _, ch := range chs {
		log.Println("task start at " + <-ch)
	}
}

 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

package main

import (
	"context"
	"log"
	"strconv"
	"time"

	"golang.org/x/time/rate"
)

func main() {
	timeLayout := "2006-01-02:15:04:05.0000"
	limiter := rate.NewLimiter(1, 5) // BT(1,5)
	log.Println("bucket current capacity: " + strconv.Itoa(limiter.Burst()))
	length := 20 // 一共请求20次
	chs := make([]chan string, length)
	for i := 0; i < length; i++ {
		chs[i] = make(chan string, 1)
		go func(taskId string, ch chan string, r *rate.Limiter) {
			err := limiter.Wait(context.TODO())
			if err != nil {
				ch <- "Task-" + taskId + " unallow " + time.Now().Format(timeLayout)
			}
			ch <- "Task-" + taskId + " run success  " + time.Now().Format(timeLayout)
			return

		}(strconv.FormatInt(int64(i), 10), chs[i], limiter)
	}
	for _, ch := range chs {
		log.Println("task start at " + <-ch)
	}
}