Golang 实现可重入锁（ReentrantLock）

Golang 实现可重入锁 (go-reentrantlock)

对 Go SDK sync 包熟悉的人都知道，其内部只提供了简单的、不可重入的 Mutex 和 RWMutex 结构。但是在一些实际应用中可能会出现线程需要多次进入临界区操作共享资源的情况，这种情况下使用 Mutex 或 RWMutex 都会导致死锁。Java 中 synchronized 修饰的代码块、JUC 包提供的 ReentrantLock 都是可重入锁。下面我们来使用 Golang 以最简单的方式去实现一个可重入锁。

（注：以下代码可以从 gitee 中克隆，https://gitee.com/elzatahmed/go-reentrant-lock）

1. 获取协程Id

GoLang 并没有像 Java 那样直接获取当前线程对象的方法，我们需要使用 Go Runtime 提供的 Stack 接口获取当前线程的栈信息，并从中抽取协程Id：

func getCurrentGoRoutineId() int64 {
	var buf [64]byte
  // 获取栈信息
	n := runtime.Stack(buf[:], false)
  // 抽取id
	idField := strings.Fields(strings.TrimPrefix(string(buf[:n]), "goroutine"))[0]
  // 转为64位整数
	id, _ := strconv.Atoi(idField)
	return int64(id)
}

2. 定义 locker 结构体并实现 sync.Locker 接口

sync.Locker 接口是 Go SDK 提供的对有获取锁 (Lock) 和释放锁 (Unlock) 方法的结构体的顶层接口。

// A Locker represents an object that can be locked and unlocked.
type Locker interface {
	Lock()
	Unlock()
}

我们实现的可重入锁也需要实现该接口以更完美地融入 Go 的开发模式。

locker 结构体的大致形状如下代码所示：

type locker struct {}

func New() sync.Locker {
	return &locker{}
}

func (l *locker) Lock() {}

func (l *locker) Unlock() {}

3. 实现 Lock 方法

为实现 Lock 方法，我们需要三种变量：

• heldCount：锁被持有的次数
• heldBy：锁被持有的协程Id
• cond：释放条件，利用 sync.Cond 实现

// locker implements sync.Locker as reentrant locker,
// it can be reaquired multiple times by the same goroutine which is holding the locker
type locker struct {
	heldCount int64      // how many times' locker has been held by the same goroutine
	heldBy    int64      // which goroutine is holding the locker
	cond      *sync.Cond // implements goroutine blocking and waking with sync.Cond
}

Lock 方法我们可以分为三种情况：

• 若当前线程持有锁：原子性地执行 heldCount + 1
• 若没有线程持有锁：原子性地替换 heldBy 并执行 heldCount + 1
• 若其他线程持有锁：利用 cond 去等待唤醒

上述三种情况的判断我们都使用atomic.CompareAndSwap去执行。

// tryOwnerAcquire acquires the lock if and only if lock is currently held by same goroutine
func (l *locker) tryOwnerAcquire(gid int64) bool {
  // CAS gid -> gid, 若成功则代表当前线程持有锁
	if atomic.CompareAndSwapInt64(&l.heldBy, gid, gid) {
		atomic.AddInt64(&l.heldCount, 1)
		return true
	}
	return false
}

// tryNoneAcquire uses CAS to try to swap l.heldBy and l.heldCount field
func (l *locker) tryNoneAcquire(gid int64) bool {
  // none = -1
  // CAS none -> gid, 若成功则代表原本没有线程持有锁，当前线程获取到了锁
	if atomic.CompareAndSwapInt64(&l.heldBy, none, gid) {
		atomic.AddInt64(&l.heldCount, 1)
		return true
	}
	return false
}

// acquireSlow waits for cond to signal
func (l *locker) acquireSlow(gid int64) {
	l.cond.L.Lock()
  // 若当前还是有其他线程持有锁
	for atomic.LoadInt64(&l.heldBy) != none {
    // 继续等待唤醒
		l.cond.Wait()
	}
  // 若当前没有线程持有锁
	l.acquireLocked(gid)
	l.cond.L.Unlock()
}

// acquireLocked does acquire by current goroutine while locked
func (l *locker) acquireLocked(gid int64) {
  // 由于sync.Cond的Signal方法每次只会唤醒一个线程，所以这里直接替换即可
	atomic.SwapInt64(&l.heldBy, gid)
	atomic.SwapInt64(&l.heldCount, 1)
}

实现好上述三种情况后，我们的 Lock 方法就如下代码所示：

// Lock locks the locker with reentrant mode
func (l *locker) Lock() {
	// get current goroutine id
	gid := getCurrentGoRoutineId()
	// if current goroutine is the lock owner
	if l.tryOwnerAcquire(gid) {
		return
	}
	// if lock is not held by any goroutine
	if l.tryNoneAcquire(gid) {
		return
	}
	// if current goroutine is not the lock owner
	l.acquireSlow(gid)
}

4. 实现 Unlock 方法

Unlock 方法没有 Lock 方法那么复杂，因为 Unlock 方法我们只允许当前持有锁的线程进入，若有其他线程进入则直接 panic 即可。持有锁的线程每 Unlock 一次我们就原子性地执行一次 heldCount - 1。如果 heldCount 变为0，就 cond 做一次 Signal 唤醒一个阻塞队列中的线程。

// Unlock unlocks the locker using reentrant mode
func (l *locker) Unlock() {
	// only the owner goroutine can enter this method
	heldBy := atomic.LoadInt64(&l.heldBy)
	if heldBy == none {
		panic("unlock of unlocked reentrantLocker")
	}
	if heldBy == getCurrentGoRoutineId() {
		l.releaseOnce()
		return
	}
	panic("unlock by different goroutine")
}

// releaseOnce release once on locker
func (l *locker) releaseOnce() {
	if l.heldCount == 0 {
		panic("unlocks more than locks")
	}
	l.heldCount--
	if l.heldCount == 0 {
		l.heldBy = -1
		// signal one and only one goroutine
		l.cond.Signal()
	}
	return
}