Go简单实现几种常用的限流

固定窗口

package main
 
import (
    "fmt"
    "sync"
    "sync/atomic"
    "time"
)
 
// 定义限流结构体
type RateLimiter struct {
    interval time.Duration // 时间窗口
    tokens   int32         // 令牌总数
    lastTime int64         // 上次更新时间
    mu       sync.Mutex    // 互斥锁
}
 
// 创建新的限流器
func NewRateLimiter(interval time.Duration, tokens int32) *RateLimiter {
    return &RateLimiter{
        interval: interval,
        tokens:   tokens,
        lastTime: time.Now().UnixNano(),
    }
}
 
// 尝试获取令牌
func (l *RateLimiter) TryAcquire() bool {
    l.mu.Lock()
    defer l.mu.Unlock()
 
    now := time.Now().UnixNano()
    // 计算当前窗口内的剩余时间
    elapsed := now - l.lastTime
    // 更新时间窗口，并添加新令牌
    l.lastTime = now
    l.tokens += elapsed / int64(l.interval)
 
    // 判断是否有足够的令牌
    if l.tokens < 1 {
        return false
    }
    l.tokens-- // 获取令牌，减1
    return true
}
 
func main() {
    // 创建一个时间窗口为1秒，总共10个令牌的限流器
    limiter := NewRateLimiter(time.Second, 10)
 
    // 模拟50个并发请求
    var wg sync.WaitGroup
    wg.Add(50)
    for i := 0; i < 50; i++ {
        go func() {
            defer wg.Done()
            if limiter.TryAcquire() {
                fmt.Println("请求被处理")
            } else {
                fmt.Println("请求被限流")
            }
        }()
    }
    wg.Wait()
}

计数器

package main

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

type CountLimiter struct {
	rate  int64         // 计数周期内运行的最大请求数
	begin time.Time     // 当前轮计数开始时间
	count int64         // 当前计数周期累计的请求数
	cycle time.Duration // 计数周期，如统计1秒内的总请求数，那么计数周期就是1秒
	lock  sync.Mutex    // 判断能否放行时需要加锁操作
}

func NewCountLimiter(rate int64, cycle time.Duration) *CountLimiter {
	return &CountLimiter{
		rate:  rate,
		begin: time.Now(),
		count: 0,
		cycle: cycle,
		lock:  sync.Mutex{},
	}
}

func (c *CountLimiter) Allow() bool {
	c.lock.Lock()
	defer c.lock.Unlock()
	if c.count == c.rate { // 达到最大限流时，判断时间是否也超过统计周期了，超了可以重置限流器了，没有超说明还在当前统计周期内，应该拦截
		if time.Now().Sub(c.begin) > c.cycle {
			c.Reset()
			return true
		} else {
			return false
		}
	} else { // 还没有达到最大限流数，那么不管时间周期是否超出统计计数周期，都可以放行
		c.count++
		return true
	}
}

func (c *CountLimiter) Reset() {
	c.begin = time.Now()
	c.count = 0
}

func main() {
	countLimiter := NewCountLimiter(3, time.Second) // 1s内不可超过3个请求
	var wg sync.WaitGroup
	wg.Add(10)
	for i := 0; i < 10; i++ {
		go func(i int) {
			defer wg.Done()
			if countLimiter.Allow() {
				fmt.Println(fmt.Sprintf("当前时间%s,请求%d通过了", time.Now().String(), i))
			} else {
				fmt.Println(fmt.Sprintf("当前时间%s,请求%d被限流了", time.Now().String(), i))
			}
		}(i)
		time.Sleep(200 * time.Millisecond)
	}
	wg.Wait()
}

令牌桶

package main

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

type TokenBucketLimiter struct {
	lock     sync.Mutex
	rate     time.Duration // 多长时间放入一个令牌，即放入令牌的速率
	capacity int64         // 令牌桶的容量，控制最多放入多少令牌，也即突发最大并发量
	tokens   int64         // 当前桶中已有的令牌数量
	lastTime time.Time     // 上次放入令牌的时间，避免开启协程定时去放入令牌，而是请求到来时懒加载的方式(now - lastTime) / rate放入令牌
}

