大战runtime.Map

礼物说 2023-04-22 ⏳3.4分钟(1.3千字)

最近业务上碰到了一个有趣的问题：有一个3.1G的模型文件，内容为2.2亿的uint64->float64,存储在Go runtime.map中占用了近10个G，内容膨胀了3倍。基于这个问题我们深入了解下runtime.map源码，并结合实际的业务场景，聊一聊如何去优化其内存使用率。

runtime.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
}

申请内存

runtime.makemap¹为入口，找到合适的B²,根据B去runtime.makeBucketArray³,（1<<B(正常桶)+1<<(B-4)(溢出桶)）* bucketSize，然后分配空间。

基于上述的内容，我们可以发现map的内存大小和count并不是线性的关系。当到达了某个零界点⁴后会直接翻倍。下面我们验证下猜想以及我们当前2.2亿的数据量处于什么位置。

// find 根据runtime.map的分配规则，找到零界点
// 2.2亿前后的零界点为:218103808,436207616
func find() {
  for i := 0; i < 30; i++ {
    fmt.Println((1 << i) * 13 / 2)
  }
}

// alloc 测试map空间分配
func alloc() {
  debug.SetGCPercent(-1)
  var m runtime.MemStats
  runtime.ReadMemStats(&m)
  startAlloc := m.Alloc

  _ = make(map[int]int, 218103808)
  runtime.ReadMemStats(&m)
  fmt.Println(m.Alloc - startAlloc) // 5133828120
  runtime.ReadMemStats(&m)
  startAlloc = m.Alloc

  _ = make(map[int]int, 218103809)
  runtime.ReadMemStats(&m)
  fmt.Println(m.Alloc - startAlloc) // 10267656216
  runtime.ReadMemStats(&m)
  startAlloc = m.Alloc

  _ = make(map[int]int, 436207616)
  runtime.ReadMemStats(&m)
  fmt.Println(m.Alloc - startAlloc) // 10267656216
}

可以发现我们2.2亿的数据量刚过零界点，内存使用直接从5G翻倍到了10G，这也就能解释为什么我们为什么最近才出的问题。基于这个点我们可以得出方案一。

扩容方式

触发扩容

if !h.growing() && (overLoadFactor(h.count+1, h.B) || tooManyOverflowBuckets(h.noverflow, h.B)) {
  hashGrow(t, h)
  goto again // Growing the table invalidates everything, so try again
 }

装载因子过高会触发倍量扩容

溢出桶过多会触发等量扩容

扩容过程

func (h *hmap) growing() bool {
  return h.oldbuckets != nil
}
...
if h.growing() {
  growWork(t, h, bucket)
}
...
func growWork(t *maptype, h *hmap, bucket uintptr) {
  // make sure we evacuate the oldbucket corresponding
  // to the bucket we're about to use
  evacuate(t, h, bucket&h.oldbucketmask())

  // evacuate one more oldbucket to make progress on growing
  if h.growing() {
      evacuate(t, h, h.nevacuate)
  }
}

mapdelete⁵ & mapassign⁶ 的时候会判断是否在扩容中，从而去进行growWork⁷(将当前命中的bucket迁移到新的bucket中，并额外迁移一个其他的bucket),全部迁移之后会将h.oldbuckets设置为nil. 基于上述的内容，我们可以发现如果map内部扩容未完成那么新老bucket会同时存在，占用的内存就是1+2.

而我们的业务就是当数据加载完之后就只读了。所以map内部可能同时存在新老bucket。下面有个扩容的例子，我们可以尝试循环设置去帮助map稳定。方案二。

// grow 测试map的扩容
// 直接申请1665(只有新bucket),使用内存81944
// 申请1664再触发扩容(新老bucket都存在),使用内存122804
// 遍历设置一遍map,触发渐进式扩容(回收老bucket),减少内存-39552
func grow() {
  debug.SetGCPercent(-1)
  var m runtime.MemStats
  runtime.ReadMemStats(&m)
  startAlloc := m.Alloc

  mm := make(map[int]int, 1664)
  for i := 0; i < 1665; i++ {
      mm[i] = i
  }
  runtime.ReadMemStats(&m)
  fmt.Println(m.Alloc - startAlloc) // 122904
  runtime.ReadMemStats(&m)
  startAlloc = m.Alloc

  for k, v := range mm {
      mm[k] = v
  }
  runtime.GC()
  runtime.ReadMemStats(&m)
  fmt.Println(int(m.Alloc) - int(startAlloc)) // -39552
  runtime.ReadMemStats(&m)
  startAlloc = m.Alloc

  _ = make(map[int]int, 1665)
  runtime.ReadMemStats(&m)
  fmt.Println(m.Alloc - startAlloc) // 81944
  mm[1666] = 1666
}

思路

基于上述的内容，我们产生了2个思路

方案一

map分shard，降低扩容的颗粒度。当数量在218103808时分配了5G内存，当数量在218103808+1时就分配了10G内存，因为你的hmap.B涨了1，我们业务个2.2亿数据正好就在这个附近，所以导致了这个问题。我们可以分成2^x个map,这样当+1出现时只有一个map的B会上涨。

方案二

由于我们的map在生成完之后只有读操作，所以可能是因为没有渐进式扩容的触发点，导致oldbucket一直存在，占用了内存。GO并没有提供推动map扩容的工具，所以我们只能尝试循环一遍map，设置m[k]=v来推动扩容。

以上方案均能解决一部分的问题，结合使用可以将10G的内存优化到6G，但还是存在内存膨胀的问题，runtime.map是一个比较通用的场景，而我们更需要一个针对我们业务的存储结构，我们可以基于业务的数据量以及hash的冲突分布去实现一个查询O(1),插入尽量快的一个方案。针对紧凑型的内存结构首先想到了slice，make(map[int]int, 220000000)占用了9.5G，make([]int, 220000000) * 2(kv都需要)占用了3.2G.感觉方案可行。

实践

数据存储

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

寻址方式

借鉴runtime.Map的hash方式，每一层数量为2^B。第一层的index为 key & (1<<B - 1) 下一层的index为 index &= 1<<B >>1 - 1

最终的寻址过程为: fastModN获取hash值当前层是否有数据无->存入，有-> hash移位定位下一层 5层slice都冲突 -> 存入map

代码实现

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)
}

最终效果

runtime.Map vs own.Map 在亿+数据量下的对比

runtime.Map内存使用:9.5G

own.Map内存使用:3.6G

runtime.Map.Set: 43.42 ns/op

own.Map.Set: 38.26 ns/op

runtime.Map.Get: 8.952 ns/op

own.Map.Get:3.981 ns/op

Set的性能测试：

Size(byte)	Runtime.Map(ns/op)	Own.Map(ns/op)
16	7.253	13.97
128	7.874	18.41
1024	12.61	22.69
8192	20.22	29.27
131072	22.67	33.59
220000000	43.42	38.26

Get的性能测试：

Size(byte)	Runtime.Map(ns/op)	Own.Map(ns/op)
16	5.337	2.013
128	6.061	2.502
1024	7.307	2.551
8192	17.37	4.871
131072	20.78	9.125
100000000	8.952	3.981

参考链接

理解 Golang 哈希表 Map 的原理