feat(完整功能): 以Apache License 2.0协议开源

NOTICE

@@ -0,0 +1,13 @@
+   Copyright 2025 肖其顿
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,2 +1,44 @@
-# limit
+# limit [![PkgGoDev](https://pkg.go.dev/badge/github.com/xiaoqidun/limit)](https://pkg.go.dev/github.com/xiaoqidun/limit)
-高性能、并发安全的动态速率限制器
+一个高性能、并发安全的 Go 语言动态速率限制器。
+
+# 安装指南
+```shell
+go get -u github.com/xiaoqidun/limit
+```
+
+# 快速开始
+```go
+package main
+
+import (
+	"fmt"
+
+	"github.com/xiaoqidun/limit"
+	"golang.org/x/time/rate"
+)
+
+func main() {
+	// 1. 创建一个新的 Limiter 实例
+	limiter := limit.New()
+	// 2. 确保在程序退出前优雅地停止后台任务，这非常重要
+	defer limiter.Stop()
+
+	// 3. 为任意键 "some-key" 获取一个速率限制器
+	//    - rate.Limit(2): 表示速率为 "每秒2个请求"
+	//    - 2: 表示桶的容量 (Burst)，允许瞬时处理2个请求
+	rateLimiter := limiter.Get("some-key", rate.Limit(2), 2)
+
+	// 4. 模拟3次连续的突发请求
+	//    由于速率和容量都为2，只有前两次请求能立即成功
+	for i := 0; i < 3; i++ {
+		if rateLimiter.Allow() {
+			fmt.Printf("请求 %d: 已允许\n", i+1)
+		} else {
+			fmt.Printf("请求 %d: 已拒绝\n", i+1)
+		}
+	}
+}
+```
+
+# 授权协议
+本项目基于 [Apache License 2.0](https://github.com/xiaoqidun/limit/blob/main/LICENSE) 授权。

example_test.go

@@ -0,0 +1,68 @@
+// Copyright 2025 肖其顿
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package limit_test
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/xiaoqidun/limit"
+	"golang.org/x/time/rate"
+)
+
+// ExampleLimiter 演示了 limit 包的基本用法。
+func ExampleLimiter() {
+	// 创建一个使用默认配置的 Limiter 实例
+	limiter := limit.New()
+	// 程序退出前，优雅地停止后台任务，这非常重要
+	defer limiter.Stop()
+	// 为一个特定的测试键获取一个速率限制器
+	// 限制为每秒2个请求，最多允许3个并发（桶容量）
+	testKey := "testKey"
+	rateLimiter := limiter.Get(testKey, rate.Limit(2), 3)
+	// 模拟连续的请求
+	for i := 0; i < 5; i++ {
+		if rateLimiter.Allow() {
+			fmt.Printf("请求 %d: 已允许\n", i+1)
+		} else {
+			fmt.Printf("请求 %d: 已拒绝\n", i+1)
+		}
+		time.Sleep(100 * time.Millisecond)
+	}
+	// 手动移除一个不再需要的限制器
+	limiter.Del(testKey)
+	// Output:
+	// 请求 1: 已允许
+	// 请求 2: 已允许
+	// 请求 3: 已允许
+	// 请求 4: 已拒绝
+	// 请求 5: 已拒绝
+}
+
+// ExampleNewWithConfig 展示了如何使用自定义配置。
+func ExampleNewWithConfig() {
+	// 自定义配置
+	config := limit.Config{
+		ShardCount: 64,               // 分片数量，必须是2的幂
+		GCInterval: 5 * time.Minute,  // GC 检查周期
+		Expiration: 15 * time.Minute, // 限制器过期时间
+	}
+	// 使用自定义配置创建一个 Limiter 实例
+	customLimiter := limit.NewWithConfig(config)
+	defer customLimiter.Stop()
+	fmt.Println("使用自定义配置的限制器已成功创建")
+	// Output:
+	// 使用自定义配置的限制器已成功创建
+}

go.mod

@@ -1,3 +1,5 @@
 module github.com/xiaoqidun/limit
 
 go 1.25.1
+
+require golang.org/x/time v0.13.0

go.sum