func NewTokenBucketLimiter(rate time.Duration, capacity int64) *TokenBucketLimiter {
	if capacity < 1 {
        panic("token bucket capacity must be large 1")
    }
	return &TokenBucketLimiter{
		lock:     sync.Mutex{},
		rate:     rate,
		capacity: capacity,
		tokens:   0,
		lastTime: time.Time{},
	}
}

func (tbl *TokenBucketLimiter) Allow() bool {
	tbl.lock.Lock() // 加锁避免并发错误
	defer tbl.lock.Unlock()

	// 如果 now 与上次请求的间隔超过了 token rate
	// 则增加令牌，更新lastTime
	now := time.Now()
	if now.Sub(tbl.lastTime) > tbl.rate {
		tbl.tokens += int64((now.Sub(tbl.lastTime)) / tbl.rate) // 放入令牌
		if tbl.tokens > tbl.capacity {
			tbl.tokens = tbl.capacity // 总令牌数不能大于桶的容量
		}
		tbl.lastTime = now // 更新上次往桶中放入令牌的时间
	}

	if tbl.tokens > 0 { // 令牌数是否充足
		tbl.tokens -= 1
		return true
	}

	return false // 令牌不足，拒绝请求
}

func main() {
	tbl := NewTokenBucketLimiter(10, 5) // 每10ms放一个令牌，1s放100个，桶容量（最大突发流量）为5
	for i := 0; i < 10; i++ {
		fmt.Println(tbl.Allow())  // 模拟突发流量10个请求，超过桶容量5
	}
	time.Sleep(100 * time.Millisecond)
	fmt.Println(tbl.Allow())
}

漏斗

package main

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

type Result struct {
	Msg string // 根据实际情况定义返回结果需要哪些字段
}

type Handler func() Result // 处理函数的形式也应该根据具体需要而定

type Task struct {
	id      int64       // 任务id
	result  chan Result // 任务的执行结果，即请求的响应结果
	handler Handler     // 请求的执行函数
}

func NewTask(id int, handler Handler) Task {
	return Task{
		handler: handler,
		result:  make(chan Result),
		id:      int64(id),
	}
}

type LeakyBucketLimiter struct {
	bucketSize int64     // 桶的大小
	workerNum  int64     // 工作者数量，即最大并发数
	taskChan   chan Task // 用于存放请求
}

func NewLeakyBucketLimiter(bucketSize, workerNum int64) *LeakyBucketLimiter {
	if capacity < 1 {
        panic("capacity must be large 1")
    }
    if workerNum < 1 {
        panic("workerNum must be large 1")
    }
	return &LeakyBucketLimiter{
		bucketSize: bucketSize,
		workerNum:  workerNum,
		taskChan:   make(chan Task, bucketSize),
	}
}

func (lbl *LeakyBucketLimiter) AddTask(task Task) bool { // 类似其他限流算法的Allow方法
	// 如果木桶已经满了，或者任务执行失败或超时了，返回false
	select {
	case lbl.taskChan <- task:  // 利用了select的特性判断是否能往通道中添加任务
	default:
		fmt.Printf("请求%d被拒绝了\n", task.id)
		return false
	}

	// 如果成功入桶，调用者会等待Task的Handler执行结果
	// 由于Task的result是无缓冲的通道，不应该让其无限等待阻塞，否则出现问题时，不往该chan写，就会一直阻塞在这里了，泄漏
	// 因此设置一个超时时间
	//resp := <-task.result
	//fmt.Printf("请求%d运行成功，结果为：%v\n", task.id, resp)
	select {
	case resp := <-task.result:
		fmt.Printf("请求%d运行成功，结果为：%v\n", task.id, resp)
	case <-time.After(5 * time.Second): // 超时时间可以稍微设置长一点点，因为任务放入桶中后，可能需要排队一点时间才被拉取出来执行
		return false // 这里超时当被限流处理
	}

	return true
}

