SegmentKeysLock 引入 TryLock 和 TryRLock (#230)

syncx/segment_key_lock.go

@@ -22,50 +22,59 @@ import (
 // SegmentKeysLock 部分key lock结构定义
 type SegmentKeysLock struct {
 	locks []*sync.RWMutex
+	size  uint32
 }
 
 // NewSegmentKeysLock 创建 SegmentKeysLock 示例
-func NewSegmentKeysLock(size int) *SegmentKeysLock {
+func NewSegmentKeysLock(size uint32) *SegmentKeysLock {
 	locks := make([]*sync.RWMutex, size)
 	for i := range locks {
 		locks[i] = &sync.RWMutex{}
 	}
 	return &SegmentKeysLock{
 		locks: locks,
+		size:  size,
 	}
 }
 
 // hash 索引锁的hash函数
 func (s *SegmentKeysLock) hash(key string) uint32 {
 	h := fnv.New32a()
-	h.Write([]byte(key))
+	_, _ = h.Write([]byte(key))
 	return h.Sum32()
 }
 
 // RLock 读锁加锁
 func (s *SegmentKeysLock) RLock(key string) {
-	hash := s.hash(key)
-	lock := s.locks[hash%uint32(len(s.locks))]
-	lock.RLock()
+	s.getLock(key).RLock()
+}
+
+// TryRLock 试着加读锁，加锁成功会返回
+func (s *SegmentKeysLock) TryRLock(key string) bool {
+	return s.getLock(key).TryRLock()
 }
 
 // RUnlock 读锁解锁
 func (s *SegmentKeysLock) RUnlock(key string) {
-	hash := s.hash(key)
-	lock := s.locks[hash%uint32(len(s.locks))]
-	lock.RUnlock()
+	s.getLock(key).RUnlock()
 }
 
 // Lock 写锁加锁
 func (s *SegmentKeysLock) Lock(key string) {
-	hash := s.hash(key)
-	lock := s.locks[hash%uint32(len(s.locks))]
-	lock.Lock()
+	s.getLock(key).Lock()
+}
+
+// TryLock 试着加锁，加锁成功会返回 true
+func (s *SegmentKeysLock) TryLock(key string) bool {
+	return s.getLock(key).TryLock()
 }
 
 // Unlock 写锁解锁
 func (s *SegmentKeysLock) Unlock(key string) {
+	s.getLock(key).Unlock()
+}
+
+func (s *SegmentKeysLock) getLock(key string) *sync.RWMutex {
 	hash := s.hash(key)
-	lock := s.locks[hash%uint32(len(s.locks))]
-	lock.Unlock()
+	return s.locks[hash%s.size]
 }

syncx/segment_key_lock_test.go

@@ -15,60 +15,53 @@
 package syncx
 
 import (
-	"sync"
-	"sync/atomic"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 )
 
-func TestSegmentKeysLock(t *testing.T) {
-	s := NewSegmentKeysLock(10)
-	key := "test_key"
-	var wg sync.WaitGroup
-	wg.Add(2)
-	var writeDone atomic.Bool
-	var readStarted atomic.Bool
-	val := false
-	cond := sync.NewCond(&sync.Mutex{})
-	cond.L.Lock()
+// 通过 TryLock 和 TryRLock 来判定加锁问题
+// 也就是只判定我们拿到了正确的锁，但是没有判定并发与互斥
 
-	// 写 goroutine
-	go func() {
-		defer wg.Done()
-		s.Lock(key)
-		val = true // 加写锁写
-		s.Unlock(key)
-		writeDone.Store(true)
-		cond.Broadcast()
-	}()
+// TestNewSegmentKeysLock_Lock 测试 Lock, UnLock 和 TryLock
+func TestNewSegmentKeysLock_Lock(t *testing.T) {
+	l := NewSegmentKeysLock(8)
+	key1 := "key1"
+	l.Lock(key1)
+	// 必然加锁失败
+	assert.False(t, l.TryLock(key1))
+	// 读锁也失败
+	assert.False(t, l.TryRLock(key1))
+	key2 := "key2"
+	// 加锁成功
+	assert.True(t, l.TryLock(key2))
+	// 解锁不会触发 panic
+	defer l.Unlock(key2)
 
-	// 读 goroutine
-	go func() {
-		defer wg.Done()
-		cond.L.Lock()
-		defer cond.L.Unlock()
-
-		// 等待写操作完成
-		for !writeDone.Load() {
-			cond.Wait()
-		}
-
-		readStarted.Store(true)
-		cond.Broadcast()
-		s.RLock(key)
-		assert.Equal(t, true, val, "Read lock err") // 加读锁读
-		defer s.RUnlock(key)
-	}()
-
-	// 等待读操作开始
-	for !readStarted.Load() {
-		cond.Wait()
-	}
+	// 释放锁
+	l.Unlock(key1)
+	// 此时应该预期自己可以再次加锁
+	assert.True(t, l.TryLock(key1))
+}
 
-	// 检查写操作是否已完成，防止意外情况导致读优先写发生
-	assert.Equal(t, true, writeDone.Load(), "Write operation did not complete before read operation started")
+func TestNewSegmentKeysLock_RLock(t *testing.T) {
+	l := NewSegmentKeysLock(8)
+	key1, key2 := "key1", "key2"
+	l.RLock(key1)
+	// 必然加锁失败
+	assert.False(t, l.TryLock(key1))
+	// 读锁可以成功
+	assert.True(t, l.TryRLock(key1))
+	// 加锁成功
+	assert.True(t, l.TryRLock(key2))
+	// 解锁不会触发 panic
+	defer l.RUnlock(key2)
 
-	cond.L.Unlock()
-	wg.Wait()
+	// 释放读锁
+	l.RUnlock(key1)
+	// 此时还有一个读锁没有释放
+	assert.False(t, l.TryLock(key1))
+	// 再次释放读锁
+	l.RUnlock(key1)
+	assert.True(t, l.TryLock(key1))
 }