为 Go Mutex 实现 TryLock 方法

目录 [−]

使用 unsafe 操作指针
实现自旋锁
使用 channel 实现
性能比较
参考资料

Go标准库的sync/Mutex、RWMutex实现了sync/Locker接口，提供了Lock()和UnLock()方法，可以获取锁和释放锁，我们可以方便的使用它来控制我们对共享资源的并发控制上。

但是标准库中的Mutex.Lock的锁被获取后，如果在未释放之前再调用Lock则会被阻塞住，这种设计在有些情况下可能不能满足我的需求。有时候我们想尝试获取锁，如果获取到了，没问题继续执行，如果获取不到，我们不想阻塞住，而是去调用其它的逻辑，这个时候我们就想要TryLock方法了。

虽然很早(13年)就有人给Go开发组提需求了，但是这个请求并没有纳入官方库中，最终在官方库的清理中被关闭了，也就是官方库目前不会添加这个方法。

顺便说一句， sync/Mutex的源代码实现可以访问这里，它应该是实现了一种自旋(spin)加休眠的方式实现，有兴趣的读者可以阅读源码，或者阅读相关的文章，比如 Go Mutex 源码剖析。这不是本文要介绍的内容，读者可以找一些资料来阅读。

好了，转入正题，看看几种实现TryLock的方式吧。

使用 `unsafe` 操作指针

如果你查看sync/Mutex的代码，会发现Mutext的数据结构如下所示：

type Mutex struct {
	state int32
	sema  uint32
}

它使用state这个32位的整数来标记锁的占用，所以我们可以使用CAS来尝试获取锁。

代码实现如下：

const mutexLocked = 1 << iota
type Mutex struct {
	sync.Mutex
}
func (m *Mutex) TryLock() bool {
	return atomic.CompareAndSwapInt32((*int32)(unsafe.Pointer(&m.Mutex)), 0, mutexLocked)
}

使用起来和标准库的Mutex用法一样。

func main() {
	var m Mutex
	m.Lock()
	go func() {
		m.Lock()
	}()
	time.Sleep(time.Second)
	fmt.Printf("TryLock: %t\n", m.TryLock()) //false
	fmt.Printf("TryLock: %t\n", m.TryLock()) // false
	m.Unlock()
	fmt.Printf("TryLock: %t\n", m.TryLock()) //true
	fmt.Printf("TryLock: %t\n", m.TryLock()) //false
	m.Unlock()
	fmt.Printf("TryLock: %t\n", m.TryLock()) //true
	m.Unlock()
}

注意TryLock不是检查锁的状态，而是尝试获取锁，所以TryLock返回true的时候事实上这个锁已经被获取了。

实现自旋锁

上面一节给了我们启发，利用 uint32和CAS操作我们可以一个自定义的锁:

type SpinLock struct {
	f uint32
}
func (sl *SpinLock) Lock() {
	for !sl.TryLock() {
		runtime.Gosched()
	}
}
func (sl *SpinLock) Unlock() {
	atomic.StoreUint32(&sl.f, 0)
}
func (sl *SpinLock) TryLock() bool {
	return atomic.CompareAndSwapUint32(&sl.f, 0, 1)
}

整体来看，它好像是标准库的一个精简版，没有休眠和唤醒的功能。

当然这个自旋锁可以在大并发的情况下CPU的占用率可能比较高，这是因为它的Lock方法使用了自旋的方式，如果别人没有释放锁，这个循环会一直执行，速度可能更快但CPU占用率高。

当然这个版本还可以进一步的优化，尤其是在复制的时候。下面是一个优化的版本:

type spinLock uint32
func (sl *spinLock) Lock() {
	for !atomic.CompareAndSwapUint32((*uint32)(sl), 0, 1) {
		runtime.Gosched() //without this it locks up on GOMAXPROCS > 1
	}
}
func (sl *spinLock) Unlock() {
	atomic.StoreUint32((*uint32)(sl), 0)
}
func (sl *spinLock) TryLock() bool {
	return atomic.CompareAndSwapUint32((*uint32)(sl), 0, 1)
}
func SpinLock() sync.Locker {
	var lock spinLock
	return &lock
}

使用 channel 实现

另一种方式是使用channel:

type ChanMutex chan struct{}
func (m *ChanMutex) Lock() {
	ch := (chan struct{})(*m)
	ch <- struct{}{}
}
func (m *ChanMutex) Unlock() {
	ch := (chan struct{})(*m)
	select {
	case <-ch:
	default:
		panic("unlock of unlocked mutex")
	}
}
func (m *ChanMutex) TryLock() bool {
	ch := (chan struct{})(*m)
	select {
	case ch <- struct{}{}:
		return true
	default:
	}
	return false
}

有兴趣的同学可以关注我的同事写的库 lrita/gosync。

性能比较

首先看看上面三种方式和标准库中的Mutex、RWMutex的Lock和Unlock的性能比较：

BenchmarkMutex_LockUnlock-4         	100000000	        16.8 ns/op	       0 B/op	       0 allocs/op
BenchmarkRWMutex_LockUnlock-4       	50000000	        36.8 ns/op	       0 B/op	       0 allocs/op
BenchmarkUnsafeMutex_LockUnlock-4   	100000000	        16.8 ns/op	       0 B/op	       0 allocs/op
BenchmarkChannMutex_LockUnlock-4    	20000000	        65.6 ns/op	       0 B/op	       0 allocs/op
BenchmarkSpinLock_LockUnlock-4      	100000000	        18.6 ns/op	       0 B/op	       0 allocs/op

可以看到单线程(goroutine)的情况下｀spinlock｀并没有比标准库好多少，反而差一点,并发测试的情况比较好，如下表中显示，这是符合预期的。

unsafe方式和标准库差不多。

channel方式的性能就比较差了。

BenchmarkMutex_LockUnlock_C-4         	20000000	        75.3 ns/op	       0 B/op	       0 allocs/op
BenchmarkRWMutex_LockUnlock_C-4       	20000000	       100 ns/op	       0 B/op	       0 allocs/op
BenchmarkUnsafeMutex_LockUnlock_C-4   	20000000	        75.3 ns/op	       0 B/op	       0 allocs/op
BenchmarkChannMutex_LockUnlock_C-4    	10000000	       231 ns/op	       0 B/op	       0 allocs/op
BenchmarkSpinLock_LockUnlock_C-4      	50000000	        32.3 ns/op	       0 B/op	       0 allocs/op

再看看三种实现TryLock方法的锁的性能：

1
2
3

BenchmarkUnsafeMutex_Trylock-4        	50000000	        34.0 ns/op	       0 B/op	       0 allocs/op
BenchmarkChannMutex_Trylock-4         	20000000	        83.8 ns/op	       0 B/op	       0 allocs/op
BenchmarkSpinLock_Trylock-4           	50000000	        30.9 ns/op	       0 B/op	       0 allocs/op