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

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) // 模拟请求间隔
	}
}

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) // 模拟请求间隔
	}
}

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) // 模拟请求间隔
	}
}