test/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2024-11-29 18:38:44 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

Make KeyLock generic (#27854)

Browse Source 
Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
This commit is contained in:
congqixia 2023-10-23 17:40:11 +08:00 committed by GitHub
parent 0c33ddb7c5
commit 4a6790b63a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 14 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
pkg/util/lock/key_lock.go
View File

@ -45,19 +45,19 @@ func newRefLock() *RefLock {
return &c
}
type KeyLock struct {
type KeyLock[K comparable] struct {
keyLocksMutex sync.Mutex
refLocks map[string]*RefLock
refLocks map[K]*RefLock
}
func NewKeyLock() *KeyLock {
keyLock := KeyLock{
refLocks: make(map[string]*RefLock),
func NewKeyLock[K comparable]() *KeyLock[K] {
keyLock := KeyLock[K]{
refLocks: make(map[K]*RefLock),
}
return &keyLock
}
func (k *KeyLock) Lock(key string) {
func (k *KeyLock[K]) Lock(key K) {
k.keyLocksMutex.Lock()
// update the key map
if keyLock, ok := k.refLocks[key]; ok {
@ -76,12 +76,12 @@ func (k *KeyLock) Lock(key string) {
}
}
func (k *KeyLock) Unlock(lockedKey string) {
func (k *KeyLock[K]) Unlock(lockedKey K) {
k.keyLocksMutex.Lock()
defer k.keyLocksMutex.Unlock()
keyLock, ok := k.refLocks[lockedKey]
if !ok {
log.Warn("Unlocking non-existing key", zap.String("key", lockedKey))
log.Warn("Unlocking non-existing key", zap.Any("key", lockedKey))
return
}
keyLock.unref()
@ -91,7 +91,7 @@ func (k *KeyLock) Unlock(lockedKey string) {
keyLock.mutex.Unlock()
}
func (k *KeyLock) RLock(key string) {
func (k *KeyLock[K]) RLock(key K) {
k.keyLocksMutex.Lock()
// update the key map
if keyLock, ok := k.refLocks[key]; ok {
@ -110,12 +110,12 @@ func (k *KeyLock) RLock(key string) {
}
}
func (k *KeyLock) RUnlock(lockedKey string) {
func (k *KeyLock[K]) RUnlock(lockedKey K) {
k.keyLocksMutex.Lock()
defer k.keyLocksMutex.Unlock()
keyLock, ok := k.refLocks[lockedKey]
if !ok {
log.Warn("Unlocking non-existing key", zap.String("key", lockedKey))
log.Warn("Unlocking non-existing key", zap.Any("key", lockedKey))
return
}
keyLock.unref()
@ -125,7 +125,7 @@ func (k *KeyLock) RUnlock(lockedKey string) {
keyLock.mutex.RUnlock()
}
func (k *KeyLock) size() int {
func (k *KeyLock[K]) size() int {
k.keyLocksMutex.Lock()
defer k.keyLocksMutex.Unlock()
return len(k.refLocks)

4
pkg/util/lock/key_lock_test.go
View File

@ -11,7 +11,7 @@ import (
func TestKeyLock(t *testing.T) {
keys := []string{"Milvus", "Blazing", "Fast"}
keyLock := NewKeyLock()
keyLock := NewKeyLock[string]()
keyLock.Lock(keys[0])
keyLock.Lock(keys[1])
@ -46,7 +46,7 @@ func TestKeyLock(t *testing.T) {
func TestKeyRLock(t *testing.T) {
keys := []string{"Milvus", "Blazing", "Fast"}
keyLock := NewKeyLock()
keyLock := NewKeyLock[string]()
keyLock.RLock(keys[0])
keyLock.RLock(keys[0])