Go TryLock实现

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

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

实现自旋锁

上面一节给了我们启发,利用 uint32CAS操作我们可以一个自定义的锁:

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

性能比较

首先看看上面三种方式和标准库中的MutexRWMutexLockUnlock的性能比较:

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方法的锁的性能:

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