Go组件库总结之无等待锁

本篇文章我们用Go封装一个无等待锁库。文章参考自：https://github.com/brewlin/net-protocol

1.锁的封装

type Mutex struct {
    v  int32
    ch chan struct{}
}

2.锁的初始化

func (m *Mutex) Init() {
    m.v = 1
    m.ch = make(chan struct{}, 1)
}

3.等待锁

如果没有拿到锁值，将会阻塞在信道上。

func (m *Mutex) Lock() {
    if atomic.AddInt32(&m.v, -1) == 0 {
        return
    }

    for {
        if v := atomic.LoadInt32(&m.v); v >= 0 && atomic.SwapInt32(&m.v, -1) == 1 {
            return
        }

        <-m.ch
    }
}

4.无等待锁

如果没有拿到锁值，会直接返回。

func (m *Mutex) TryLock() bool {
    v := atomic.LoadInt32(&m.v)
    if v <= 0 {
        return false
    }
    return atomic.CompareAndSwapInt32(&m.v, 1, 0)
}

5.释放锁

释放锁时，向信道发送值以唤醒阻塞的协程。

func (m *Mutex) Unlock() {
    if atomic.SwapInt32(&m.v, 1) == 0 {
        return
    }

    select {
    case m.ch <- struct{}{}:
    default:
    }
}

6.使用示例

通过对上锁时机的控制，我们让GetCount()直接返回了-1，实现了无等待锁。

var (
    count      int
    countGuard *Mutex
)

func GetCount() int {
    if countGuard.TryLock() {
        countGuard.Unlock()
        return count
    } else {
        return -1
    }
}

func SetCount(c int) {
    countGuard.Lock()
    time.Sleep(time.Second * 2)
    count = c
    countGuard.Unlock()
}

func main() {
    countGuard = new(Mutex)
    countGuard.Init()

    go SetCount(1)
    time.Sleep(time.Second)
    fmt.Println(GetCount())
}