Golang实现内存数据库

Concurrent Hash Map

KV 内存数据库的核心是并发安全的哈希表，常见的设计有几种：

• sync.map: golang 官方提供的并发哈希表, 适合读多写少的场景。但是在 m.dirty 刚被提升后会将 m.read 复制到新的 m.dirty 中，在数据量较大的情况下复制操作会阻塞所有协程，存在较大的隐患。

• juc.ConcurrentHashMap: java 的并发哈希表采用分段锁实现。在进行扩容时访问哈希表线程都将协助进行 rehash 操作，在 rehash 结束前所有的读写操作都会阻塞。因为缓存数据库中键值对数量巨大且对读写操作响应时间要求较高，使用juc的策略是不合适的。

• memcached hashtable: 在后台线程进行 rehash 操作时，主线程会判断要访问的哈希槽是否已被 rehash 从而决定操作 old_hashtable 还是操作 primary_hashtable。

memcached hashtable 的渐进式 Rehash 策略使主线程和rehash线程之间的 data race 限制在哈希槽内，最小化rehash操作对读写操作的影响，这是最理想的实现方式。但由于作者才疏学浅无法使用 golang 实现该策略故忍痛放弃（主要原因在于 golang 没有 volatile 关键字, 保证可见性的操作非常复杂），欢迎各位读者讨论。

本文采用 golang 社区广泛使用的分段锁策略。我们将 key 分散到固定数量的 shard 中避免 rehash 操作。shard 是有锁保护的 map, 当 shard 进行 rehash 时会阻塞shard内的读写，但不会对其他 shard 造成影响。

这种策略简单可靠易于实现，但由于需要两次 hash 性能略差。这个 dict 完整源码在github.com/hdt3213/godis/datastruct/dict/concurrent.go 可以独立使用（虽然也没有什么用。。。）。

定义数据结构:

type ConcurrentDict struct {
    table []*Shard
    count int32
}

type Shard struct {
    m     map[string]interface{}
    mutex sync.RWMutex
}

在构造时初始化 shard，这个操作相对比较耗时:

func computeCapacity(param int) (size int) {
	if param <= 16 {
		return 16
	}
	n := param - 1
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	if n < 0 {
		return math.MaxInt32
	} else {
		return int(n + 1)
	}
}

func MakeConcurrent(shardCount int) *ConcurrentDict {
    shardCount = computeCapacity(shardCount)
    table := make([]*Shard, shardCount)
    for i := 0; i < shardCount; i++ {
        table[i] = &Shard{
            m: make(map[string]interface{}),
        }
    }
    d := &ConcurrentDict{
        count: 0,
        table: table,
    }
    return d
}

哈希算法选择FNV算法:

const prime32 = uint32(16777619)

func fnv32(key string) uint32 {
    hash := uint32(2166136261)
    for i := 0; i < len(key); i++ {
        hash *= prime32
        hash ^= uint32(key[i])
    }
    return hash
}

定位shard, 当n为2的整数幂时 h % n == (n - 1) & h

func (dict *ConcurrentDict) spread(hashCode uint32) uint32 {
	if dict == nil {
		panic("dict is nil")
	}
	tableSize := uint32(len(dict.table))
	return (tableSize - 1) & uint32(hashCode)
}

func (dict *ConcurrentDict) getShard(index uint32) *Shard {
	if dict == nil {
		panic("dict is nil")
	}
	return dict.table[index]
}

Get 和 Put 方法实现:

func (dict *ConcurrentDict) Get(key string) (val interface{}, exists bool) {
	if dict == nil {
		panic("dict is nil")
	}
	hashCode := fnv32(key)
	index := dict.spread(hashCode)
	shard := dict.getShard(index)
	shard.mutex.RLock()
	defer shard.mutex.RUnlock()
	val, exists = shard.m[key]
	return
}

func (dict *ConcurrentDict) Len() int {
	if dict == nil {
		panic("dict is nil")
	}
	return int(atomic.LoadInt32(&dict.count))
}