@@ -0,0 +1,2 @@
+golang.org/x/time v0.13.0 h1:eUlYslOIt32DgYD6utsuUeHs4d7AsEYLuIAdg7FlYgI=
+golang.org/x/time v0.13.0/go.mod h1:eL/Oa2bBBK0TkX57Fyni+NgnyQQN4LitPmob2Hjnqw4=

limit.go

@@ -0,0 +1,286 @@
+// Copyright 2025 肖其顿
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package limit 提供了一个高性能、并发安全的动态速率限制器。
+// 它使用分片锁来减少高并发下的锁竞争，并能自动清理长期未使用的限制器。
+package limit
+
+import (
+	"hash"
+	"hash/fnv"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"golang.org/x/time/rate"
+)
+
+// defaultShardCount 是默认的分片数量，设为2的幂可以优化哈希计算。
+const defaultShardCount = 32
+
+// Config 定义了 Limiter 的可配置项。
+type Config struct {
+	// ShardCount 指定分片数量，必须是2的幂。如果为0或无效值，则使用默认值32。
+	ShardCount int
+	// GCInterval 指定GC周期，即检查并清理过期限制器的间隔。如果为0，则使用默认值10分钟。
+	GCInterval time.Duration
+	// Expiration 指定过期时间，即限制器在最后一次使用后能存活多久。如果为0，则使用默认值30分钟。
+	Expiration time.Duration
+}
+
+// Limiter 是一个高性能、分片实现的动态速率限制器。
+// 它的实例在并发使用时是安全的。
+type Limiter struct {
+	// 存储所有分片
+	shards []*shard
+	// 配置信息
+	config Config
+	// 标记限制器是否已停止
+	stopped atomic.Bool
+	// 确保Stop方法只执行一次
+	stopOnce sync.Once
+}
+
+// New 使用默认配置创建一个新的 Limiter 实例。
+func New() *Limiter {
+	return NewWithConfig(Config{})
+}
+
+// NewWithConfig 根据提供的配置创建一个新的 Limiter 实例。
+func NewWithConfig(config Config) *Limiter {
+	// 如果未设置，则使用默认值
+	if config.ShardCount == 0 {
+		config.ShardCount = defaultShardCount
+	}
+	if config.GCInterval == 0 {
+		config.GCInterval = 10 * time.Minute
+	}
+	if config.Expiration == 0 {
+		config.Expiration = 30 * time.Minute
+	}
+
+	// 确保分片数量是2的幂，以便进行高效的位运算
+	if config.ShardCount <= 0 || (config.ShardCount&(config.ShardCount-1)) != 0 {
+		config.ShardCount = defaultShardCount
+	}
+
+	l := &Limiter{
+		shards: make([]*shard, config.ShardCount),
+		config: config,
+	}
+
+	// 初始化所有分片
+	for i := 0; i < config.ShardCount; i++ {
+		l.shards[i] = newShard(config.GCInterval, config.Expiration)
+	}
+	return l
+}
+
+// Get 获取或创建一个与指定键关联的速率限制器。
+// 如果限制器已存在，它会根据传入的 r (速率) 和 b (并发数) 更新其配置。
+// 如果 Limiter 实例已被 Stop 方法关闭，此方法将返回 nil。
+func (l *Limiter) Get(k string, r rate.Limit, b int) *rate.Limiter {
+	// 快速路径检查，避免在已停止时进行哈希和查找
+	if l.stopped.Load() {
+		return nil
+	}
+	// 定位到具体分片进行操作
+	return l.getShard(k).get(k, r, b)
+}
+
+// Del 手动移除一个与指定键关联的速率限制器。
+// 如果 Limiter 实例已被 Stop 方法关闭，此方法不执行任何操作。
+func (l *Limiter) Del(k string) {
+	// 快速路径检查
+	if l.stopped.Load() {
+		return
+	}
+	// 定位到具体分片进行操作
+	l.getShard(k).del(k)
+}
+
+// Stop 停止 Limiter 的所有后台清理任务，并释放相关资源。
+// 此方法对于并发调用是安全的，并且可以被多次调用。
+func (l *Limiter) Stop() {
+	l.stopOnce.Do(func() {
+		l.stopped.Store(true)
+		for _, s := range l.shards {
+			s.stop()
+		}
+	})
+}
+
+// getShard 根据key的哈希值获取对应的分片。
+func (l *Limiter) getShard(key string) *shard {
+	hasher := fnvHasherPool.Get().(hash.Hash32)
+	defer func() {
+		hasher.Reset()
+		fnvHasherPool.Put(hasher)
+	}()
+	_, _ = hasher.Write([]byte(key)) // FNV-1a never returns an error.
+	// 使用位运算代替取模，提高效率
+	return l.shards[hasher.Sum32()&(uint32(l.config.ShardCount)-1)]
+}
+
+// shard 代表 Limiter 的一个分片，它包含独立的锁和数据，以减少全局锁竞争。
+type shard struct {
+	mutex     sync.Mutex
+	stopCh    chan struct{}
+	limiter   map[string]*session
+	stopOnce  sync.Once
+	waitGroup sync.WaitGroup
+}
+
+// newShard 创建一个新的分片实例，并启动其gc任务。
+func newShard(gcInterval, expiration time.Duration) *shard {
+	s := &shard{
+		// mutex 会被自动初始化为其零值（未锁定状态）
+		stopCh:  make(chan struct{}),
+		limiter: make(map[string]*session),
+	}
+	s.waitGroup.Add(1)
+	go s.gc(gcInterval, expiration)
+	return s
+}
+
+// gc 定期清理分片中过期的限制器。
+func (s *shard) gc(interval, expiration time.Duration) {
+	defer s.waitGroup.Done()
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+	for {
+		// 优先检查停止信号，确保能快速响应
+		select {
+		case <-s.stopCh:
+			return
+		default:
+		}
+
+		select {
+		case <-ticker.C:
+			s.mutex.Lock()
+			// 再次检查分片是否已停止，防止在等待锁期间被停止
+			if s.limiter == nil {
+				s.mutex.Unlock()
+				return
+			}
+			for k, v := range s.limiter {
+				// 清理过期的限制器
+				if time.Since(v.lastGet) > expiration {
+					// 将 session 对象放回池中前，重置其状态
+					v.limiter = nil
+					v.lastGet = time.Time{}
+					sessionPool.Put(v)
+					delete(s.limiter, k)
+				}
+			}
+			s.mutex.Unlock()
+		case <-s.stopCh:
+			// 收到停止信号，退出goroutine
+			return
+		}
+	}
+}
+
+// get 获取或创建一个新的速率限制器，如果已存在则更新其配置。
+func (s *shard) get(k string, r rate.Limit, b int) *rate.Limiter {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+	// 检查分片是否已停止
+	if s.limiter == nil {
+		return nil
+	}
+	sess, ok := s.limiter[k]
+	if !ok {
+		// 从池中获取 session 对象
+		sess = sessionPool.Get().(*session)
+		sess.limiter = rate.NewLimiter(r, b)
+		s.limiter[k] = sess
+	} else {
+		// 如果已存在，则更新其速率和并发数
+		sess.limiter.SetLimit(r)
+		sess.limiter.SetBurst(b)
+	}
+	sess.lastGet = time.Now()
+	return sess.limiter
+}
+
+// del 从分片中移除一个键的速率限制器。
+func (s *shard) del(k string) {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+	// 检查分片是否已停止
+	if s.limiter == nil {
+		return
+	}
+	if sess, ok := s.limiter[k]; ok {
+		// 将 session 对象放回池中前，重置其状态
+		sess.limiter = nil
+		sess.lastGet = time.Time{}
+		sessionPool.Put(sess)
+		delete(s.limiter, k)
+	}
+}
+
+// stop 停止分片的gc任务，并同步等待其完成后再清理资源。
+func (s *shard) stop() {
+	// 使用 sync.Once 确保 channel 只被关闭一次，彻底避免并发风险
+	s.stopOnce.Do(func() {
+		close(s.stopCh)
+	})
+
+	// 等待 gc goroutine 完全退出
+	s.waitGroup.Wait()
+
+	// 锁定并进行最终的资源清理
+	// 因为 gc 已经退出，所以此时只有 Get/Del 会竞争锁
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	// 检查是否已被清理，防止重复操作
+	if s.limiter == nil {
+		return
+	}
+
+	// 将所有 session 对象放回对象池
+	for _, sess := range s.limiter {
+		sess.limiter = nil
+		sess.lastGet = time.Time{}
+		sessionPool.Put(sess)
+	}
+	// 清理map，释放内存，并作为停止标记
+	s.limiter = nil
+}
+
+// session 存储每个键的速率限制器实例和最后访问时间。
+type session struct {
+	// 最后一次访问时间
+	lastGet time.Time
+	// 速率限制器
+	limiter *rate.Limiter
+}
+
+// sessionPool 使用 sync.Pool 来复用 session 对象，以减少 GC 压力。
+var sessionPool = sync.Pool{
+	New: func() interface{} {
+		return new(session)
+	},
+}
+
+// fnvHasherPool 使用 sync.Pool 来复用 FNV-1a 哈希对象，以减少高并发下的内存分配。
+var fnvHasherPool = sync.Pool{
+	New: func() interface{} {
+		return fnv.New32a()
+	},
+}

