diff --git a/common/utils/limit/LICENSE b/common/utils/limit/LICENSE new file mode 100644 index 000000000..800f2c7c2 --- /dev/null +++ b/common/utils/limit/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + 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. diff --git a/common/utils/limit/NOTICE b/common/utils/limit/NOTICE new file mode 100644 index 000000000..73d68a92f --- /dev/null +++ b/common/utils/limit/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. diff --git a/common/utils/limit/README.md b/common/utils/limit/README.md new file mode 100644 index 000000000..cb911dbaf --- /dev/null +++ b/common/utils/limit/README.md @@ -0,0 +1,42 @@ +# 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) 授权协议 \ No newline at end of file diff --git a/common/utils/limit/example_test.go b/common/utils/limit/example_test.go new file mode 100644 index 000000000..5f96cbba9 --- /dev/null +++ b/common/utils/limit/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: + // 使用自定义配置的限制器已成功创建 +} diff --git a/common/utils/limit/go.mod b/common/utils/limit/go.mod new file mode 100644 index 000000000..fed5e3c10 --- /dev/null +++ b/common/utils/limit/go.mod @@ -0,0 +1,5 @@ +module github.com/xiaoqidun/limit + +go 1.18 + +require golang.org/x/time v0.8.0 diff --git a/common/utils/limit/go.sum b/common/utils/limit/go.sum new file mode 100644 index 000000000..d06eb417b --- /dev/null +++ b/common/utils/limit/go.sum @@ -0,0 +1,2 @@ +golang.org/x/time v0.8.0 h1:9i3RxcPv3PZnitoVGMPDKZSq1xW1gK1Xy3ArNOGZfEg= +golang.org/x/time v0.8.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= diff --git a/common/utils/limit/limit.go b/common/utils/limit/limit.go new file mode 100644 index 000000000..fd89e16f3 --- /dev/null +++ b/common/utils/limit/limit.go @@ -0,0 +1,278 @@ +// 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() + }, +} diff --git a/common/utils/limit/limit_test.go b/common/utils/limit/limit_test.go new file mode 100644 index 000000000..9b5e43e53 --- /dev/null +++ b/common/utils/limit/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() + }) +} diff --git a/go.work b/go.work index df7fa9f54..51d741e7f 100644 --- a/go.work +++ b/go.work @@ -13,6 +13,7 @@ use ( ./common/utils/go-jsonrpc ./common/utils/go-sensitive-word-1.3.3 ./common/utils/goja + ./common/utils/limit ./common/utils/log ./common/utils/sturc ./common/utils/xml diff --git a/logic/main.go b/logic/main.go index 84818e11a..44ad21557 100644 --- a/logic/main.go +++ b/logic/main.go @@ -10,6 +10,7 @@ import ( _ "blazing/contrib/drivers/pgsql" "blazing/logic/service/fight" + "blazing/logic/service/player" "blazing/cool" @@ -30,6 +31,12 @@ func PprofWeb() { } func signalHandlerForMain(sig os.Signal) { fight.Fightpool.Release() + + player.Mainplayer.Range(func(key uint32, value *player.Player) bool { + value.Save() + + return true + }) fmt.Println("MainProcess is shutting down due to signal:", sig.String()) } diff --git a/login/go.sum b/login/go.sum index c9c7a4aa3..c8f0075eb 100644 --- a/login/go.sum +++ b/login/go.sum @@ -8,6 +8,7 @@ github.com/clbanning/mxj/v2 v2.7.0 h1:WA/La7UGCanFe5NpHF0Q3DNtnCsVoxbPKuyBNHWRyM github.com/clbanning/mxj/v2 v2.7.0/go.mod h1:hNiWqW14h+kc+MdF9C6/YoRfjEJoR3ou6tn/Qo+ve2s= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM= @@ -49,6 +50,7 @@ github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY= github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/redis/go-redis/v9 v9.5.1 h1:H1X4D3yHPaYrkL5X06Wh6xNVM/pX0Ft4RV0vMGvLBh8= github.com/redis/go-redis/v9 v9.5.1/go.mod h1:hdY0cQFCN4fnSYT6TkisLufl/4W5UIXyv0b/CLO2V2M= github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= @@ -59,6 +61,8 @@ github.com/rogpeppe/go-internal v1.11.0 h1:cWPaGQEPrBb5/AsnsZesgZZ9yb1OQ+GOISoDN github.com/rogpeppe/go-internal v1.11.0/go.mod h1:ddIwULY96R17DhadqLgMfk9H9tvdUzkipdSkR5nkCZA= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/xiaoqidun/limit v1.0.1 h1:tJBDRv+JkziRghRCKHUVyxLGSL1CkkY+aPKMorWZJ20= +github.com/xiaoqidun/limit v1.0.1/go.mod h1:9blqqk5PcsnkWyIBwkkPHUQJjlx2p+Cf0TxVGz0MDFg= go.opentelemetry.io/otel v1.24.0 h1:0LAOdjNmQeSTzGBzduGe/rU4tZhMwL5rWgtp9Ku5Jfo= go.opentelemetry.io/otel v1.24.0/go.mod h1:W7b9Ozg4nkF5tWI5zsXkaKKDjdVjpD4oAt9Qi/MArHo= go.opentelemetry.io/otel/metric v1.24.0 h1:6EhoGWWK28x1fbpA4tYTOWBkPefTDQnb8WSGXlc88kI= @@ -77,6 +81,7 @@ golang.org/x/text v0.22.0 h1:bofq7m3/HAFvbF51jz3Q9wLg3jkvSPuiZu/pD1XwgtM= golang.org/x/text v0.22.0/go.mod h1:YRoo4H8PVmsu+E3Ou7cqLVH8oXWIHVoX0jqUWALQhfY= golang.org/x/time v0.8.0 h1:9i3RxcPv3PZnitoVGMPDKZSq1xW1gK1Xy3ArNOGZfEg= golang.org/x/time v0.8.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= + gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= diff --git a/login/internal/cmd/cmd.go b/login/internal/cmd/cmd.go index 757cec089..23669d4a1 100644 --- a/login/internal/cmd/cmd.go +++ b/login/internal/cmd/cmd.go @@ -12,11 +12,13 @@ import ( "github.com/gogf/gf/v2/net/ghttp" "github.com/gogf/gf/v2/os/gcmd" "github.com/gogf/gf/v2/os/gfile" + "github.com/xiaoqidun/limit" "golang.org/x/time/rate" ) var ( - Main = gcmd.Command{ + limiter = limit.New() + Main = gcmd.Command{ Name: "main", Usage: "main", Brief: "start http server", @@ -25,7 +27,10 @@ var ( if cool.IsRedisMode { go cool.ListenFunc(ctx) } + // 1. 创建一个新的 Limiter 实例 + // 2. 确保在程序退出前优雅地停止后台任务，这非常重要 + defer limiter.Stop() s := g.Server() s.Use(Limiter, ghttp.MiddlewareHandlerResponse) s.EnableAdmin() @@ -52,12 +57,12 @@ func beforeServeHook(r *ghttp.Request) { r.Response.CORSDefault() } -var limiter = rate.NewLimiter(rate.Limit(10), 1) - // Limiter is a middleware that implements rate limiting for all HTTP requests. // It returns HTTP 429 (Too Many Requests) when the rate limit is exceeded. func Limiter(r *ghttp.Request) { - if !limiter.Allow() { + rateLimiter := limiter.Get(r.GetClientIp(), rate.Limit(10), 2) + + if !rateLimiter.Allow() { r.Response.WriteStatusExit(429) // Return 429 Too Many Requests r.ExitAll() } diff --git a/login/main.go b/login/main.go index a813b6521..f051660b6 100644 --- a/login/main.go +++ b/login/main.go @@ -30,7 +30,6 @@ import ( ) func main() { - //Test_kick() cmd.Main.Run(gctx.New()) }