别等大佬了——用TryLock实现可重入锁

Posted 2023-01-22 文大侠666

Go1.18 中，Russ Cox 大佬终于妥协，Mutex新增了 TryLock 方法，历时9年。

生产中，大家常说的go中Mutex有两个痛点

• 锁是不可重入的
• 不支持TryLock方法

同样是大佬抗拒的痛点，TryLock有了，可重入锁还会远吗？TryLock没啥好说的，简单介绍下，再看看TryLock对优化可重入锁实现的作用，自己动手，丰衣足食，别等大佬了，人生有多少个9年！

要是本文对您有帮助的话，欢迎【关注】作者，【点赞】+【收藏】，保持交流！

TryLock

TryLock 常见的场景

尽管引入了TryLock，具体的应用场景并不明确，官方的说法是因为很多库实现了自己的TryLock，与其让你们瞎搞，不如我提供一个官方的版本好了，算是和Java对齐了。
一种可能的场景
有时我们希望尝试获取锁，如果获取到了则继续执行，如果获取不到，我们也不想阻塞住，而是去调用其它的逻辑，这个时候我们就想要 TryLock 方法：即尝试获取锁，获取不到也不堵塞。
这里不做过多讨论，大家知道有这个方法就够了。

TryLock实现原理

如下，因为Mutext的实现原理就是基于自旋+位标记的方式，TryLock的官方实现很简单，就是判断下位状态就够了，和之前官方TryLock实现之前的自定义实现差不多。

func (m *Mutex) TryLock() bool 
    // 已加锁/饥饿状态返回false
    old := m.state
    if old&(mutexLocked|mutexStarving) != 0 
        return false
    

    // 竞争失败则返回false，否则标记锁状态
    if !atomic.CompareAndSwapInt32(&m.state, old, old|mutexLocked) 
        return false
    

    if race.Enabled 
        race.Acquire(unsafe.Pointer(m))
    
    return true

可重入锁

Mutex不可重入导致死锁

用惯了java等中的mutex，在使用go中并发锁很容易踩中这个坑，最常见场景如下：
同一个协程中，实现一个加锁方法，嵌套调用另一个加锁方法，运行出现死锁。

type Calc struct 
	m sync.Mutex


func NewCalc() *Calc 
	return &Calcsync.Mutex


func (c *Calc) Do(wg *sync.WaitGroup) 
	defer wg.Done()

	c.m.Lock()

	log.Println("Do start")
	c.detailDo()
	log.Println("Do end")

	c.m.Unlock()


func (c *Calc) detailDo() 
	c.m.Lock()

	log.Println("detailDo start")

	c.m.Unlock()


func TestRecursiveLock(t *testing.T) 
	c := NewCalc()

	wg := sync.WaitGroup
	wg.Add(1)

	go c.Do(&wg)

	wg.Wait()

这里看起来好像没必要，单实际代码中，往往是函数特别长逻辑特别复杂公用多个锁的业务场景，可能还存在多个历史版本，往往很难仔细梳理逻辑，这时候很多时候就是无脑加锁，就会出现这种可重入的要求，毕竟写业务代码，上线紧时能加锁解决的都不是事。

但是，Go强调大道至简，做正确的事。关于为什么不支持，Russ Cox说了很多，总结就一句话——业务自己保证别出现嵌套调用，别TMD瞎加锁。

自己实现可重入锁

当然我们要尽可能按照Go最佳实践来，不要出现锁的重入，但是难免出现部分场景下需要这个特性，毕竟隔壁Java都有，凭啥Go不能有。

可重入锁实现也不难，一个通用的实现是

1. 锁对象增加当前持有的协程id和重入次数
2. 每次Lock，判断当前锁被同一协程持有则重入次数+1，不同协程持有则等待
3. UnLock时，重入次数-1直到=0时释放锁

如下

