程序层面限流 - 图解四种限流方法实现（Golang版）

如果❤️我的文章有帮助，欢迎点赞、关注。这是对我继续技术创作最大的鼓励。

本文为 应对突防流量 系列文章，以下为相关基友篇

• 面对突发流量：系统、软件层限流手段
• 程序层面限流 - 必要性
• 程序层面限流 - 图解四种限流方法实现（Golang版）

计数器限流

package main

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

type RequestLimiter struct {
	Interval time.Duration // 重新计数时间
	MaxCount int           // 最大计数
	Lock     sync.Mutex
	ReqCount int // 目前的请求数
}

func (reqLimiter *RequestLimiter) IsAvailable() bool {
	reqLimiter.Lock.Lock()
	defer reqLimiter.Lock.Unlock()

	return reqLimiter.ReqCount < reqLimiter.MaxCount
}

// 非阻塞
func (reqLimiter *RequestLimiter) AddRequestCount() bool {
	reqLimiter.Lock.Lock()
	defer reqLimiter.Lock.Unlock()
	if reqLimiter.ReqCount < reqLimiter.MaxCount {
		reqLimiter.ReqCount += 1
		return true
	}
	return false
}

func NewRequestLimitService(interval time.Duration, maxCount int) *RequestLimiter {
	reqLimit := &RequestLimiter{
		Interval: interval,
		MaxCount: maxCount,
	}
	go func() {
		ticker := time.NewTicker(interval)
		for true {
			<-ticker.C
			reqLimit.Lock.Lock()
			fmt.Println("reset Count ...")
			reqLimit.ReqCount = 0
			reqLimit.Lock.Unlock()
		}
	}()

	return reqLimit
}

/**
	原理：协程不阻塞 + 死循环 + ticker 定时执行 ReqCount 归零
	ticker := time.NewTicker(2*time.Second)
	for {
		currentTime := <-ticker.C
		fmt.Println("当前时间为:", currentTime)
	}

	输出：
		当前时间为: 2021-05-19 11:02:12.3603475 +0800 CST m=+2.002147201
		当前时间为: 2021-05-19 11:02:14.3611806 +0800 CST m=+4.002980301
		当前时间为: 2021-05-19 11:02:16.360625 +0800 CST m=+6.002424701
 */
func main() {
	service := NewRequestLimitService(time.Second, 2)
	for true {
		hasToken := service.AddRequestCount()
		if hasToken {
			fmt.Println(time.Now())
		}
	}
}

滑动窗口

package main

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

type WindowLimiter struct {
	Interval    time.Duration // 总计数时间
	WinCount    []int         // 每个窗口的访问数量
	TicketSize  int           // 窗口最大容量
	TicketCount int           // 窗口数量
	Lock        sync.Mutex
	CurIndex    int // 目前使用哪个窗口
}

func (reqLimiter *WindowLimiter) IsAvailable() bool {
	reqLimiter.Lock.Lock()
	defer reqLimiter.Lock.Unlock()
	return reqLimiter.WinCount[reqLimiter.CurIndex] < reqLimiter.TicketSize
}

func (reqLimiter *WindowLimiter) AddRequestCount() bool {
	reqLimiter.Lock.Lock()
	defer reqLimiter.Lock.Unlock()
	if reqLimiter.WinCount[reqLimiter.CurIndex] < reqLimiter.TicketSize {
		reqLimiter.WinCount[reqLimiter.CurIndex]++
		return true
	}
	return false
}

func NewRequestLimitService(interval time.Duration, ticketCount int, ticketSize int) *WindowLimiter {
	reqLimit := &WindowLimiter{
		Interval:    interval,
		WinCount:    make([]int, ticketCount, ticketCount),
		TicketSize:  ticketSize,
		TicketCount: ticketCount,
		CurIndex:    0,
	}
	go func() {
		ticker := time.NewTicker(time.Duration(interval.Nanoseconds() / int64(ticketCount)))
		for true {
			<-ticker.C
			reqLimit.Lock.Lock()
			reqLimit.CurIndex = (reqLimit.CurIndex + 1) % reqLimit.TicketCount
			reqLimit.WinCount[reqLimit.CurIndex] = 0
			fmt.Println("reset Count ...")
			reqLimit.Lock.Unlock()
		}
	}()

	return reqLimit
}

func main() {
	service := NewRequestLimitService(time.Second, 2, 1)
	for true {
		hasToken := service.AddRequestCount()
		if hasToken {
			fmt.Println(time.Now())
		}
	}
}

漏桶算法

package main

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

type BucketLimiter struct {
	Timestamp time.Time // 当前注水的时间戳
	Capacity  float64   // 桶的容量
	Rate      float64   // 速度
	Water     float64   // 当前水量
	Lock      sync.Mutex
}

func AddWater(bucket *BucketLimiter) bool {
	now := time.Now()
	leftWater := math.Max(0, bucket.Water-now.Sub(bucket.Timestamp).Seconds()*bucket.Rate)
	bucket.Lock.Lock()
	defer bucket.Lock.Unlock()
	if leftWater+1 < bucket.Capacity {
		// 尝试加水，此时水桶未满
		bucket.Timestamp = now
		bucket.Water = leftWater + 1
		return true
	} else {
		// 水满了，拒绝访问
		return false
	}
}

func main() {
	service := &BucketLimiter{
		Timestamp: time.Now(),
		Capacity:  2,
		Rate:      1,
		Water:     0,
	}
	for true {
		hasToken := AddWater(service)
		if hasToken {
			fmt.Println(time.Now())
		}

	}
}

令牌桶算法

package main

import (
	"math"
	"sync"
	"time"
)

// 定义令牌桶结构
type tokenBucket struct {
	timestamp time.Time // 当前时间戳
	capacity  float64   // 桶的容量（存放令牌的最大量）
	rate      float64   // 令牌放入速度
	tokens    float64   // 当前令牌总量
	lock      sync.Mutex
}

// 判断是否获取令牌（若能获取，则处理请求）
func getToken(bucket tokenBucket) bool {
	now := time.Now()
	bucket.lock.Lock()
	defer bucket.lock.Unlock()
	// 先添加令牌
	leftTokens := math.Max(bucket.capacity, bucket.tokens+now.Sub(bucket.timestamp).Seconds()*bucket.rate)
	if leftTokens < 1 {
		// 若桶中一个令牌都没有了，则拒绝
		return false
	} else {
		// 桶中还有令牌，领取令牌
		bucket.tokens -= 1
		bucket.timestamp = now
		return true
	}
}