基于 Uber AtomicLimiter 的高性能按 Key 限流设计与实践

基于 Uber AtomicLimiter 的高性能按 Key 限流设计与实践

背景介绍

在现代分布式系统中，API 调用频率控制（限流）是保障系统稳定性和防止资源滥用的关键手段。尤其是在多租户、多业务场景下，往往需要对不同的业务维度（如用户、接口、区域等）进行独立限流。

本文结合 Uber 开源的高性能原子限流器 AtomicLimiter ，设计并实现了一个支持按 Key 动态限流的限流管理器，满足多维度限流需求，且具备良好的并发性能和扩展性。

---

限流设计思路

1. 按 Key 限流

• 不同业务维度对应不同的限流速率。
• 例如：用户 A 限流 10 QPS，用户 B 限流 5 QPS。
• 需要动态维护每个 Key 对应的限流器。

2. 高性能并发安全

• 限流器内部使用无锁原子操作，保证高并发下性能。
• 管理器维护限流器映射，需保证并发安全。

3. 动态创建与复用

• 第一次访问某个 Key 时创建对应限流器。
• 后续复用，避免重复创建。

---

Uber AtomicLimiter 简介

AtomicLimiter 是 Uber 开源的基于原子操作的限流器，核心特点：

• 通过原子指针和 CAS 操作实现无锁并发安全。
• 维护上次允许请求时间和待等待时间，平滑控制请求速率。
• 支持阻塞调用 Take()，保证请求间隔平均为 1/rate。

核心代码片段：

func (t *AtomicLimiter) Take() time.Time {
	for !taken {
		now := t.clock.Now()
		oldState := (*state)(atomic.LoadPointer(&t.state))
		newState := state{last: now, sleepFor: oldState.sleepFor}

		if oldState.last.IsZero() {
			taken = atomic.CompareAndSwapPointer(&t.state, previousStatePointer, unsafe.Pointer(&newState))
			continue
		}

		newState.sleepFor += t.perRequest - now.Sub(oldState.last)
		if newState.sleepFor < t.maxSlack {
			newState.sleepFor = t.maxSlack
		}
		if newState.sleepFor > 0 {
			newState.last = newState.last.Add(newState.sleepFor)
			interval, newState.sleepFor = newState.sleepFor, 0
		}
		taken = atomic.CompareAndSwapPointer(&t.state, previousStatePointer, unsafe.Pointer(&newState))
	}
	t.clock.Sleep(interval)
	return newState.last
}

---

按 Key 限流管理器设计

基于 AtomicLimiter，我们设计了如下限流管理器：

package datasource

import (
	"k8s.io/klog/v2"
	"sync"
	"time"
)

const (
	// DefaultLimit 默认流控
	DefaultLimit int = 15
)

// YunAPIRateLimit 定义按 Key 限流管理器
type YunAPIRateLimit struct {
	RateLimiterMap map[string]*ratelimit.AtomicLimiter
	Locker         *sync.Mutex
}

// NewYunAPIRateLimit 初始化限流管理器
func NewYunAPIRateLimit() *YunAPIRateLimit {
	return &YunAPIRateLimit{
		RateLimiterMap: map[string]*ratelimit.AtomicLimiter{},
		Locker:         &sync.Mutex{},
	}
}

// Check 执行限流，阻塞直到允许通过
func (y *YunAPIRateLimit) Check(key string, rate int) {
	// 先加锁判断限流器是否存在，若不存在则创建
	y.Locker.Lock()
	limiter, exists := y.RateLimiterMap[key]
	if !exists {
		klog.Infof("create limiter for key: %s with rate: %d", key, rate)
		limiter = ratelimit.NewAtomicBased(rate)
		y.RateLimiterMap[key] = limiter
	}
	y.Locker.Unlock()

	// 调用限流器 Take()，阻塞等待
	start := time.Now()
	now := limiter.Take()
	klog.Infof("key: %s, wait time: %v", key, now.Sub(start))
}

设计要点

