Go中实现用户每日限额的妙招

后端

2023-10-02 05:22:26

前言

在互联网应用中，经常会遇到需要对用户请求进行限流的情况，比如：

每个手机号每天只能发5条验证码短信
每个用户每小时只能连续尝试3次密码
每个会员每天只能领3次福利

这些场景都需要对用户请求进行限流，以防止恶意用户对系统造成过大的压力。

固定时间窗口限流算法

固定时间窗口限流算法是一种简单有效的限流算法，它的基本原理是：在固定的时间窗口内，只允许一定数量的请求通过，超过这个数量的请求将被拒绝。

固定时间窗口限流算法的实现非常简单，只需要一个计数器和一个时间窗口。当一个请求到来时，首先检查计数器是否已经达到上限，如果达到上限，则拒绝这个请求；否则，将计数器加1，并记录下当前时间。当时间窗口过去后，将计数器重置为0。

固定时间窗口限流算法的优点是简单易懂，实现方便。缺点是它的灵活性较差，如果需要对不同用户或不同类型的请求进行不同的限流策略，则需要使用更复杂的限流算法。

滑动窗口限流算法

滑动窗口限流算法是固定时间窗口限流算法的改进，它的基本原理是：将时间窗口划分为多个小的时间段，每个小的时间段都有自己的计数器。当一个请求到来时，首先检查当前小的时间段的计数器是否已经达到上限，如果达到上限，则拒绝这个请求；否则，将计数器加1。当当前小的时间段过去后，将计数器重置为0，并移动到下一个小的时间段。

滑动窗口限流算法的优点是灵活性强，可以对不同用户或不同类型的请求进行不同的限流策略。缺点是它的实现比固定时间窗口限流算法复杂一些。

Go 中的实现

在 Go 中，可以使用 sync.Map 和 time 库来实现固定时间窗口限流算法和滑动窗口限流算法。

以下是一个使用 sync.Map 和 time 库实现的固定时间窗口限流算法的示例代码：

package main

import (
	"sync"
	"time"
)

// RateLimiter is a fixed time window rate limiter.
type RateLimiter struct {
	sync.Mutex
	windowSize time.Duration
	maxRequests int
	requests map[string]int
}

// NewRateLimiter creates a new rate limiter.
func NewRateLimiter(windowSize time.Duration, maxRequests int) *RateLimiter {
	return &RateLimiter{
		windowSize: windowSize,
		maxRequests: maxRequests,
		requests: make(map[string]int),
	}
}

// Allow checks if a request is allowed.
func (r *RateLimiter) Allow(key string) bool {
	r.Lock()
	defer r.Unlock()

	now := time.Now()

	// Remove expired requests.
	for k, v := range r.requests {
		if now.Sub(v) > r.windowSize {
			delete(r.requests, k)
		}
	}

	// Check if the request is allowed.
	if r.requests[key] < r.maxRequests {
		r.requests[key]++
		return true
	}

	return false
}

func main() {
	// Create a new rate limiter.
	rateLimiter := NewRateLimiter(time.Hour, 100)

	// Check if a request is allowed.
	if rateLimiter.Allow("user1") {
		// The request is allowed.
	} else {
		// The request is not allowed.
	}
}

以下是一个使用 sync.Map 和 time 库实现的滑动窗口限流算法的示例代码：

package main

import (
	"sync"
	"time"
)

// RateLimiter is a sliding window rate limiter.
type RateLimiter struct {
	sync.Mutex
	windowSize time.Duration
	numWindows int
	maxRequests int
	windows map[string][]int
}

// NewRateLimiter creates a new rate limiter.
func NewRateLimiter(windowSize time.Duration, numWindows int, maxRequests int) *RateLimiter {
	return &RateLimiter{
		windowSize: windowSize,
		numWindows: numWindows,
		maxRequests: maxRequests,
		windows: make(map[string][]int),
	}
}

// Allow checks if a request is allowed.
func (r *RateLimiter) Allow(key string) bool {
	r.Lock()
	defer r.Unlock()

	now := time.Now()

	// Remove expired windows.
	for k, v := range r.windows {
		if now.Sub(v[0]) > r.windowSize {
			delete(r.windows, k)
		}
	}

	// Get the current window.
	window := r.windows[key]

	// If the current window is full, move to the next window.
	if len(window) >= r.numWindows {
		window = window[1:]
	}

	// Add the current time to the window.
	window = append(window, now)

	// Check if the request is allowed.
	if len(window) <= r.maxRequests {
		r.windows[key] = window
		return true
	}

	return false
}

func main() {
	// Create a new rate limiter.
	rateLimiter := NewRateLimiter(time.Second, 10, 100)

	// Check if a request is allowed.
	if rateLimiter.Allow("user1") {
		// The request is allowed.
	} else {
		// The request is not allowed.
	}
}