如何在Golang中实现基于指定key的互斥锁（Lock/Unlock）功能？

核心实现代码 #

package main

import "sync"

// KeyMutex 基于字符串key的互斥锁实现
type KeyMutex struct {
	globalMu sync.Mutex
	keyLocks map[string]*keyLockEntry
}

type keyLockEntry struct {
	mu     sync.Mutex
	refCnt int // 引用计数，标记当前有多少个goroutine在等待/持有该锁
}

// NewKeyMutex 初始化KeyMutex实例
func NewKeyMutex() *KeyMutex {
	return &KeyMutex{
		keyLocks: make(map[string]*keyLockEntry),
	}
}

// Lock 对指定key加锁
func (km *KeyMutex) Lock(key string) {
	km.globalMu.Lock()
	entry, ok := km.keyLocks[key]
	if !ok {
		entry = &keyLockEntry{}
		km.keyLocks[key] = entry
	}
	entry.refCnt++
	km.globalMu.Unlock()

	entry.mu.Lock()
}

// Unlock 对指定key解锁
func (km *KeyMutex) Unlock(key string) {
	km.globalMu.Lock()
	entry := km.keyLocks[key]
	entry.mu.Unlock()
	entry.refCnt--
	// 引用计数为0时清理对应key，避免内存泄漏
	if entry.refCnt == 0 {
		delete(km.keyLocks, key)
	}
	km.globalMu.Unlock()
}

使用示例 #

首先初始化全局的KeyMutex实例，然后在你需要控制并发的函数中直接调用即可：

// 全局初始化锁实例
var keyLocker = NewKeyMutex()

func somefunction(key string, wg *sync.WaitGroup) {
	defer wg.Done()
	// 按key加锁
	keyLocker.Lock(key)
	// 用defer保证解锁，避免业务逻辑panic导致死锁
	defer keyLocker.Unlock(key)

	// 你的业务代码段
}

实现效果说明 #

• 相同key的goroutine同一时间仅能有一个执行业务代码，完全符合需求：比如key为A的goroutine1运行结束后，3、4才能依次执行；key为B的goroutine2运行结束后5才能执行
• 不同key的goroutine之间完全无互斥，可以并行运行
• 无内存泄漏问题：没有goroutine使用的key会自动从内部map中清理

注意事项 #

• 解锁和加锁必须使用相同的key，否则会出现逻辑异常甚至panic
• 必须保证每个Lock调用都对应一次Unlock调用，推荐用defer来处理解锁逻辑

内容的提问来源于stack exchange，提问作者killsteenkill