func (lbl *LeakyBucketLimiter) Start(ctx context.Context) {
	// 开启workerNum个协程从木桶拉取任务执行
	for i := 0; int64(i) < lbl.workerNum; i++ {
		go func(ctx context.Context) {
			defer func() { // 铁则：开启的子协程一定要捕获异常，否则一旦出现异常不捕获，会导致程序退出
				if err := recover(); err != any(nil) {
					fmt.Println("捕获到异常")
				}
			}()

			for { // 持续监听，拉取任务执行
				select {
				case <-ctx.Done():
					fmt.Println("退出工作")
					return
				default:
					task := <-lbl.taskChan
					result := task.handler()
					task.result <- result // 处理结果写入对应Task的结果通道
				}
			}
		}(ctx)
	}
}

func main() {
	bucket := NewLeakyBucketLimiter(10, 4)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	bucket.Start(ctx) // 开启消费者
	// 模拟20个并发请求
	var wg sync.WaitGroup
	wg.Add(20)
	for i := 0; i < 20; i++ {
		go func(id int) {
			defer wg.Done()
			task := NewTask(id, func() Result { // 这里的func应该根据实际需要定义为handler，写具体的业务逻辑和返回的result
				time.Sleep(300 * time.Millisecond)  // 模拟业务逻辑消耗的时间
				return Result{}
			})
			bucket.AddTask(task)  // 请求入桶
		}(i)
	}
	wg.Wait()
	time.Sleep(10 * time.Second)
}

滑动窗口

package main

import (
    "errors"
    "sync"
    "time"
)

// SlidingWindowLimiter 滑动窗口限流器，用于控制请求的速率。
type SlidingWindowLimiter struct {
    limit        int           // 窗口内允许的最大请求数
    window       int64         // 窗口时间大小（纳秒）
    smallWindow  int64         // 小窗口时间大小（纳秒）
    smallWindows int64         // 窗口内小窗口的数量
    counters     map[int64]int // 每个小窗口的请求计数
    mutex        sync.RWMutex  // 使用读写锁提高并发性能
}

// NewSlidingWindowLimiter 创建并初始化滑动窗口限流器。
func NewSlidingWindowLimiter(limit int, window, smallWindow time.Duration) (*SlidingWindowLimiter, error) {
    if int64(window%smallWindow) != 0 {
       return nil, errors.New("window size must be divisible by the small window size")
    }
    return &SlidingWindowLimiter{
       limit:        limit,
       window:       int64(window),
       smallWindow:  int64(smallWindow),
       smallWindows: int64(window / smallWindow),
       counters:     make(map[int64]int),
    }, nil
}

// TryAcquire 尝试在当前窗口内获取一个请求的机会。
func (l *SlidingWindowLimiter) TryAcquire() bool {
    l.mutex.RLock() // 读锁，允许多个并发读取
    defer l.mutex.RUnlock()

    now := time.Now().UnixNano()
    currentSmallWindow := now / l.smallWindow * l.smallWindow // 当前小窗口的起始点

    // 清理过期的小窗口计数器
    l.cleanExpiredWindows(now)

    // 检查并更新当前小窗口的计数
    l.mutex.Lock() // 写锁，更新计数器
    defer l.mutex.Unlock()

    count, exists := l.counters[currentSmallWindow]
    if !exists || count < l.limit {
       l.counters[currentSmallWindow] = count + 1
       return true
    }
    return false
}

// cleanExpiredWindows 清理已过期的小窗口计数器。
func (l *SlidingWindowLimiter) cleanExpiredWindows(now int64) {
    startSmallWindow := now/l.smallWindow*l.smallWindow - l.window
    for smallWindow := range l.counters {
       if smallWindow < startSmallWindow {
          delete(l.counters, smallWindow)
       }
    }
}

func main() {
    // 创建限流器，窗口大小为10秒，小窗口大小为1秒，窗口内最大请求数为5
    limiter, err := limiter.NewSlidingWindowLimiter(5, 10*time.Second, time.Second)
    if err != nil {
        fmt.Println("Error creating limiter:", err)
        return
    }

    // 模拟请求
    for i := 0; i < 10; i++ {
        if limiter.TryAcquire() {
            fmt.Println("Request", i+1, "allowed")
        } else {
            fmt.Println("Request", i+1, "denied")
        }
        time.Sleep(100 * time.Millisecond) // 等待100毫秒
    }
}