// return the number of new inserted key-value
func (dict *ConcurrentDict) Put(key string, val interface{}) (result int) {
	if dict == nil {
		panic("dict is nil")
	}
	hashCode := fnv32(key)
	index := dict.spread(hashCode)
	shard := dict.getShard(index)
	shard.mutex.Lock()
	defer shard.mutex.Unlock()

	if _, ok := shard.m[key]; ok {
		shard.m[key] = val
		return 0
	} else {
		shard.m[key] = val
		dict.addCount()
		return 1
	}
}

LockMap

实现的ConcurrentMap 可以保证对单个 key 操作的并发安全性，但是仍然无法满足需求:

1. Incr 命令需要完成: 读取 -> 做加法 -> 写入 三步操作
2. MSETNX 命令当且仅当所有给定键都不存在时所有给定键设置值, 因此我们需要锁定所有给定的键直到完成所有键的检查和设置

因此我们需要实现 db.Locker 用于锁定一个或一组 key 并在我们需要的时候释放锁。

实现 db.Locker 最直接的想法是使用一个 map[string]*sync.RWMutex, 加锁过程分为两步: 初始化对应的锁 -> 加锁， 解锁过程也分为两步: 解锁 -> 释放对应的锁。那么存在一个无法解决的并发问题:

时间	协程A	协程B
1		locker["a"].Unlock()
2	locker["a"] = &sync.RWMutex{}
3		delete(locker["a"])
4	locker["a"].Lock()

由于 t3 时协程B释放了锁，t4 时协程A试图加锁会失败。

若我们在解锁时不释放锁就可以避免该异常的发生，但是每个曾经使用过的锁都无法释放从而造成严重的内存泄露。

我们注意到哈希表的长度远少于可能的键的数量，反过来说多个键可以共用一个哈希槽。若我们不为单个键加锁而是为它所在的哈希槽加锁，因为哈希槽的数量非常少即使不释放锁也不会占用太多内存。

作者根据这种思想实现了 LockMap 来解决并发控制问题。

type Locks struct {
    table []*sync.RWMutex
}

func Make(tableSize int) *Locks {
    table := make([]*sync.RWMutex, tableSize)
    for i := 0; i < tableSize; i++ {
        table[i] = &sync.RWMutex{}
    }
    return &Locks{
        table: table,
    }
}

func (locks *Locks)Lock(key string) {
    index := locks.spread(fnv32(key))
    mu := locks.table[index]
    mu.Lock()
}

func (locks *Locks)UnLock(key string) {
    index := locks.spread(fnv32(key))
    mu := locks.table[index]
    mu.Unlock()
}

哈希算法已经在 Dict 一节介绍过不再赘述。

在锁定多个key时需要注意，若协程A持有键a的锁试图获得键b的锁，此时协程B持有键b的锁试图获得键a的锁则会形成死锁。

解决方法是所有协程都按照相同顺序加锁，若两个协程都想获得键a和键b的锁，那么必须先获取键a的锁后获取键b的锁，这样就可以避免循环等待。

func (locks *Locks) toLockIndices(keys []string, reverse bool) []uint32 {
    indexMap := make(map[uint32]bool)
    for _, key := range keys {
        index := locks.spread(fnv32(key))
        indexMap[index] = true
    }
    indices := make([]uint32, 0, len(indexMap))
    for index := range indexMap {
        indices = append(indices, index)
    }
    sort.Slice(indices, func(i, j int) bool {
        if !reverse {
            return indices[i] < indices[j]
        } else {
            return indices[i] > indices[j]
        }
    })
    return indices
}

// 允许 writeKeys 和 readKeys 中存在重复的 key
func (locks *Locks) RWLocks(writeKeys []string, readKeys []string) {
	keys := append(writeKeys, readKeys...)
	indices := locks.toLockIndices(keys, false)
	writeIndices := locks.toLockIndices(writeKeys, false)
	writeIndexSet := make(map[uint32]struct{})
	for _, idx := range writeIndices {
		writeIndexSet[idx] = struct{}{}
	}
	for _, index := range indices {
		_, w := writeIndexSet[index]
		mu := locks.table[index]
		if w {
			mu.Lock()
		} else {
			mu.RLock()
		}
	}
}

