go 中 sync.Map 源码刨铣

func syncMapDemo() {

	var smp sync.Map

	// 数据写入
	smp.Store("name", "小红")
	smp.Store("age", 18)

	// 数据读取
	name, _ := smp.Load("name")
	fmt.Println(name)

	age, _ := smp.Load("age")
	fmt.Println(age)

	// 遍历
	smp.Range(func(key, value interface{}) bool {
		fmt.Println(key, value)
		return true
	})

	// 删除
	smp.Delete("age")
	age, ok := smp.Load("age")
	fmt.Println("删除后的查询")
	fmt.Println(age, ok)

	// 读取或写入,存在就读取，不存在就写入
	smp.LoadOrStore("age", 100)
	age, _ = smp.Load("age")
	fmt.Println("不存在")
	fmt.Println(age)

	smp.LoadOrStore("age", 99)
	age, _ = smp.Load("age")
	fmt.Println("存在")
	fmt.Println(age)
}

// sync/map.go
type Map struct {
	// 当写read map 或读写dirty map时 需要上锁
	mu Mutex

	// read map的 k v(entry) 是不变的，删除只是打标记，插入新key会加锁写到dirty中
	// 因此对read map的读取无需加锁
	read atomic.Value // 保存readOnly结构体

	// dirty map 对dirty map的操作需要持有mu锁
	dirty map[interface{}]*entry

	// 当Load操作在read map中未找到，尝试从dirty中进行加载时(不管是否存在)，misses+1
	// 当misses达到diry map len时，dirty被提升为read 并且重新分配dirty
	misses int
}

// read map数据结构
type readOnly struct {
	m       map[interface{}]*entry
	// 为true时代表dirty map中含有m中没有的元素
	amended bool
}

type entry struct {
	// 指向实际的interface{}
	// p有三种状态:
	// p == nil: 键值已经被删除，此时，m.dirty==nil 或 m.dirty[k]指向该entry
	// p == expunged: 键值已经被删除， 此时, m.dirty!=nil 且 m.dirty不存在该键值
	// 其它情况代表实际interface{}地址 如果m.dirty!=nil 则 m.read[key] 和 m.dirty[key] 指向同一个entry
	// 当删除key时，并不实际删除，先CAS entry.p为nil 等到每次dirty map创建时(dirty提升后的第一次新建Key)，会将entry.p由nil CAS为expunged
	p unsafe.Pointer // *interface{}
}

func (m *Map) Load(key interface{}) (value interface{}, ok bool) {
	// 首先在通过atomic的原子操作读取内容
	read, _ := m.read.Load().(readOnly)
	e, ok := read.m[key]
	// 如果没在 read 中找到，并且 amended 为 true，即 dirty 中存在 read 中没有的 key
	if !ok && read.amended {
		// read调用了atomic的原子性，所以不用加锁，但是dirty map[interface{}]*entry就需要了，用的互斥锁
		m.mu.Lock()
		// double check，避免在加锁的时候dirty map提升为read map
		read, _ = m.read.Load().(readOnly)
		e, ok = read.m[key]
		// 还是没有找到
		if !ok && read.amended {
			// 从 dirty 中找
			e, ok = m.dirty[key]
			// 不管dirty中有没有找到 都增加misses计数 该函数可能将dirty map提升为readmap
			m.missLocked()
		}
		m.mu.Unlock()
	}
	if !ok {
		return nil, false
	}
	return e.load()
}

// 从entry中atomic load实际interface{}
func (e *entry) load() (value interface{}, ok bool) {
	p := atomic.LoadPointer(&e.p)
	if p == nil || p == expunged {
		return nil, false
	}
	return *(*interface{})(p), true
}

// 增加misses计数，并在必要的时候提升dirty map
// 如果 misses 值小于 m.dirty 的长度，就直接返回。否则，将 m.dirty 晋升为 read，并清空 dirty，清空 misses 计数值。这样，之前
// 一段时间新加入的 key 都会进入到 read 中，从而能够提升 read 的命中率。
func (m *Map) missLocked() {
	m.misses++
	if m.misses < len(m.dirty) {
		return
	}
	// 提升过程很简单，直接将m.dirty赋给m.read.m
	// 提升完成之后 amended == false m.dirty == nil
	// m.dirty并不立即创建被拷贝元素，而是延迟创建
	m.read.Store(readOnly{m: m.dirty})
	m.dirty = nil
	m.misses = 0
}

// Store sets the value for a key.
func (m *Map) Store(key, value interface{}) {
	// 如果read map中存在该key  则尝试直接更改(由于修改的是entry内部的pointer，因此dirty map也可见)
	read, _ := m.read.Load().(readOnly)
	if e, ok := read.m[key]; ok && e.tryStore(&value) {
		return
	}

	m.mu.Lock()
	read, _ = m.read.Load().(readOnly)
	if e, ok := read.m[key]; ok {
		if e.unexpungeLocked() {
			// 如果 read map 中存在该 key，但 p == expunged，则说明 m.dirty != nil 并且 m.dirty 中不存在该 key 值 此时:
			//    a. 将 p 的状态由 expunged  更改为 nil
			//    b. dirty map 插入 key
			m.dirty[key] = e
		}
		// 更新 entry.p = value (read map 和 dirty map 指向同一个 entry)
		e.storeLocked(&value)
	} else if e, ok := m.dirty[key]; ok {
		// 如果 read map 中不存在该 key，但 dirty map 中存在该 key，直接写入更新 entry(read map 中仍然没有这个 key)
		e.storeLocked(&value)
	} else {
		// 如果read map和dirty map中都不存在该key，则:
		//    a. 如果dirty map为空，则需要创建dirty map，并从read map中拷贝未删除的元素
		//    b. 更新amended字段，标识dirty map中存在read map中没有的key
		//    c. 将k v写入dirty map中，read.m不变
		if !read.amended {

			m.dirtyLocked()
			m.read.Store(readOnly{m: read.m, amended: true})
		}
		m.dirty[key] = newEntry(value)
	}
	m.mu.Unlock()
}

