Go語言如何實現限制用戶請求

更新時間：2025年01月21日 08:36:20 作者：Ai 編碼

這篇文章主要為大家詳細介紹了Go語言如何實現限制用戶 1 分鐘內最多請求 1000 次,文中為大家整理了三個常用的方法,希望對大家有所幫助

在 Go 語言中，限制用戶每分鐘最多請求 1000 次的常見做法是使用限流算法（Rate Limiting）。有多種算法可以實現這一目標，其中最常見的包括令牌桶算法 (Token Bucket)、漏桶算法 (Leaky Bucket) 和計數器算法 (Counter)。每種算法有其特點和適用場景，下面將逐個介紹，并附上相應的 Go 語言實現。

1. 令牌桶算法 (Token Bucket)

令牌桶算法是常見的限流算法，適用于需要平滑流量控制的場景。令牌桶維護一個存儲令牌的桶，每個請求需要消耗一個令牌。如果桶內有足夠的令牌，請求可以繼續(xù)；如果沒有令牌，則請求被拒絕。令牌按固定速率生成，當桶滿時，額外的令牌會丟棄。

令牌桶算法的實現

package main

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

type TokenBucket struct {
	rate      int           // 生成令牌的速率，單位是令牌/秒
	capacity  int           // 桶的容量
	tokens    int           // 當前令牌數量
	lastToken time.Time     // 上次生成令牌的時間
	mutex     sync.Mutex    // 用于并發(fā)控制
}

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() {
	// 創(chuàng)建令牌桶，令牌生成速率為每秒 1000 個，容量為 1000 個令牌
	tb := NewTokenBucket(1000, 1000)

	// 模擬用戶發(fā)起請求
	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    // 用于并發(fā)控制
}

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() {
	// 創(chuàng)建漏桶，排水速率為每秒 1000 個，桶的容量為 1000 個
	lb := NewLeakyBucket(1000, 1000)

	// 模擬用戶發(fā)起請求
	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)

計數器算法是最簡單的一種限流算法。在每個時間窗口內，記錄請求的數量。當請求數達到限制時，就會拒絕進一步的請求。它適用于簡單的限流場景，但對于高并發(fā)時可能會出現窗口突發(fā)的情況。

計數器算法的實現

package main

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

type Counter struct {
	limit      int           // 請求限制次數
	windowSize time.Duration // 時間窗口大小
	mu         sync.Mutex    // 用于并發(fā)控制
	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() {
	// 創(chuàng)建計數器，限制每分鐘 1000 次請求
	counter := NewCounter(1000, time.Minute)

	// 模擬用戶發(fā)起請求
	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()：每次請求時，檢查當前窗口內請求數是否達到限制。