func (locks *Locks) RWUnLocks(writeKeys []string, readKeys []string) {
	keys := append(writeKeys, readKeys...)
	indices := locks.toLockIndices(keys, true)
	writeIndices := locks.toLockIndices(writeKeys, true)
	writeIndexSet := make(map[uint32]struct{})
	for _, idx := range writeIndices {
		writeIndexSet[idx] = struct{}{}
	}
	for _, index := range indices {
		_, w := writeIndexSet[index]
		mu := locks.table[index]
		if w {
			mu.Unlock()
		} else {
			mu.RUnlock()
		}
	}
}

TTL

Time To Live (TTL) 功能可以为 key 设置失效时间。它的核心是存储 key -> expireTime 的 map 以及自动删除过期的 key 的时间轮。

完整代码在database/single_db.go

func genExpireTask(key string) string {
	return "expire:" + key
}

// Expire sets ttlCmd of key
func (db *DB) Expire(key string, expireTime time.Time) {
	db.stopWorld.Wait()
	db.ttlMap.Put(key, expireTime) // 记录过期时间
	taskKey := genExpireTask(key)
    // 通过时间轮设置定时任务，在 key 过期后自动删除
	timewheel.At(expireTime, taskKey, func() {
        // 采用 check-lock-check 防止在等待锁期间，其它协程修改了 ttl 设置
		db.Lock(key)
		defer db.UnLock(key)
		logger.Info("expire " + key)
		rawExpireTime, ok := db.ttlMap.Get(key)
		if !ok {
			return
		}
		expireTime, _ := rawExpireTime.(time.Time)
		expired := time.Now().After(expireTime)
		if expired {
			db.Remove(key)
		}
	})
}

// Persist cancel ttlCmd of key
func (db *DB) Persist(key string) {
	db.stopWorld.Wait()
	db.ttlMap.Remove(key)
	taskKey := genExpireTask(key)
	timewheel.Cancel(taskKey)
}

// IsExpired check whether a key is expired
func (db *DB) IsExpired(key string) bool {
    db.stopWorld.Wait()
    db.Lock(key)
	defer db.UnLock(key)
	rawExpireTime, ok := db.ttlMap.Get(key)
	if !ok {
		return false
	}
	expireTime, _ := rawExpireTime.(time.Time)
	expired := time.Now().After(expireTime)
	if expired {
		db.Remove(key)
	}
	return expired
}

内存数据库代码结构

最后为打算阅读源码的朋友们简单介绍一下 godis/database 包的代码结构：

MultiDB 实现了 database 接口，server.Handler 会持有一个 MultiDB 实例作为 Redis 存储引擎，并通过 Exec 函数将命令行传递给 MultiDB。

MultiDB 是支持 select 命令的多数据库引擎，它持有多个单数据库实例(godis.DB)以及支持发布订阅、AOF持久化所需的 pubsub.Hub 和 aof.Handler。

MultiDB.Exec 函数是 MultiDB 的总入口，它会自行处理鉴权、发布订阅、AOF相关命令以及FLUSHALL等系统命令，其它命令会调用被选中数据库的 DB.Exec 函数执行。

godis.DB.Exec 函数会自行处理 multi、exec 等事务控制命令，并调用 DB.execNormalCommand 函数处理常规命令。这里的常规命令是指仅读写有限个 key、可以在事务内执行且可以回滚的命令，比如 get、set、lpush 等命令，flushdb、keys 等命令不属于常规命令。

godis/database/router.go 中的 RegisterCommand函数负责注册常规命令。

实现一个常规命令需要提供3个函数:

• ExecFunc 是实际执行命令的函数，如: execHSet
• PrepareFunc 在 ExecFunc 前执行，负责分析命令行读写了哪些 key 便于进行加锁
• UndoFunc 仅在事务中被使用，负责准备 undo logs 以备事务执行过程中遇到错误需要回滚。