diff --git a/NOTICE b/NOTICE new file mode 100644 index 0000000..73d68a9 --- /dev/null +++ b/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/README.md b/README.md index e84614a..7a4492f 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,50 @@ # klock -一个基于键的高性能、并发安全的 Go 语言读写锁 +一个键级别的高性能、并发安全的 Go 语言读写锁 + +klock 可以为每一个数据(由一个唯一的“键”来标识)提供一把专属的锁。在高并发时,这意味着操作不同数据的任务可以并行,互不影响,从而大幅提升程序性能。它特别适合需要精细化控制单个资源访问的场景。 + +# 安装指南 +```shell +go get -u github.com/xiaoqidun/klock +``` + +# 快速开始 +下面的示例演示了 klock 的核心作用:保护一个共享计数器在并发修改下的数据一致性 +```go +package main + +import ( + "fmt" + "sync" + + "github.com/xiaoqidun/klock" +) + +func main() { + // 1. 创建 klock 实例 + kl := klock.New() + // 2. 准备一个共享的计数器和用于它的锁的键 + var counter int + lockKey := "counter_lock" + // 3. 模拟高并发修改计数器 + var wg sync.WaitGroup + concurrency := 1000 // 增加并发量以突显问题 + wg.Add(concurrency) + for i := 0; i < concurrency; i++ { + go func() { + defer wg.Done() + kl.Lock(lockKey) // 获取锁 + counter++ // 安全修改 + kl.Unlock(lockKey) // 释放锁 + }() + } + wg.Wait() + // 4. 验证结果 + // 如果没有 klock 保护,由于竞态条件,counter 的最终值将是一个小于1000的不确定数字。 + // 有了 klock,结果一定是 1000。 + fmt.Printf("最终结果: %d\n", counter) +} +``` + +# 授权协议 +本项目使用 [Apache License 2.0](https://github.com/xiaoqidun/klock/blob/main/LICENSE) 授权协议 \ No newline at end of file diff --git a/example_test.go b/example_test.go new file mode 100644 index 0000000..8aceeb8 --- /dev/null +++ b/example_test.go @@ -0,0 +1,142 @@ +// 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 klock_test + +import ( + "fmt" + "sync" + "time" + + "github.com/xiaoqidun/klock" +) + +// ExampleKeyLock_Lock 演示了如何使用 Lock 和 Unlock 来保护对共享资源的并发写操作。 +func ExampleKeyLock_Lock() { + kl := klock.New() + var wg sync.WaitGroup + // 共享数据 + sharedData := make(map[string]int) + // 模拟10个并发请求更新同一个键的值 + for i := 0; i < 10; i++ { + wg.Add(1) + go func() { + defer wg.Done() + // 为 "dataKey" 加锁,确保同一时间只有一个 goroutine 能修改其值 + kl.Lock("dataKey") + sharedData["dataKey"]++ + kl.Unlock("dataKey") + }() + } + wg.Wait() + fmt.Printf("键 'dataKey' 的最终计数值是: %d\n", sharedData["dataKey"]) + // Output: 键 'dataKey' 的最终计数值是: 10 +} + +// ExampleKeyLock_RLock 演示了多个 goroutine 如何同时获取读锁以并发地读取数据。 +func ExampleKeyLock_RLock() { + kl := klock.New() + sharedResource := "这是一个共享资源" + // 启动5个 goroutine 并发读取数据 + var wg sync.WaitGroup + for i := 0; i < 5; i++ { + wg.Add(1) + go func(gid int) { + defer wg.Done() + kl.RLock("resourceKey") + fmt.Printf("Goroutine %d 读取资源: %s\n", gid, sharedResource) + kl.RUnlock("resourceKey") + }(i) + } + wg.Wait() + // 注意:由于 goroutine 调度顺序不确定,输出的顺序可能不同。 + // 但这个示例的核心是展示它们可以并发执行,而不会像写锁一样互相等待。 +} + +// ExampleKeyLock_TryLock 演示了如何尝试非阻塞地获取锁。 +// 如果锁已被占用,它不会等待,而是立即返回 false。 +func ExampleKeyLock_TryLock() { + kl := klock.New() + key := "resource_key" + // 第一次尝试,应该成功 + if kl.TryLock(key) { + fmt.Println("第一次 TryLock 成功获取锁") + // 第二次尝试,因为锁已被持有,所以应该失败 + if !kl.TryLock(key) { + fmt.Println("第二次 TryLock 失败,因为锁已被占用") + } + kl.Unlock(key) + } + // Output: + // 第一次 TryLock 成功获取锁 + // 第二次 TryLock 失败,因为锁已被占用 +} + +// ExampleKeyLock_LockWithTimeout 演示了如何在指定时间内尝试获取锁,避免无限期等待。 +func ExampleKeyLock_LockWithTimeout() { + kl := klock.New() + key := "task_key" + // 主 goroutine 先获取锁 + kl.Lock(key) + fmt.Println("主 goroutine 持有锁") + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + // 尝试在 10ms 内获取锁,此时主 goroutine 仍持有锁,所以会失败 + fmt.Println("Goroutine 尝试在 10ms 内获取锁...") + if !kl.LockWithTimeout(key, 10*time.Millisecond) { + fmt.Println("Goroutine 获取锁超时,任务取消") + } + }() + // 等待 20ms 后释放锁 + time.Sleep(20 * time.Millisecond) + kl.Unlock(key) + fmt.Println("主 goroutine 释放锁") + wg.Wait() + // Output: + // 主 goroutine 持有锁 + // Goroutine 尝试在 10ms 内获取锁... + // Goroutine 获取锁超时,任务取消 + // 主 goroutine 释放锁 +} + +// ExampleKeyLock_Status 演示了如何查询一个键的当前锁定状态。 +func ExampleKeyLock_Status() { + kl := klock.New() + key := "status_key" + var wg sync.WaitGroup + kl.Lock(key) + // 启动两个 goroutine 在后台等待锁 + for i := 0; i < 2; i++ { + wg.Add(1) + go func() { + defer wg.Done() + kl.Lock(key) + time.Sleep(1 * time.Millisecond) + kl.Unlock(key) + }() + } + // 等待一小段时间,确保两个 goroutine 已经处于等待状态 + time.Sleep(10 * time.Millisecond) + holders, waiters := kl.Status(key) + fmt.Printf("当前持有者: %d, 等待者: %d\n", holders, waiters) + kl.Unlock(key) + wg.Wait() + holders, waiters = kl.Status(key) + fmt.Printf("所有任务完成后,持有者: %d, 等待者: %d\n", holders, waiters) + // Output: + // 当前持有者: 1, 等待者: 2 + // 所有任务完成后,持有者: 0, 等待者: 0 +} diff --git a/klock.go b/klock.go new file mode 100644 index 0000000..7983ae8 --- /dev/null +++ b/klock.go @@ -0,0 +1,304 @@ +// 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 klock 提供了一个基于键的高性能、并发安全的 Go 语言读写锁库。 +// +// 它允许为不同的键(如资源ID、任务标识符)创建独立的读写锁, +// 在高并发场景下,通过分片技术减少锁竞争,从而显著提升性能。 +// 该库能够自动管理锁对象的生命周期,防止在长时间运行的服务中发生内存泄漏。 +// +// 功能完善,支持阻塞式锁定 (Lock/RLock)、非阻塞式锁定 (TryLock/TryRLock)、 +// 带超时的锁定 (LockWithTimeout/RLockWithTimeout) 以及锁状态查询 (Status)。 +package klock + +import ( + "sync" + "time" +) + +// shardCount 是分片数量,必须是2的幂次方以获得最佳性能。 +const shardCount uint32 = 256 + +// fnv-1a 哈希算法使用的常量 +const ( + prime32 = 16777619 + offset32 = 2166136261 +) + +// Config 用于自定义 KeyLock 实例的行为。 +// 通过此配置可以调整超时锁定功能的轮询策略,以适应不同的负载场景。 +type Config struct { + // MaxPollInterval 是超时锁定功能中轮询的最大等待间隔。 + MaxPollInterval time.Duration + // InitialPollInterval 是超时锁定功能中轮询的初始等待间隔。 + InitialPollInterval time.Duration +} + +// defaultConfig 返回一个包含推荐默认值的配置。 +func defaultConfig() Config { + return Config{ + MaxPollInterval: 16 * time.Millisecond, + InitialPollInterval: 1 * time.Millisecond, + } +} + +// KeyLock 提供了基于字符串键的高性能、分片读写锁功能。 +type KeyLock struct { + shards []*keyLockShard + config Config +} + +// lockEntry 包含一个读写锁和两个引用计数器,用于安全地自动清理。 +type lockEntry struct { + rw sync.RWMutex + holders int // 持有锁的 goroutine 数量 + waiters int // 正在等待获取锁的 goroutine 数量 +} + +// keyLockShard 代表一个独立的锁分片。 +// 每个分片都有自己的互斥锁来保护其内部的锁映射。 +type keyLockShard struct { + mu sync.Mutex + locks map[string]*lockEntry +} + +// New 使用默认配置创建一个新的高性能 KeyLock 实例。 +func New() *KeyLock { + return NewWithConfig(defaultConfig()) +} + +// NewWithConfig 使用自定义配置创建一个新的高性能 KeyLock 实例。 +func NewWithConfig(config Config) *KeyLock { + kl := &KeyLock{ + shards: make([]*keyLockShard, shardCount), + config: config, + } + for i := uint32(0); i < shardCount; i++ { + kl.shards[i] = &keyLockShard{ + locks: make(map[string]*lockEntry), + } + } + return kl +} + +// hash 使用 FNV-1a 算法计算字符串的32位哈希值,此实现避免了不必要的内存分配。 +func hash(s string) uint32 { + h := uint32(offset32) + for i := 0; i < len(s); i++ { + h ^= uint32(s[i]) + h *= prime32 + } + return h +} + +// getShard 根据键的哈希值返回对应的分片。 +func (kl *KeyLock) getShard(key string) *keyLockShard { + return kl.shards[hash(key)&(shardCount-1)] +} + +// commitLock 将一个成功获取到锁的 goroutine 从“等待者”状态转换成“持有者”状态。 +func (kl *KeyLock) commitLock(key string, le *lockEntry) { + shard := kl.getShard(key) + shard.mu.Lock() + le.waiters-- + le.holders++ + shard.mu.Unlock() +} + +// cancelLock 在获取锁失败(如 TryLock 失败或超时)后,安全地撤销“等待”状态。 +// 如果撤销后锁不再被需要,则会将其清理。 +func (kl *KeyLock) cancelLock(key string, le *lockEntry) { + shard := kl.getShard(key) + shard.mu.Lock() + le.waiters-- + if le.holders == 0 && le.waiters == 0 { + delete(shard.locks, key) + } + shard.mu.Unlock() +} + +// prepareLock 获取或创建指定键的锁条目,并增加其等待者计数。 +// 这是所有加锁操作的第一步,确保了锁条目在后续操作中不会被意外回收。 +func (kl *KeyLock) prepareLock(key string) *lockEntry { + shard := kl.getShard(key) + shard.mu.Lock() + le, ok := shard.locks[key] + if !ok { + le = &lockEntry{} + shard.locks[key] = le + } + le.waiters++ + shard.mu.Unlock() + return le +} + +// releaseLock 是一个通用的解锁辅助函数。 +// 它处理持有者计数、锁条目的生命周期管理(清理)以及实际的读写锁释放。 +func (kl *KeyLock) releaseLock(key string, unlockFunc func(*sync.RWMutex)) { + shard := kl.getShard(key) + shard.mu.Lock() + le, ok := shard.locks[key] + if !ok { + shard.mu.Unlock() + panic("klock: unlock of unlocked key") + } + le.holders-- + isLast := le.holders == 0 && le.waiters == 0 + if isLast { + delete(shard.locks, key) + } + shard.mu.Unlock() + unlockFunc(&le.rw) +} + +// Lock 为指定的键获取一个写锁。 +// 如果锁已被其他 goroutine 持有,则此调用将阻塞直到锁可用。 +func (kl *KeyLock) Lock(key string) { + le := kl.prepareLock(key) + le.rw.Lock() + kl.commitLock(key, le) +} + +// TryLock 尝试非阻塞地为指定的键获取写锁。 +// 如果成功获取锁,则返回 true;否则立即返回 false。 +func (kl *KeyLock) TryLock(key string) bool { + le := kl.prepareLock(key) + if le.rw.TryLock() { + kl.commitLock(key, le) + return true + } + kl.cancelLock(key, le) + return false +} + +// LockWithTimeout 尝试在给定的时间内为指定的键获取写锁。 +// 如果在超时前成功获取锁,则返回 true;否则返回 false。 +func (kl *KeyLock) LockWithTimeout(key string, timeout time.Duration) bool { + le := kl.prepareLock(key) + // 立即尝试一次,以避免在锁可用时产生不必要的延迟。 + if le.rw.TryLock() { + kl.commitLock(key, le) + return true + } + var acquired bool + defer func() { + if !acquired { + kl.cancelLock(key, le) + } + }() + timer := time.NewTimer(timeout) + defer timer.Stop() + pollInterval := kl.config.InitialPollInterval + maxPollInterval := kl.config.MaxPollInterval + for { + select { + case <-timer.C: + return false + default: + time.Sleep(pollInterval) + if le.rw.TryLock() { + kl.commitLock(key, le) + acquired = true + return true + } + pollInterval *= 2 + if pollInterval > maxPollInterval { + pollInterval = maxPollInterval + } + } + } +} + +// RLock 为指定的键获取一个读锁。 +// 如果锁已被其他 goroutine 持有写锁,则此调用将阻塞直到锁可用。 +// 多个 goroutine 可以同时持有同一个键的读锁。 +func (kl *KeyLock) RLock(key string) { + le := kl.prepareLock(key) + le.rw.RLock() + kl.commitLock(key, le) +} + +// TryRLock 尝试非阻塞地为指定的键获取读锁。 +// 如果成功获取锁,则返回 true;否则立即返回 false。 +func (kl *KeyLock) TryRLock(key string) bool { + le := kl.prepareLock(key) + if le.rw.TryRLock() { + kl.commitLock(key, le) + return true + } + kl.cancelLock(key, le) + return false +} + +// RLockWithTimeout 尝试在给定的时间内为指定的键获取读锁。 +// 如果在超时前成功获取锁,则返回 true;否则返回 false。 +func (kl *KeyLock) RLockWithTimeout(key string, timeout time.Duration) bool { + le := kl.prepareLock(key) + // 立即尝试一次 + if le.rw.TryRLock() { + kl.commitLock(key, le) + return true + } + var acquired bool + defer func() { + if !acquired { + kl.cancelLock(key, le) + } + }() + timer := time.NewTimer(timeout) + defer timer.Stop() + pollInterval := kl.config.InitialPollInterval + maxPollInterval := kl.config.MaxPollInterval + for { + select { + case <-timer.C: + return false + default: + time.Sleep(pollInterval) + if le.rw.TryRLock() { + kl.commitLock(key, le) + acquired = true + return true + } + pollInterval *= 2 + if pollInterval > maxPollInterval { + pollInterval = maxPollInterval + } + } + } +} + +// Unlock 释放指定键的写锁。 +// 如果对未锁定的键调用 Unlock,将会引发 panic。 +func (kl *KeyLock) Unlock(key string) { + kl.releaseLock(key, (*sync.RWMutex).Unlock) +} + +// RUnlock 释放指定键的读锁。 +// 如果对未进行读锁定的键调用 RUnlock,将会引发 panic。 +func (kl *KeyLock) RUnlock(key string) { + kl.releaseLock(key, (*sync.RWMutex).RUnlock) +} + +// Status 返回指定键的当前锁状态。 +// 它返回两个值:持有该锁的 goroutine 数量,以及正在等待获取该锁的 goroutine 数量。 +func (kl *KeyLock) Status(key string) (holders int, waiters int) { + shard := kl.getShard(key) + shard.mu.Lock() + defer shard.mu.Unlock() + if le, ok := shard.locks[key]; ok { + return le.holders, le.waiters + } + return 0, 0 +} diff --git a/klock_test.go b/klock_test.go new file mode 100644 index 0000000..d99ae63 --- /dev/null +++ b/klock_test.go @@ -0,0 +1,443 @@ +// 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 klock + +import ( + "fmt" + "sync" + "sync/atomic" + "testing" + "time" +) + +// TestLockUnlock 测试基本的写锁获取与释放功能。 +func TestLockUnlock(t *testing.T) { + kl := New() + key := "test_key" + kl.Lock(key) + kl.Unlock(key) +} + +// TestRLockRUnlock 测试基本的读锁获取与释放功能,并验证多个读锁可以共存。 +func TestRLockRUnlock(t *testing.T) { + kl := New() + key := "test_key" + kl.RLock(key) + kl.RLock(key) + kl.RUnlock(key) + kl.RUnlock(key) +} + +// TestLockBlocksRLock 测试写锁会阻塞后续的读锁请求。 +func TestLockBlocksRLock(t *testing.T) { + kl := New() + key := "test_key" + ch := make(chan struct{}) + kl.Lock(key) + go func() { + kl.RLock(key) + kl.RUnlock(key) + close(ch) + }() + select { + case <-ch: + t.Fatal("RLock 应该被 Lock 阻塞") + case <-time.After(10 * time.Millisecond): + // 预期的行为 + } + kl.Unlock(key) + select { + case <-ch: + // 预期的行为 + case <-time.After(10 * time.Millisecond): + t.Fatal("RLock 在 Unlock 后应该被解除阻塞") + } +} + +// TestRLockBlocksLock 测试读锁会阻塞后续的写锁请求。 +func TestRLockBlocksLock(t *testing.T) { + kl := New() + key := "test_key" + ch := make(chan struct{}) + kl.RLock(key) + go func() { + kl.Lock(key) + kl.Unlock(key) + close(ch) + }() + select { + case <-ch: + t.Fatal("Lock 应该被 RLock 阻塞") + case <-time.After(10 * time.Millisecond): + // 预期的行为 + } + kl.RUnlock(key) + select { + case <-ch: + // 预期的行为 + case <-time.After(10 * time.Millisecond): + t.Fatal("Lock 在所有 RUnlock 后应该被解除阻塞") + } +} + +// TestConcurrentLock 测试多个 goroutine 并发获取同一个写锁的互斥性。 +func TestConcurrentLock(t *testing.T) { + kl := New() + key := "test_key" + var counter int32 + var wg sync.WaitGroup + numGoroutines := 100 + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func() { + defer wg.Done() + kl.Lock(key) + atomic.AddInt32(&counter, 1) + kl.Unlock(key) + }() + } + wg.Wait() + if counter != int32(numGoroutines) { + t.Errorf("期望计数器为 %d, 但得到 %d", numGoroutines, counter) + } +} + +// TestConcurrentRead 测试多个 goroutine 可以并发地持有读锁并读取数据。 +func TestConcurrentRead(t *testing.T) { + kl := New() + key := "test_key" + var wg sync.WaitGroup + numGoroutines := 100 + kl.Lock(key) + sharedData := 42 + kl.Unlock(key) + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func() { + defer wg.Done() + kl.RLock(key) + if sharedData != 42 { + t.Errorf("读取到错误数据: 得到 %d, 期望 42", sharedData) + } + kl.RUnlock(key) + }() + } + wg.Wait() +} + +// TestConcurrentReadWrite 测试读写锁在并发读写场景下的正确性。 +func TestConcurrentReadWrite(t *testing.T) { + kl := New() + key := "test_key" + var data int + var wg sync.WaitGroup + numGoroutines := 100 + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < numGoroutines; i++ { + kl.Lock(key) + data++ + kl.Unlock(key) + } + }() + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < numGoroutines; j++ { + kl.RLock(key) + _ = data // 只是读取 + kl.RUnlock(key) + } + }() + } + wg.Wait() + if data != numGoroutines { + t.Errorf("期望最终数据为 %d, 得到 %d", numGoroutines, data) + } +} + +// TestDifferentKeysDoNotBlock 测试对不同键的操作不会相互阻塞。 +func TestDifferentKeysDoNotBlock(t *testing.T) { + kl := New() + key1, key2 := "key1", "key2" + ch := make(chan struct{}) + kl.Lock(key1) + go func() { + kl.Lock(key2) + kl.Unlock(key2) + close(ch) + }() + select { + case <-ch: + // 预期的行为,不应该被阻塞 + case <-time.After(20 * time.Millisecond): + t.Fatal("锁定不同的键不应该被阻塞") + } + kl.Unlock(key1) +} + +// TestTryLock 测试非阻塞获取写锁的正确性。 +func TestTryLock(t *testing.T) { + kl := New() + key := "test_key" + if !kl.TryLock(key) { + t.Fatal("TryLock 在键未锁定时应该成功") + } + if kl.TryLock(key) { + t.Fatal("TryLock 在键已锁定时应该失败") + } + kl.Unlock(key) + if !kl.TryLock(key) { + t.Fatal("TryLock 在键解锁后应该成功") + } + kl.Unlock(key) +} + +// TestTryRLock 测试非阻塞获取读锁的正确性。 +func TestTryRLock(t *testing.T) { + kl := New() + key := "test_key" + if !kl.TryRLock(key) { + t.Fatal("TryRLock 在键未锁定时应该成功") + } + if !kl.TryRLock(key) { + t.Fatal("TryRLock 在键已被读锁定时应该成功") + } + kl.RUnlock(key) + kl.RUnlock(key) + kl.Lock(key) + if kl.TryRLock(key) { + t.Fatal("TryRLock 在键已被写锁定时应该失败") + } + kl.Unlock(key) +} + +// TestLockWithTimeout 测试带超时的写锁获取功能。 +func TestLockWithTimeout(t *testing.T) { + kl := New() + key := "test_key" + if !kl.LockWithTimeout(key, 10*time.Millisecond) { + t.Fatal("LockWithTimeout 在键未锁定时应该立即成功") + } + kl.Unlock(key) + kl.Lock(key) + go func() { + time.Sleep(5 * time.Millisecond) + kl.Unlock(key) + }() + if !kl.LockWithTimeout(key, 20*time.Millisecond) { + t.Fatal("LockWithTimeout 在锁于超时前被释放时应该成功") + } + kl.Unlock(key) + kl.Lock(key) + defer kl.Unlock(key) + if kl.LockWithTimeout(key, 10*time.Millisecond) { + t.Fatal("LockWithTimeout 在超时后应该失败") + } +} + +// TestRLockWithTimeout 测试带超时的读锁获取功能。 +func TestRLockWithTimeout(t *testing.T) { + kl := New() + key := "test_key" + kl.Lock(key) + go func() { + time.Sleep(10 * time.Millisecond) + kl.Unlock(key) + }() + if !kl.RLockWithTimeout(key, 20*time.Millisecond) { + t.Fatal("RLockWithTimeout 在锁于超时前被释放时应该成功") + } + kl.RUnlock(key) + kl.Lock(key) + defer kl.Unlock(key) + if kl.RLockWithTimeout(key, 10*time.Millisecond) { + t.Fatal("RLockWithTimeout 在超时后应该失败") + } +} + +// TestPanicOnUnlockOfUnlockedKey 测试对未锁定的键执行 Unlock 操作是否会引发 panic。 +func TestPanicOnUnlockOfUnlockedKey(t *testing.T) { + defer func() { + if r := recover(); r == nil { + t.Errorf("代码在 Unlock 未锁定的键时没有 panic") + } + }() + kl := New() + kl.Unlock("non_existent_key") +} + +// TestPanicOnRUnlockOfUnlockedKey 测试对未读锁定的键执行 RUnlock 操作是否会引发 panic。 +func TestPanicOnRUnlockOfUnlockedKey(t *testing.T) { + defer func() { + if r := recover(); r == nil { + t.Errorf("代码在 RUnlock 未读锁定的键时没有 panic") + } + }() + kl := New() + kl.RUnlock("non_existent_key") +} + +// TestStatus 测试 Status 方法是否能正确返回锁的持有者和等待者数量。 +func TestStatus(t *testing.T) { + kl := New() + key := "test_key" + var wg sync.WaitGroup + holders, waiters := kl.Status(key) + if holders != 0 || waiters != 0 { + t.Fatalf("期望 0 个持有者和 0 个等待者, 得到 %d 和 %d", holders, waiters) + } + kl.Lock(key) + holders, waiters = kl.Status(key) + if holders != 1 || waiters != 0 { + t.Fatalf("期望 1 个持有者和 0 个等待者, 得到 %d 和 %d", holders, waiters) + } + wg.Add(2) + go func() { + defer wg.Done() + kl.Lock(key) + kl.Unlock(key) + }() + go func() { + defer wg.Done() + kl.RLock(key) + kl.RUnlock(key) + }() + time.Sleep(10 * time.Millisecond) + holders, waiters = kl.Status(key) + if holders != 1 || waiters != 2 { + t.Fatalf("期望 1 个持有者和 2 个等待者, 得到 %d 和 %d", holders, waiters) + } + kl.Unlock(key) + wg.Wait() + holders, waiters = kl.Status(key) + if holders != 0 || waiters != 0 { + t.Fatalf("期望解锁后有 0 个持有者和 0 个等待者, 得到 %d 和 %d", holders, waiters) + } +} + +// TestLockChurn 测试在高并发和键频繁创建销毁的情况下,引用计数是否能正常工作,防止内存泄漏。 +func TestLockChurn(t *testing.T) { + kl := New() + numGoroutines := 100 + iterations := 1000 + var wg sync.WaitGroup + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(gid int) { + defer wg.Done() + for j := 0; j < iterations; j++ { + key := fmt.Sprintf("key-%d-%d", gid, j%10) // 限制键的数量以增加竞争 + kl.Lock(key) + kl.Unlock(key) + } + }(i) + } + wg.Wait() + for i := uint32(0); i < shardCount; i++ { + shard := kl.shards[i] + shard.mu.Lock() + if len(shard.locks) != 0 { + t.Errorf("分片 %d 存在锁泄漏, 数量: %d", i, len(shard.locks)) + } + shard.mu.Unlock() + } +} + +// --- 基准测试 --- + +// BenchmarkLockUnlock_SingleKey 测试单个键上高竞争的写锁性能。 +func BenchmarkLockUnlock_SingleKey(b *testing.B) { + kl := New() + key := "key" + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + kl.Lock(key) + kl.Unlock(key) + } + }) +} + +// BenchmarkRLockRUnlock_SingleKey 测试单个键上高竞争的读锁性能。 +func BenchmarkRLockRUnlock_SingleKey(b *testing.B) { + kl := New() + key := "key" + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + kl.RLock(key) + kl.RUnlock(key) + } + }) +} + +// BenchmarkLockUnlock_MultipleKeys 测试多个键上低竞争的写锁性能。 +func BenchmarkLockUnlock_MultipleKeys(b *testing.B) { + kl := New() + var counter uint64 + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + key := fmt.Sprintf("key-%d", atomic.AddUint64(&counter, 1)%1024) + kl.Lock(key) + kl.Unlock(key) + } + }) +} + +// BenchmarkRLockRUnlock_MultipleKeys 测试多个键上低竞争的读锁性能。 +func BenchmarkRLockRUnlock_MultipleKeys(b *testing.B) { + kl := New() + var counter uint64 + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + key := fmt.Sprintf("key-%d", atomic.AddUint64(&counter, 1)%1024) + kl.RLock(key) + kl.RUnlock(key) + } + }) +} + +// BenchmarkReadWrite_SingleKey 测试单个键上混合读写的性能。 +func BenchmarkReadWrite_SingleKey(b *testing.B) { + kl := New() + key := "key" + b.SetParallelism(10) + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + kl.RLock(key) + kl.RUnlock(key) + } + }) + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + kl.Lock(key) + kl.Unlock(key) + } + }) +} + +// BenchmarkTryLock_Contended 测试在锁已被持有的情况下,TryLock 的性能。 +func BenchmarkTryLock_Contended(b *testing.B) { + kl := New() + key := "key" + kl.Lock(key) + defer kl.Unlock(key) + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + kl.TryLock(key) + } + }) +}