Go实现可重入锁的两种办法

972 阅读3分钟

开启掘金成长之旅!这是我参与「掘金日新计划 · 2 月更文挑战」的第 1 天,点击查看活动详情

关于可重入锁

用过Java的朋友们都知道,ReentrantLock是一个可重入锁。什么是可重入?

当一个线程获取锁时,如果没有其它线程拥有这个锁,那么,这个线程就成功获取到这个锁。之后,如果其它线程再请求这个锁,就会处于阻塞等待的状态。但是,如果拥有这把锁的线程再请求这把锁的话,不会阻塞,而是成功返回,所以叫可重入锁(有时候也叫做递归锁)。只要你拥有这把锁,你可以可着劲儿地调用,比如通过递归实现一些算法,调用者不会阻塞或者死锁。

go语言的mutex只记录了加锁状态,没有记录锁的所有者,所以不支持可重入,自己加的锁别人也可以打开。下面介绍两种办法实现go的可重入锁。

方法一:将锁与gorouting ID绑定

对mutex进行修改,利用gorouting ID把锁与锁绑定,此gorouting可以重入,并防止其他gorouting开锁。

  • 获取gorouting ID

    • 可以利用runtime.Stack解析字符串获得
    func GoID() int {
        var buf [64]byte
        n := runtime.Stack(buf[:], false)
        // 得到id字符串
        idField := strings.Fields(strings.TrimPrefix(string(buf[:n]), "goroutine "))[0]
        id, err := strconv.Atoi(idField)
        if err != nil {
            panic(fmt.Sprintf("cannot get goroutine id: %v", err))
        }
        return id
    }
    
    • 也可以利用第三方库获取:petermattis/goid
    // RecursiveMutex 包装一个Mutex,实现可重入
    type RecursiveMutex struct {
        sync.Mutex
        owner     int64 // 当前持有锁的goroutine id
        recursion int32 // 这个goroutine 重入的次数
    }
    
    func (m *RecursiveMutex) Lock() {
        gid := goid.Get()
        // 如果当前持有锁的goroutine就是这次调用的goroutine,说明是重入
        if atomic.LoadInt64(&m.owner) == gid {
            m.recursion++
            return
        }
        m.Mutex.Lock()
        // 获得锁的goroutine第一次调用,记录下它的goroutine id,调用次数加1
        atomic.StoreInt64(&m.owner, gid)
        m.recursion = 1
    }
    
    func (m *RecursiveMutex) Unlock() {
        gid := goid.Get()
        // 非持有锁的goroutine尝试释放锁,错误的使用
        if atomic.LoadInt64(&m.owner) != gid {
            panic(fmt.Sprintf("wrong the owner(%d): %d!", m.owner, gid))
        }
        // 调用次数减1
        m.recursion--
        if m.recursion != 0 { // 如果这个goroutine还没有完全释放,则直接返回
            return
        }
        // 此goroutine最后一次调用,需要释放锁
        atomic.StoreInt64(&m.owner, -1)
        m.Mutex.Unlock()
    }
    

方法二:利用自定义token绑定锁

调用者自己提供一个 token,获取锁的时候把这个 token 传入,释放锁的时候也需要把这个 token 传入。通过用户传入的 token 替换方案一中 goroutine id,其它逻辑和方案一一致。

// Token方式的递归锁
type TokenRecursiveMutex struct {
    sync.Mutex
    token     int64
    recursion int32
}

// 请求锁,需要传入token
func (m *TokenRecursiveMutex) Lock(token int64) {
    if atomic.LoadInt64(&m.token) == token { //如果传入的token和持有锁的token一致,说明是递归调用
        m.recursion++
        return
    }
    m.Mutex.Lock() // 传入的token不一致,说明不是递归调用
    // 抢到锁之后记录这个token
    atomic.StoreInt64(&m.token, token)
    m.recursion = 1
}

// 释放锁
func (m *TokenRecursiveMutex) Unlock(token int64) {
    if atomic.LoadInt64(&m.token) != token { // 释放其它token持有的锁
        panic(fmt.Sprintf("wrong the owner(%d): %d!", m.token, token))
    }
    m.recursion-- // 当前持有这个锁的token释放锁
    if m.recursion != 0 { // 还没有回退到最初的递归调用
        return
    }
    atomic.StoreInt64(&m.token, 0) // 没有递归调用了,释放锁
    m.Mutex.Unlock()
}

这种方法就不满足Locker接口了,需要注意。