go map 探索和内存优化

2023-06-11 ⏳2.3分钟(0.9千字)

本文介绍 golang runtime.map 相关和自定一个内存占用相对少的 map.

缘起

最近业务碰到了一个内存 OOM 的问题, 业务要加载一个模型, 容量大概为 2.2 亿条的 map, 其中 key 是 int64, value 是 float64.

源二进制文件(紧凑的int64 float64 int64 float64 ….)大概是 3.1G, 但是使用 runtime.map 加载到内存后发现, 占用 10G 左右, 而实例的内存是 16G, 这个文件每天更新, 由于每次更新的 value 是变化的为了保持版本一致无法增量更新, 所以在加载模型的时候会导致 OOM.

为什么 runtime.Map 占用那么大?

我们知道 map 数据结构不是一个定容量的数据结构, 在一定条件下会触发扩容, 扩容一般会翻倍。如果我们加载的键值对刚刚超过阈值, 而 map 扩容了一倍这个时候是确实会浪费内存, 去看 runtime.map 分配内存 的源码:

base := bucketShift(b)
nbuckets := base
// For small b, overflow buckets are unlikely.
// Avoid the overhead of the calculation.
if b >= 4 {
        // Add on the estimated number of overflow buckets
        // required to insert the median number of elements
        // used with this value of b.
        nbuckets += bucketShift(b - 4)
        sz := t.bucket.size * nbuckets
        up := roundupsize(sz)
	if up != sz {
        nbuckets = up / t.bucket.size
	}
}

其中 bucketShift 在计算 2^N , 当 b (b 代表找到的能存放目标简直对的 2 的幂) > 4 的时候, 会创建一个 2^(b-4) 的溢出桶, 用来处理冲突的键值对。

我们再来看 map 的具体数据结构:

// map struct
type hmap struct {
  count     int     // 元素数量
  flags     uint8   // 二进制位,用来判断
                    // 是否isWriting/是否是sameSizeGrow
  B         uint8   // 当前持有2^B个bucket 
                    // 可以容纳loadFactor[^装载因子] * 2^B个元素
  noverflow uint16  // 溢出桶的数量
  hash0     uint32  // hash seed

  buckets    unsafe.Pointer // 当前bucket
  oldbuckets unsafe.Pointer // 老bucket
  nevacuate  uintptr // TODO

  extra *mapextra
}

// 溢出桶信息
type mapextra struct {
  overflow    *[]*bmap
  oldoverflow *[]*bmap

  nextOverflow *bmap
}

// bucket struct
type bmap struct {
  tophash [bucketCnt]uint8
}

其中 buckets 是存放当前的 key-value 值, 而 oldbuckets 保留了老的 key-value 值, 所以当触发扩缩的时候, 会出现占用两倍内存的情况.

发生扩缩的原因会有两种

那什么时候会把 oldbuckets 清空呢？ 扩容的时候不会一次性把 oldbuckets 直接迁移到 buckets, 这是一个随着访问慢慢迁移的过程. 那么

未迁移完成前那么会一直多占用内存。

思考

  1. 我们加载的模型完成后其实不会再更改, 只会整体替换成新模型

  2. 业务上发现 value 的值在 -1 ～ 1 之间徘徊, 不需要 float64 的精度

  3. 检查了下冲突率, 大部分冲突的情况是出现了两次, 两次以上的情况很少.

考虑 1 的情况, runtime.map 的扩缩带来了浪费, 完全不需要 oldbuckets. 而且额外创建的溢出桶也不是必要的, 可以不要或者不要那么多.

考虑 2 的情况, float64 可以换成 float32 (ps: 立省25% 内存…

假设不自实现 map , 可以考虑想办法让扩容后的 map 进行快速迁移, 考虑到量级很大, 迁移行为可能影响线上业务 - pass

那么考虑自实现 map .

方案

5层slice + map: 引入map是为了解决在slice中多次冲突的数据，期望在保证slice层数的同时落入map中的数据尽可能的少。 slice为倒金字塔形，下一层的长度为当前的一半，期望每层尽可能的紧凑存储数据，落到下一层的冲突数据越来越少。测试了3-7层的数据分布,3层的map元素数量在3.5kw，5层的map元素在1kw，7层的map在0.7kw。最终选取了5层保证一定的查询效率。

package xmap

// 层数 过大查询效率低 过小冲突率高
const layer = 5

type Map struct {
	vs    [][]float32
	ks    [][]uint
	m     map[uint]float32
	hint  uint // 最高层个数 1<<b
	b     uint // 最高层数
	l     uint // 最低层数
	count int
}

func New(hint uint) (m *Map) {
	b := findB(hint)

	m = &Map{
		m:    make(map[uint]float32, int(float32(hint)*0.04)),
		hint: uint(1 << b),
		b:    b,
	}

	if b < layer {
		m.l = 0
	} else {
		m.l = b - layer
	}
	for i := b; i > m.l; i-- {
		m.vs = append(m.vs, make([]float32, 1<<i))
		m.ks = append(m.ks, make([]uint, 1<<i))
	}
	return
}

func findB(hint uint) uint {
	i := uint(1)
	for {
		if 1<<i > hint {
			return i - 1
		}
		i++
	}
}

func (m *Map) Get(k uint) float32 {
	mr := fastModN(k, m.b)
	hint := m.hint
	for i := 0; i < int(m.b-m.l); i++ {
		if m.ks[i][mr] == k {
			return m.vs[i][mr]
		}
		hint = hint >> 1
		mr &= hint - 1
	}

	return m.m[k]
}

func (m *Map) Set(k uint, v float32) {
	m.count++
	mr := fastModN(k, m.b)
	hint := m.hint
	for i := 0; i < int(m.b-m.l); i++ {
		if m.ks[i][mr] == 0 {
			m.ks[i][mr] = k
			m.vs[i][mr] = v
			return
		}
		hint = hint >> 1
		mr &= hint - 1
	}
	m.m[k] = v
}

func (m *Map) Len() int {
	if m == nil {
		return 0
	}
	return m.count
}

func fastModN(x uint, b uint) uint {
	return x & (1<<b - 1)
}