type RecursiveMutex struct 
	internalMutex    sync.Mutex // 保护 currentGoRoutine + lockCount
	currentGoRoutine int64      // 当前持有锁的协程 id
	lockCount        uint64     // 当前锁重复加锁次数


func (rm *RecursiveMutex) Lock() 
	// 获取当前协程id
	goRoutineID := gls.GoID()

	// 自旋锁实现mutex，效率相对较低
    // 一直自旋等到 currentGoRoutine=0 则成功获取锁
	for 
		rm.internalMutex.Lock()

		if rm.currentGoRoutine == 0 
			// 锁当前没被任何协程持有，加锁成功
			rm.currentGoRoutine = goRoutineID
			break
		 else if rm.currentGoRoutine == goRoutineID 
			// 当前协程重复加锁，直接退出，加锁成功
			break
		 else 
			// 不同的协程获取锁，等待其它持有锁的协程释放(currentGoRoutine=0)
			rm.internalMutex.Unlock()

			time.Sleep(time.Millisecond)
			continue
		
	

	// 增加当前锁重入次数
	rm.lockCount++
	rm.internalMutex.Unlock()


func (rm *RecursiveMutex) Unlock() 
	rm.internalMutex.Lock()
	defer rm.internalMutex.Unlock()

	// 减少锁重入次数，=0时释放锁
	rm.lockCount--
	if rm.lockCount == 0 
		rm.currentGoRoutine = 0
	

之前的Mutex换成RecursiveMutex，如下测试运行，完美！

func TestRecursiveLock2(t *testing.T) 
	c := NewCalc2()

	wg := sync.WaitGroup
	wg.Add(2)

	go c.Do(&wg)
	go c.Do(&wg)

	wg.Wait()

TryLock优化可重入锁

如上实现，每次自旋等待都伴随加锁解锁，且看起来相对不直观，这里借助TryLock优化下，如下

type RecursiveMutex2 struct 
	internalMutex    sync.Mutex // 不同协程竞争这一把锁
	currentGoRoutine int64      // 当前持有锁的协程 id
	lockCount        uint64     // 当前锁重复加锁次数


func (rm *RecursiveMutex2) Lock() 
	// 获取当前协程id
	goRoutineID := gls.GoID()

	// 自旋等待，直到TryLock=true
	for 
		bLock := rm.internalMutex.TryLock()

		if bLock 
			// 锁当前没被任何协程持有，加锁成功
			atomic.StoreInt64(&rm.currentGoRoutine, goRoutineID)
			break
		 else 
			if atomic.CompareAndSwapInt64(&rm.currentGoRoutine, goRoutineID, goRoutineID) 
				// 当前协程重复加锁，直接退出，加锁成功
				break
			 else 
				// 不同的协程获取锁，等待其它持有锁的协程释放(TryLock=true)
				time.Sleep(time.Millisecond)
				continue
			
		
	

	// 增加当前锁重入次数
	rm.lockCount++


func (rm *RecursiveMutex2) Unlock() 
	// 减少锁重入次数，=0时释放锁
	rm.lockCount--
	if rm.lockCount == 0 
		atomic.StoreInt64(&rm.currentGoRoutine, 0)
		rm.internalMutex.Unlock()
	

相对之前internalMutex用作竞争保护，这里internalMutex用于锁的竞争状态获取，实现上更直观，效率也相对较高。

参考

代码 https://gitee.com/wenzhou1219/go-in-prod/tree/master/mutex_lock

TryLock历史 https://zhuanlan.zhihu.com/p/467654600
官方TryLock之前的自定义实现参考 https://colobu.com/2017/03/09/implement-TryLock-in-Go/
Mutex实现原理解析 https://zhuanlan.zhihu.com/p/534617247

为什么不支持可重入锁
https://zhuanlan.zhihu.com/p/447295720
https://stackoverflow.com/questions/14670979/recursive-locking-in-go
老版本可重入锁实现
https://medium.com/@bytecraze.com/recursive-locking-in-go-9c1c2a106a38