Golang轻量级（细粒度）锁的实现

Z时代
2024-01-10
分类：综合

背景

最近在项目中遇到一个多线程处理db的一个场景，就是Check-And-Set，需要先查看redis里的一个key值，然后根据查看的结果，决定是否更新，这是典型的容易导致并发脏数据的场景，多个goroutine同时查看，都满足条件，最后都更新导致产生脏数据。因此需要加锁。

golang自身有sync.Mutex互斥锁，可以达到效果，但是这个锁的粒度太大。想要针对字符串加锁就有点无能为力了。

因此有了下面的sycn.Map的解决方案，实现了对字符串进行加锁，比sync.Mutex对struct加锁的粒度要小很多。

实现

原理很简单就是利用sync.Map是线程安全的，上锁就是给一个key-value，解锁，就是删掉对应的key-value


type KeyLock struct {
	m sync.Map
}
func (k *KeyLock) TryLock(key interface{}) bool {
	_, ok := k.m.LoadOrStore(key, struct{}{})
	return !ok
}
func (k *KeyLock) WaitLock(key interface{}, retry int) bool {
	for i := 0; i < retry; i++ {
		if k.TryLock(key) {
			return true
		} else {
			time.Sleep(10 * time.Microsecond)
		}
	}
	return false
}
func (k *KeyLock) UnLock(key interface{}) {
	k.m.Delete(key)
}

测试

package main
import (
	"fmt"
	"sync"
	"time"
)
func main() {
	//var tmpInt
	tmp := 0
	tmp2 := 0
	wg := sync.WaitGroup{}
	lock := &KeyLock{}
	for i:=0; i< 1000; i++{
		wg.Add(1)
		go func()  {
			for{
				if lock.TryLock("hello"){
					tmp ++
					lock.UnLock("hello")
					break
				}else{
					//time.Sleep(1*time.Second)
				}
			}
			wg.Done()
		}()
		wg.Add(1)
		go func() {
			tmp2 ++
			wg.Done()
		}()
	}
	wg.Wait()
	fmt.Println("tmp =", tmp)
	fmt.Println("tmp2 =", tmp2)
}