限流

// 利用channel生产/消费令牌
type TokenBucketLimit struct {
	intervalPerSecond int64
	buckets           chan struct{}
	closeCh           chan struct{}
}

func NewTokenBucketLimit(interval int64, capacity int) *TokenBucketLimit {
	return &TokenBucketLimit{
		intervalPerSecond: interval,
		buckets:           make(chan struct{}, capacity),
	}
}

func (l *TokenBucketLimit) NewServerInterceptor() grpc.UnaryServerInterceptor {
	interval := time.Second / time.Duration(l.intervalPerSecond)
	ticker := time.NewTicker(interval)
	go func() {
        defer ticker.Stop()
		for {
			select {
			case <-l.closeCh:
				return
			case <-ticker.C:
				select {
				case l.buckets <- struct{}{}:

				default:

				}
			}
		}
	}()
	// 此处可调整限流的粒度
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
		select {
		case <-l.buckets:
			return handler(ctx, req)
		default: // 并发越高，越不能阻塞，如果不高，此处也可以阻塞等待，直到超时
			return nil, errors.New("too many requests")
		}
	}
}

// Close 仅可调用一次
func (l *TokenBucketLimit) Close() {
	close(l.closeCh)
}

// 用锁保证令牌数的并发安全，也可换原子变量等
type TokenBucket struct {
	mu         sync.Mutex
	capacity   int              // 桶的容量
	tokens     int              // 桶中当前的令牌数
	refillRate float64          // 是令牌的每秒的填充速率
	lastRefill time.Time        // 记录上次填充令牌的时间
}

func NewTokenBucket(capacity int, refillRate float64) *TokenBucket {
	// 初始化时桶可被填满，tokens 和 capacity 相等
	return &TokenBucket{
		capacity:   capacity,
		tokens:     capacity,
		refillRate: refillRate,
		lastRefill: time.Now(),
	}
}

// Allow 方法用于判断当前请求是否被允许。
func (t *TokenBucket) Allow() bool {
	t.mu.Lock()
	defer t.mu.Unlock()

	now := time.Now()

	// 计算自上次填充以来经过的秒数
	timeElapsed := float64(now.Unix() - t.lastRefill.Unix())

	// 计算应该添加的令牌数。
	tokensToAdd := t.refillRate * timeElapsed

	// 更新令牌数，但不超过桶的容量
	t.tokens += int(tokensToAdd)
	if t.tokens > t.capacity {
		t.tokens = t.capacity
	}

	// 如果桶中有令牌，则移除一个令牌并允许请求通过
	if t.tokens > 0 {
		t.tokens--
		t.lastRefill = now
		return true
	}

	// 拒绝请求
	return false
}

// 类似于令牌桶的实现，只不过桶的容量为0
type TokenBucketLimit struct {
	intervalPerSecond int64
	buckets           chan struct{}
	closeCh           chan struct{}
}

func NewTokenBucketLimit(interval int64) *TokenBucketLimit {
	return &TokenBucketLimit{
		intervalPerSecond: interval,
		buckets:           make(chan struct{}),
	}
}

func (l *TokenBucketLimit) process() {
    interval := time.Second / time.Duration(l.intervalPerSecond)
	ticker := time.NewTicker(interval)
    
    for {
        defer ticker.Stop()
        select {
        case <-l.closeCh:
            return
        case <-ticker.C:
            <-l.buckets
        }
    }
}

func (l *TokenBucketLimit) NewServerInterceptor() grpc.UnaryServerInterceptor {
	go l.process()
    
	// 此处可调整限流的粒度
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
		select {
		case l.buckets <-struct{}{}:
			return handler(ctx, req)
		default: // 并发越高，越不能阻塞，如果不高，此处也可以阻塞等待，直到超时
			return nil, errors.New("too many requests")
		}
	}
}

// Close 仅可调用一次
func (l *TokenBucketLimit) Close() {
	close(l.closeCh)
}

// SlidingWindowLimiter 结构体实现滑动窗口限流算法。
type SlidingWindowLimiter struct {
	mutex          sync.Mutex
	counters       []int
	limit          int
	windowStart    time.Time
	windowDuration time.Duration
	interval       time.Duration
}

// NewSlidingWindowLimiter 构造函数初始化 SlidingWindowLimiter 实例。
func NewSlidingWindowLimiter(limit int, windowDuration time.Duration, interval time.Duration) *SlidingWindowLimiter {
	buckets := int(windowDuration / interval)
	return &SlidingWindowLimiter{
		counters:       make([]int, buckets),
		limit:          limit,
		windowStart:    time.Now(),
		windowDuration: windowDuration,
		interval:       interval,
	}
}

// Allow 方法用于判断当前请求是否被允许，并实现滑动窗口的逻辑。
func (s *SlidingWindowLimiter) Allow() bool {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	// 检查是否需要滑动窗口
	if time.Since(s.windowStart) > s.windowDuration {
		s.slideWindow()
	}

	now := time.Now()
	index := int((now.UnixNano() - s.windowStart.UnixNano()) / s.interval.Nanoseconds()) % len(s.counters)

	if s.counters[index] < s.limit {
		s.counters[index]++
		return true
	}
	return false
}

// slideWindow 方法实现滑动窗口逻辑，移除最旧的时间段并重置计数器。
func (s *SlidingWindowLimiter) slideWindow() {
	// 滑动窗口，忽略最旧的时间段
	copy(s.counters, s.counters[1:])
	// 重置最后一个时间段的计数器
	s.counters[len(s.counters)-1] = 0
	// 更新窗口开始时间
	s.windowStart = time.Now()
}