mirror of
https://github.com/xiaoqidun/klock.git
synced 2025-10-11 11:40:23 +08:00
feat(完整功能): 以Apache License 2.0协议开源
This commit is contained in:
13
NOTICE
Normal file
13
NOTICE
Normal 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.
|
50
README.md
50
README.md
@@ -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
142
example_test.go
Normal 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
304
klock.go
Normal 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
443
klock_test.go
Normal 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)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
Reference in New Issue
Block a user