// 尝试直接更新entry 如果p == expunged 返回false
func (e *entry) tryStore(i *interface{}) bool {
	p := atomic.LoadPointer(&e.p)
	if p == expunged {
		return false
	}
	for {
		if atomic.CompareAndSwapPointer(&e.p, p, unsafe.Pointer(i)) {
			return true
		}
		p = atomic.LoadPointer(&e.p)
		if p == expunged {
			return false
		}
	}
}

func (e *entry) unexpungeLocked() (wasExpunged bool) {
	return atomic.CompareAndSwapPointer(&e.p, expunged, nil)
}

// 如果 dirty map为nil，则从read map中拷贝元素到dirty map
func (m *Map) dirtyLocked() {
	if m.dirty != nil {
		return
	}

	read, _ := m.read.Load().(readOnly)
	m.dirty = make(map[interface{}]*entry, len(read.m))
	for k, e := range read.m {
		// a. 将所有为 nil的 p 置为 expunged
		// b. 只拷贝不为expunged 的 p
		if !e.tryExpungeLocked() {
			m.dirty[k] = e
		}
	}
}

func (e *entry) tryExpungeLocked() (isExpunged bool) {
	p := atomic.LoadPointer(&e.p)
	for p == nil {
		if atomic.CompareAndSwapPointer(&e.p, nil, expunged) {
			return true
		}
		p = atomic.LoadPointer(&e.p)
	}
	return p == expunged
}

// Delete deletes the value for a key.
func (m *Map) Delete(key interface{}) {
	// 从read map中寻找
	read, _ := m.read.Load().(readOnly)
	e, ok := read.m[key]
	// 没找到
	if !ok && read.amended { // read.amended为true代表dirty map中含有m中没有的元素
		m.mu.Lock()
		// double check
		read, _ = m.read.Load().(readOnly)
		e, ok = read.m[key]
		// 第二次仍然没找到，但dirty map中存在，则直接从dirty map删除
		if !ok && read.amended {
			delete(m.dirty, key)
		}
		m.mu.Unlock()
	}
	// 如果read存在，将entry.p 置为 nil
	if ok {
		e.delete()
	}
}

func (e *entry) delete() (hadValue bool) {
	for {
		p := atomic.LoadPointer(&e.p)
		if p == nil || p == expunged {
			return false
		}
		if atomic.CompareAndSwapPointer(&e.p, p, nil) {
			return true
		}
	}
}

func (m *Map) LoadOrStore(key, value interface{}) (actual interface{}, loaded bool) {
	// 首先还是先去read map中查询
	read, _ := m.read.Load().(readOnly)
	if e, ok := read.m[key]; ok {
		actual, loaded, ok := e.tryLoadOrStore(value)
		if ok {
			return actual, loaded
		}
	}

	m.mu.Lock()
	// double check
	read, _ = m.read.Load().(readOnly)
	if e, ok := read.m[key]; ok {
		if e.unexpungeLocked() {
			// 如果 read map 中存在该 key，但 p == expunged，则说明 m.dirty != nil 并且 m.dirty 中不存在该 key 值 此时:
			//    a. 将 p 的状态由 expunged  更改为 nil
			//    b. dirty map 插入 key
			m.dirty[key] = e
		}
		actual, loaded, _ = e.tryLoadOrStore(value)
	} else if e, ok := m.dirty[key]; ok {
		// 如果 read map 中不存在该 key，但 dirty map 中存在该 key
		actual, loaded, _ = e.tryLoadOrStore(value)
		// 不管dirty中有没有找到 都增加misses计数 该函数可能将dirty map提升为readmap
		m.missLocked()
	} else {
		if !read.amended {
			// 如果read map和dirty map中都不存在该key，则:
			//    a. 如果dirty map为空，则需要创建dirty map，并从read map中拷贝未删除的元素
			//    b. 更新amended字段，标识dirty map中存在read map中没有的key
			//    c. 将k v写入dirty map中，read.m不变
			m.dirtyLocked()
			m.read.Store(readOnly{m: read.m, amended: true})
		}
		m.dirty[key] = newEntry(value)
		actual, loaded = value, false
	}
	m.mu.Unlock()

	return actual, loaded
}

// 如果entry is expunged则不处理，返回false
func (e *entry) tryLoadOrStore(i interface{}) (actual interface{}, loaded, ok bool) {
	p := atomic.LoadPointer(&e.p)
	if p == expunged {
		return nil, false, false
	}
	if p != nil {
		return *(*interface{})(p), true, true
	}

	// Copy the interface after the first load to make this method more amenable
	// to escape analysis: if we hit the "load" path or the entry is expunged, we
	// shouldn't bother heap-allocating.
	ic := i
	for {
		if atomic.CompareAndSwapPointer(&e.p, nil, unsafe.Pointer(&ic)) {
			return i, false, true
		}
		p = atomic.LoadPointer(&e.p)
		if p == expunged {
			return nil, false, false
		}
		if p != nil {
			return *(*interface{})(p), true, true
		}
	}
}