分布式KV存储中如何实现键级锁？

我正在构建一个分布式KV存储以深入学习分布式系统与并发技术。该KV存储为全事务型，采用内存事务日志，存储层也为纯内存实现以简化开发。API提供GET、INSERT、UPDATE、REMOVE接口，所有接口仅操作单个键而非键范围。

目前我通过全局锁控制并发，但效率极低——比如读取K1时若正在更新K2，二者无关却仍需等待。我知道部分数据库采用更细粒度的锁（比如MySQL的行级锁），想请教如何实现键级锁？

现有存储结构体如下：

type Storage struct {
  store map[string]int32
}

我考虑改成如下结构，但担心locks映射需要和store保持同步；另一种方案是合并两个映射，但仍会遇到REMOVE请求先于GET时，在持有锁的情况下删除映射条目的问题：

type Storage struct {
  store map[string]int32
  locks map[string]mutex.Lock
}

可行的键级锁实现方案 #

1. 带引用计数的独立键锁（精准粒度） #

核心思路是为每个键维护独立锁，通过引用计数管理锁的生命周期，彻底解决REMOVE操作导致的锁提前删除问题：

import (
    "sync"
    "sync/atomic"
)

type Storage struct {
    store   map[string]int32
    lockMap sync.Map // key: string, value: *lockedEntry
    mu      sync.Mutex // 保护lockMap的初始化竞争
}

type lockedEntry struct {
    mu        sync.Mutex
    refCount  int32 // 原子操作计数，避免并发修改冲突
}

// 获取指定键的锁，自动初始化不存在的锁
func (s *Storage) getLock(key string) *lockedEntry {
    entry, ok := s.lockMap.Load(key)
    if !ok {
        s.mu.Lock()
        defer s.mu.Unlock()
        // 二次检查，防止并发初始化重复创建锁
        entry, ok = s.lockMap.Load(key)
        if !ok {
            newEntry := &lockedEntry{}
            s.lockMap.Store(key, newEntry)
            entry = newEntry
        }
    }
    atomic.AddInt32(&entry.(*lockedEntry).refCount, 1)
    return entry.(*lockedEntry)
}

// 释放锁，当引用计数为0时自动清理锁条目
func (s *Storage) releaseLock(key string) {
    entry, ok := s.lockMap.Load(key)
    if !ok {
        return
    }
    le := entry.(*lockedEntry)
    // 引用计数减1后检查是否为0
    if atomic.AddInt32(&le.refCount, -1) == 0 {
        s.mu.Lock()
        defer s.mu.Unlock()
        // 再次确认计数，避免并发释放时误删
        if atomic.LoadInt32(&le.refCount) == 0 {
            s.lockMap.Delete(key)
        }
    }
}

使用示例（以GET接口为例）：

func (s *Storage) Get(key string) (int32, bool) {
    le := s.getLock(key)
    le.mu.Lock()
    defer func() {
        le.mu.Unlock()
        s.releaseLock(key)
    }()
    val, ok := s.store[key]
    return val, ok
}

2. 分片锁（折中高效方案） #

如果不需要极致精细的锁粒度，可以将键哈希到固定数量的分片上，每个分片对应一把锁。这种方案实现简单，能大幅降低锁竞争，同时避免每个键维护独立锁的开销：

import (
    "sync"
    "hash/fnv"
)

const shardCount = 32 // 可根据并发量调整分片数

type Storage struct {
    shards []*shard
}

type shard struct {
    store map[string]int32
    mu    sync.Mutex
}

func NewStorage() *Storage {
    s := &Storage{shards: make([]*shard, shardCount)}
    for i := range s.shards {
        s.shards[i] = &shard{store: make(map[string]int32)}
    }
    return s
}

// 根据键哈希到对应分片
func (s *Storage) getShard(key string) *shard {
    hash := fnv.New32a()
    hash.Write([]byte(key))
    return s.shards[hash.Sum32()%shardCount]
}

// GET接口示例
func (s *Storage) Get(key string) (int32, bool) {
    shard := s.getShard(key)
    shard.mu.Lock()
    defer shard.mu.Unlock()
    val, ok := shard.store[key]
    return val, ok
}

方案对比 #

• 带引用计数的键级锁：锁粒度最细，完全消除无关键的锁竞争，但需要额外维护引用计数，实现稍复杂。
• 分片锁：实现简单，性能足以应对多数场景，锁竞争概率远低于全局锁，适合快速迭代开发。

关键注意事项 #

• 所有接口（包括REMOVE）必须先获取对应键/分片的锁，再操作store，确保并发场景下的数据一致性。
• 带引用计数的方案中，长时间未访问的键的锁会被自动清理，避免内存泄漏。

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