golang限流

golang限流其实官方提供了golang.org/x/time/rate包，直接rate.NewLimiter(2, 5)，然后判断

func limit(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        if limiter.Allow() == false {
            http.Error(w, http.StatusText(429), http.StatusTooManyRequests)
            return
        }
        next.ServeHTTP(w, r)
    })
}

自己也可以简单实现一下go的限流算法

简单计数器

// 简单计数器
package main

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

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

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
}

滑动窗口

// 滑动窗口
package main

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

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

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
}

漏斗算法

// 漏斗算法
package main

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

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

	}
}

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
	}


}

令牌桶算法

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
	}
}