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

This commit is contained in:
2025-09-24 02:58:49 +08:00
parent 1a53a5ab8f
commit 0d90c09d71
5 changed files with 951 additions and 1 deletions

13
NOTICE Normal file
View File

@@ -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.

View File

@@ -1,2 +1,50 @@
# klock # 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) 授权协议

142
example_test.go Normal file
View File

@@ -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
}

304
klock.go Normal file
View File

@@ -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
}

443
klock_test.go Normal file
View File

@@ -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)
}
})
}