Go语言如何实现限制用户请求_Golang

在 go 语言中，限制用户每分钟最多请求 1000 次的常见做法是使用限流算法（rate limiting）。有多种算法可以实现这一目标，其中最常见的包括令牌桶算法 (token bucket)、漏桶算法 (leaky bucket) 和计数器算法 (counter)。每种算法有其特点和适用场景，下面将逐个介绍，并附上相应的 go 语言实现。

1. 令牌桶算法 (token bucket)

令牌桶算法是常见的限流算法，适用于需要平滑流量控制的场景。令牌桶维护一个存储令牌的桶，每个请求需要消耗一个令牌。如果桶内有足够的令牌，请求可以继续；如果没有令牌，则请求被拒绝。令牌按固定速率生成，当桶满时，额外的令牌会丢弃。

令牌桶算法的实现

package main

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

type tokenbucket struct {
	rate      int           // 生成令牌的速率，单位是令牌/秒
	capacity  int           // 桶的容量
	tokens    int           // 当前令牌数量
	lasttoken time.time     // 上次生成令牌的时间
	mutex     sync.mutex    // 用于并发控制
}

func newtokenbucket(rate, capacity int) *tokenbucket {
	return &tokenbucket{
		rate:     rate,
		capacity: capacity,
		tokens:   capacity, // 初始时，桶里有满的令牌
	}
}

func (tb *tokenbucket) refill() {
	// 计算过去时间段内生成的令牌数
	now := time.now()
	elapsed := now.sub(tb.lasttoken)
	tb.lasttoken = now
	// 按速率生成令牌
	newtokens := int(elapsed.seconds()) * tb.rate
	if newtokens > 0 {
		// 桶中令牌数增加
		tb.tokens += newtokens
		if tb.tokens > tb.capacity {
			// 超过桶容量，令牌数只能是桶的最大容量
			tb.tokens = tb.capacity
		}
	}
}

func (tb *tokenbucket) allow() bool {
	tb.mutex.lock()
	defer tb.mutex.unlock()
	// 补充令牌
	tb.refill()
	if tb.tokens > 0 {
		// 有令牌可以消耗
		tb.tokens--
		return true
	}
	// 没有令牌可用，限制请求
	return false
}

func main() {
	// 创建令牌桶，令牌生成速率为每秒 1000 个，容量为 1000 个令牌
	tb := newtokenbucket(1000, 1000)

	// 模拟用户发起请求
	for i := 0; i < 10; i++ {
		if tb.allow() {
			fmt.println("request", i+1, "allowed")
		} else {
			fmt.println("request", i+1, "rejected")
		}
		time.sleep(100 * time.millisecond) // 模拟请求间隔
	}
}

说明：

rate：每秒生成的令牌数。

capacity：桶的最大容量。

tokens：当前桶中可用的令牌数。

每次请求时，allow() 方法会检查桶中是否有令牌，如果有，则消耗一个令牌并允许请求；如果没有令牌，则拒绝请求。

2. 漏桶算法 (leaky bucket)

漏桶算法是另一种常用的限流算法，适用于流量平滑控制。在漏桶算法中，桶里有水（请求），水按固定速率流出。当请求到来时，如果桶满了，新的请求会被丢弃；如果桶未满，新的请求会被加入桶中，并在固定速率下流出。

漏桶算法的实现

package main

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

type leakybucket struct {
	rate      int           // 水流出速率，单位是请求/秒
	capacity  int           // 桶的容量
	water     int           // 当前桶中水的数量
	lastdrain time.time     // 上次排水时间
	mutex     sync.mutex    // 用于并发控制
}

func newleakybucket(rate, capacity int) *leakybucket {
	return &leakybucket{
		rate:     rate,
		capacity: capacity,
		water:    0, // 初始时，桶里没有水
	}
}

func (lb *leakybucket) drain() {
	// 计算过去时间段内排出的请求数
	now := time.now()
	elapsed := now.sub(lb.lastdrain)
	lb.lastdrain = now
	// 按排出速率流出请求
	drained := int(elapsed.seconds()) * lb.rate
	if drained > 0 {
		lb.water -= drained
		if lb.water < 0 {
			lb.water = 0
		}
	}
}

func (lb *leakybucket) allow() bool {
	lb.mutex.lock()
	defer lb.mutex.unlock()
	// 排水
	lb.drain()
	if lb.water < lb.capacity {
		// 桶未满，允许请求
		lb.water++
		return true
	}
	// 桶已满，拒绝请求
	return false
}

func main() {
	// 创建漏桶，排水速率为每秒 1000 个，桶的容量为 1000 个
	lb := newleakybucket(1000, 1000)

	// 模拟用户发起请求
	for i := 0; i < 10; i++ {
		if lb.allow() {
			fmt.println("request", i+1, "allowed")
		} else {
			fmt.println("request", i+1, "rejected")
		}
		time.sleep(100 * time.millisecond) // 模拟请求间隔
	}
}

说明：

rate：请求的排出速率。

capacity：桶的最大容量。

water：当前桶中水（请求）的数量。

drain()：排水操作，控制请求的流出速率。

3. 计数器算法 (fixed window counter)

计数器算法是最简单的一种限流算法。在每个时间窗口内，记录请求的数量。当请求数达到限制时，就会拒绝进一步的请求。它适用于简单的限流场景，但对于高并发时可能会出现窗口突发的情况。

计数器算法的实现

package main

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

type counter struct {
	limit      int           // 请求限制次数
	windowsize time.duration // 时间窗口大小
	mu         sync.mutex    // 用于并发控制
	requests   int           // 当前请求计数
	windowstart time.time    // 当前时间窗口开始时间
}

func newcounter(limit int, windowsize time.duration) *counter {
	return &counter{
		limit:      limit,
		windowsize: windowsize,
		requests:   0,
		windowstart: time.now(),
	}
}

func (c *counter) allow() bool {
	c.mu.lock()
	defer c.mu.unlock()
	// 判断是否在当前时间窗口内
	now := time.now()
	if now.sub(c.windowstart) > c.windowsize {
		// 如果超过了窗口时间，则重置请求计数器和窗口开始时间
		c.windowstart = now
		c.requests = 0
	}
	if c.requests < c.limit {
		// 如果请求数未达到限制，允许请求
		c.requests++
		return true
	}
	// 否则，拒绝请求
	return false
}

func main() {
	// 创建计数器，限制每分钟 1000 次请求
	counter := newcounter(1000, time.minute)

	// 模拟用户发起请求
	for i := 0; i < 10; i++ {
		if counter.allow() {
			fmt.println("request", i+1, "allowed")
		} else {
			fmt.println("request", i+1, "rejected")
		}
		time.sleep(100 * time.millisecond) // 模拟请求间隔
	}
}

说明：

limit：时间窗口内允许的最大请求次数。

windowsize：时间窗口的大小（比如 1 分钟）。

requests：当前时间窗口内已处理的请求数量。

allow()：每次请求时，检查当前窗口内请求数是否达到限制。