You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
72 lines
1.2 KiB
Go
72 lines
1.2 KiB
Go
package builtin
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
|
|
ctxpkg "github.com/noahlann/nnet/pkg/context"
|
|
routerpkg "github.com/noahlann/nnet/pkg/router"
|
|
)
|
|
|
|
// RateLimiter 限流器
|
|
type RateLimiter struct {
|
|
tokens int
|
|
capacity int
|
|
rate time.Duration
|
|
lastCheck time.Time
|
|
mu sync.Mutex
|
|
}
|
|
|
|
// NewRateLimiter 创建限流器
|
|
func NewRateLimiter(capacity int, rate time.Duration) *RateLimiter {
|
|
return &RateLimiter{
|
|
tokens: capacity,
|
|
capacity: capacity,
|
|
rate: rate,
|
|
lastCheck: time.Now(),
|
|
}
|
|
}
|
|
|
|
// Allow 检查是否允许
|
|
func (rl *RateLimiter) Allow() bool {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
now := time.Now()
|
|
elapsed := now.Sub(rl.lastCheck)
|
|
|
|
// 根据时间间隔补充令牌
|
|
if elapsed > rl.rate {
|
|
tokensToAdd := int(elapsed / rl.rate)
|
|
rl.tokens += tokensToAdd
|
|
if rl.tokens > rl.capacity {
|
|
rl.tokens = rl.capacity
|
|
}
|
|
rl.lastCheck = now
|
|
}
|
|
|
|
if rl.tokens > 0 {
|
|
rl.tokens--
|
|
return true
|
|
}
|
|
|
|
return false
|
|
}
|
|
|
|
// RateLimitMiddleware 限流中间件
|
|
func RateLimitMiddleware(limiter *RateLimiter) routerpkg.Handler {
|
|
return func(ctx ctxpkg.Context) error {
|
|
if limiter == nil {
|
|
return nil
|
|
}
|
|
|
|
if !limiter.Allow() {
|
|
ctx.Response().WriteString("Rate limit exceeded\n")
|
|
return nil
|
|
}
|
|
|
|
return nil
|
|
}
|
|
}
|
|
|