Go 标准库源码分析 - sync 包的 Map

Go 普通的map是线程不安全的，在并发读写时可能会出现脏数据；sync包中的Map是线程安全的，并发读写时无需额外加锁，适用于多读少写的情况。

一、数据结构

type Map struct {
	mu Mutex  // 互斥锁，用于锁定dirty map

	read atomic.Value  // 优先读map，支持原子操作

	dirty map[interface{}]*entry  // dirty是当前最新的map，支持读写

	misses int // 记录read读不到数据而加锁读read map及dirty map的次数
}

// 通过原子操作存储在read中的结构
type readOnly struct {
	m       map[interface{}]*entry
	amended bool  // 如果数据在dirty而在read中读取不到，则该值为true
}

// map中某个键存储的值
type entry struct {
    // p == nil：该键值对已被删除
    // p == expunged：也表示被删除，但该键只存在read中而不在dirty中，发生于read复制于dirty中，会先将值为nil的修改为expunged，然后不将其复制到dirty中
    // p为其他表示真正存储的数据
	p unsafe.Pointer 
}

 二、使用方法

1. Load

func (m *Map) Load(key interface{}) (value interface{}, ok bool) {
    // 先尝试在read map中读取
	read, _ := m.read.Load().(readOnly)
	e, ok := read.m[key]
    // 如果读取不到而且dirty map中存在read map没有的数据，则上锁进一步读取
	if !ok && read.amended {
		m.mu.Lock()
        // 再次读取read map，避免在加锁过程中dirty map升级为read map
		read, _ = m.read.Load().(readOnly)
		e, ok = read.m[key]
        // 如果还是读取不到则读取dirty map中的数据
		if !ok && read.amended {
			e, ok = m.dirty[key]
			// 不管元素存不存在都要记录miss数
			m.missLocked()
		}
		m.mu.Unlock()
	}
    // 元素不在则直接返回
	if !ok {
		return nil, false
	}
	return e.load()
}

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
}

2. Store

func (m *Map) Store(key, value interface{}) {
	read, _ := m.read.Load().(readOnly)
    // 如果read map中有该键则尝试写入，成功则直接返回
	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 {
        // 如果read map中有该键，则修改对应的值
		if e.unexpungeLocked() {
			// 如果元素被标记为expunged则表明dirty map中没有该元素
			m.dirty[key] = e
		}
		e.storeLocked(&value)
	} else if e, ok := m.dirty[key]; ok {
        // 如果dirty map中有该键，则修改对应的值
		e.storeLocked(&value)
	} else {
        // read map和dirty map均无该键时，在dirty map中添加键值对
		if !read.amended {
			// amender为false表明dirty map为空，先复制read map到dirty map，再修改amended表明dirty map有数据
			m.dirtyLocked()
			m.read.Store(readOnly{m: read.m, amended: true})
		}
		m.dirty[key] = newEntry(value)
	}
	m.mu.Unlock()
}

// 如果对应的值表示为expunged即已删除，则存储失败，否则将新的值存入
func (e *entry) tryStore(i *interface{}) bool {
	for {
		p := atomic.LoadPointer(&e.p)
		if p == expunged {
			return false
		}
		if atomic.CompareAndSwapPointer(&e.p, p, unsafe.Pointer(i)) {
			return true
		}
	}
}

// 存储新的值
func (e *entry) storeLocked(i *interface{}) {
	atomic.StorePointer(&e.p, unsafe.Pointer(i))
}

3. LoadOrStore

func (m *Map) LoadOrStore(key, value interface{}) (actual interface{}, loaded bool) {
	// 尝试对read map中进行load or store操作，成功则返回
	read, _ := m.read.Load().(readOnly)
	if e, ok := read.m[key]; ok {
		actual, loaded, ok := e.tryLoadOrStore(value)
		if ok {
			return actual, loaded
		}
	}

    // 以下同Store的思路
	m.mu.Lock()
	read, _ = m.read.Load().(readOnly)
	if e, ok := read.m[key]; ok {
		if e.unexpungeLocked() {
			m.dirty[key] = e
		}
		actual, loaded, _ = e.tryLoadOrStore(value)
	} else if e, ok := m.dirty[key]; ok {
		actual, loaded, _ = e.tryLoadOrStore(value)
		m.missLocked()
	} else {
		if !read.amended {
			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
}

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
	}

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

4. LoadAndDelete

func (m *Map) LoadAndDelete(key interface{}) (value interface{}, loaded bool) {
    // 第一次检测
	read, _ := m.read.Load().(readOnly)
	e, ok := read.m[key]
	if !ok && read.amended {
		m.mu.Lock()
        // 第二次检测
		read, _ = m.read.Load().(readOnly)
		e, ok = read.m[key]
		if !ok && read.amended {
            // 如果read map中读取不到且dirty map不为空，则至dirty map读取
			e, ok = m.dirty[key]
            // 不管key在不在dirty map中都会执行delete
			delete(m.dirty, key)
			m.missLocked()
		}
		m.mu.Unlock()
	}
    // 如果元素在read map中则标记为删除
	if ok {
		return e.delete()
	}
	return nil, false
}

func (e *entry) delete() (value interface{}, ok bool) {
	for {
		p := atomic.LoadPointer(&e.p)
		if p == nil || p == expunged {
			return nil, false
		}
        // 将元素标记为删除，并返回
		if atomic.CompareAndSwapPointer(&e.p, p, nil) {
			return *(*interface{})(p), true
		}
	}
}

5. Range

func (m *Map) Range(f func(key, value interface{}) bool) {
	read, _ := m.read.Load().(readOnly)
	if read.amended {		m.mu.Lock()
		read, _ = m.read.Load().(readOnly)
		if read.amended {
            // 如果dirty map不为空，即为最新的数据，则升级为read map
			read = readOnly{m: m.dirty}
			m.read.Store(read)
			m.dirty = nil
			m.misses = 0
		}
		m.mu.Unlock()
	}

	for k, e := range read.m {
		v, ok := e.load()
		if !ok {
			continue
		}
        // 元素存在则执行函数，返回false时终止
		if !f(k, v) {
			break
		}
	}
}

6. 其他

1）missLocked

// 记录miss数，在read map中读取不到数据而去读取dirty map时，miss+1
func (m *Map) missLocked() {
	m.misses++
	if m.misses < len(m.dirty) {
		return
	}
    // 如果miss数等于dirty map长度，则将dirty map升级为read map
	m.read.Store(readOnly{m: m.dirty})
	m.dirty = nil
	m.misses = 0
}

2）unexpungedLocked

// 判断元素是否被标记为已删除
func (e *entry) unexpungeLocked() (wasExpunged bool) {
	return atomic.CompareAndSwapPointer(&e.p, expunged, nil)
}

3）dirtyLocked

// 将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 {
        // 如果标记为nil或expunged则不复制到dirty
		if !e.tryExpungeLocked() {
			m.dirty[k] = e
		}
	}
}

4）tryExpungedLocked

// 判断是否标记为已删除
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
}