limit_test.go

@@ -0,0 +1,95 @@
+// Copyright 2025 肖其顿
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package limit
+
+import (
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	"golang.org/x/time/rate"
+)
+
+// TestLimiter 覆盖了 Limiter 的主要功能。
+func TestLimiter(t *testing.T) {
+	// 子测试：验证基本的允许/拒绝逻辑
+	t.Run("基本功能测试", func(t *testing.T) {
+		limiter := New()
+		defer limiter.Stop()
+		key := "测试键"
+		// 创建一个每秒2个令牌，桶容量为1的限制器
+		rl := limiter.Get(key, rate.Limit(2), 1)
+		if rl == nil {
+			t.Fatal("limiter.Get() 意外返回 nil，测试无法继续")
+		}
+		if !rl.Allow() {
+			t.Error("rl.Allow(): 首次调用应返回 true, 实际为 false")
+		}
+		if rl.Allow() {
+			t.Error("rl.Allow(): 超出突发容量的调用应返回 false, 实际为 true")
+		}
+		time.Sleep(500 * time.Millisecond)
+		if !rl.Allow() {
+			t.Error("rl.Allow(): 令牌补充后的调用应返回 true, 实际为 false")
+		}
+	})
+
+	// 子测试：验证 Del 方法的功能
+	t.Run("删除功能测试", func(t *testing.T) {
+		limiter := New()
+		defer limiter.Stop()
+		key := "测试键"
+		rl1 := limiter.Get(key, rate.Limit(2), 1)
+		if !rl1.Allow() {
+			t.Fatal("获取限制器后的首次 Allow() 调用失败")
+		}
+		limiter.Del(key)
+		rl2 := limiter.Get(key, rate.Limit(2), 1)
+		if !rl2.Allow() {
+			t.Error("Del() 后重新获取的限制器未能允许请求")
+		}
+	})
+
+	// 子测试：验证 Stop 方法的功能
+	t.Run("停止功能测试", func(t *testing.T) {
+		limiter := New()
+		limiter.Stop()
+		if rl := limiter.Get("任意键", 1, 1); rl != nil {
+			t.Error("Stop() 后 Get() 应返回 nil, 实际返回了有效实例")
+		}
+		// 多次调用 Stop 不应引发 panic
+		limiter.Stop()
+	})
+
+	// 子测试：验证并发安全性
+	t.Run("并发安全测试", func(t *testing.T) {
+		limiter := New()
+		defer limiter.Stop()
+		var wg sync.WaitGroup
+		numGoroutines := 100
+		for i := 0; i < numGoroutines; i++ {
+			wg.Add(1)
+			go func(i int) {
+				defer wg.Done()
+				key := fmt.Sprintf("并发测试键-%d", i)
+				if limiter.Get(key, rate.Limit(10), 5) == nil {
+					t.Errorf("并发获取键 '%s' 时, Get() 意外返回 nil", key)
+				}
+			}(i)
+		}
+		wg.Wait()
+	})
+}