• 使用互斥锁保护限流器映射，保证并发安全。
• 动态创建限流器，避免预先初始化所有 Key。
• 调用 Take() 阻塞等待，保证限流效果。
• 记录等待时间，方便监控和调优。

---

实际应用场景示例

假设我们有一个云服务平台，提供多租户 API 接口。每个租户（Tenant）有不同的调用限额：

• Tenant A：20 QPS
• Tenant B：10 QPS
• Tenant C：5 QPS

我们可以在 API 网关或服务层集成 YunAPIRateLimit，按租户 ID 限流：

func HandleRequest(tenantID string, rateLimiter *YunAPIRateLimit) {
	// 这里 rate 可以从配置中心动态获取
	rate := getRateForTenant(tenantID)
	rateLimiter.Check(tenantID, rate)

	// 业务处理逻辑
	processRequest()
}

这样，系统自动对不同租户进行独立限流，防止某个租户流量暴涨影响整体服务。

---

性能与优化建议

• 锁粒度优化：当前锁保护整个映射，可能成为瓶颈。可考虑使用 sync.Map 或分段锁优化。
• 限流器生命周期管理：支持限流器过期回收，避免内存泄漏。
• 动态速率调整：支持动态调整限流速率，满足业务弹性需求。
• 日志采样：减少日志打印频率，避免日志过载。

---

Uber AtomicLimiter 代码核心实现

以下是 Uber AtomicLimiter 的核心代码，方便理解限流器内部实现：

package ratelimit

import (
	"sync/atomic"
	"time"
	"unsafe"
)

type state struct {
	last     time.Time
	sleepFor time.Duration
}

type AtomicLimiter struct {
	state unsafe.Pointer
	padding [56]byte // 避免伪共享

	perRequest time.Duration
	maxSlack   time.Duration
	clock      Clock
}

func NewAtomicBased(rate int, opts ...Option) *AtomicLimiter {
	config := buildConfig(opts)
	perRequest := config.per / time.Duration(rate)
	l := &AtomicLimiter{
		perRequest: perRequest,
		maxSlack:   -1 * time.Duration(config.slack) * perRequest,
		clock:      config.clock,
	}

	initialState := state{
		last:     time.Time{},
		sleepFor: 0,
	}
	atomic.StorePointer(&l.state, unsafe.Pointer(&initialState))
	return l
}

func (t *AtomicLimiter) Take() time.Time {
	var (
		newState state
		taken    bool
		interval time.Duration
	)
	for !taken {
		now := t.clock.Now()
		previousStatePointer := atomic.LoadPointer(&t.state)
		oldState := (*state)(previousStatePointer)

		newState = state{
			last:     now,
			sleepFor: oldState.sleepFor,
		}

		if oldState.last.IsZero() {
			taken = atomic.CompareAndSwapPointer(&t.state, previousStatePointer, unsafe.Pointer(&newState))
			continue
		}

		newState.sleepFor += t.perRequest - now.Sub(oldState.last)
		if newState.sleepFor < t.maxSlack {
			newState.sleepFor = t.maxSlack
		}
		if newState.sleepFor > 0 {
			newState.last = newState.last.Add(newState.sleepFor)
			interval, newState.sleepFor = newState.sleepFor, 0
		}
		taken = atomic.CompareAndSwapPointer(&t.state, previousStatePointer, unsafe.Pointer(&newState))
	}
	t.clock.Sleep(interval)
	return newState.last
}

func (t *AtomicLimiter) GetLastTask() time.Time {
	previousStatePointer := atomic.LoadPointer(&t.state)
	oldState := (*state)(previousStatePointer)
	return oldState.last
}

---

总结

本文介绍了基于 Uber AtomicLimiter 的高性能按 Key 限流设计，结合实际业务场景，展示了如何实现灵活且高效的多维度限流。该方案适用于多租户、分布式服务等复杂环境，帮助保障系统稳定性和公平性。

欢迎大家试用并提出宝贵意见！

---

参考链接

• Uber go.uber.org/ratelimit
• Go sync/atomic 包文档
• Go 语言并发编程实战