让 sync.Map 支持泛型

渐渐地， Go泛型越来越多应用的Go的标准库中了。一些标准库的类型，比如container/heap、container/list、container/ring、math都是有机会支持泛型的，但是考虑到Go向下兼容的情况，这些包可能不会直接修改，最可能就是新建一些并发的包，或者放在扩展包中。

本篇文章将讲一个相对复杂的例子，也就是对sync.Map的修改，让它支持泛型。

找出变换量，将其改变为type parameter

sync.Map是一种map，所以它遵循map的K-V的映射关系。键是一种可以比较的类型，以便进行是否相等的检查，值是任意类型。内建的map类型在变量定义的时候就执行了键和值的类型，也算是最初级的泛型，sync.Map的键和值是any类型的，但实际上键是comparable类型才可以。下面这个代码能编译，但是运行时会出错。

var m sync.Map

key := func() {}
m.Store(key, "value")

键的类型和值的类型是变化量，可以更改成类型参数，所以我们先对这两个变化量下手，修改成类型参数。

标准库中主要使用两个map对象做read和dirty数据中的存储，所以我们使用类型参数创建这两个map对象。
entry也改成泛型的支持。

支持泛型的sync.Map

type Map[K comparable, V any] struct { // 定义K,V 类型参数
	mu sync.Mutex

	read atomic.Pointer[readOnly[K, V]] // readonly 也支持泛型，因为它内部也是使用了一个map

	dirty map[K]*entry[V] // map改造成泛型

	misses int
}

type readOnly[K comparable, V any] struct {
	m       map[K]*entry[V] // 这个改成泛型
	amended bool
}

var expunged = unsafe.Pointer(new(any)) // 一个特殊的标记为删除的值

type entry[V any] struct {
	p atomic.Pointer[V] // 泛型类型，atomic.Pointer支持泛型
}

标准库中的sync.Map

type Map struct {
	mu Mutex

	read atomic.Pointer[readOnly]
	dirty map[any]*entry

	misses int
}

type readOnly struct {
	m       map[any]*entry
	amended bool 
}

var expunged = new(any)

type entry struct {
	p atomic.Pointer[any]
}

基本上我们通过少量的修改，把sync.Map的相关的类型都改成了支持泛型的类型的类型。

改造方法支持泛型

相关的类型改造成支持泛型后，编译就过不去了，因为相关的方法也需要修改。

泛型函数比如newEntry编译器或者编辑器中报错说是entry是泛型的，我们根据这个提示改成泛型的：

func newEntry[V any](i V) *entry[V] {
	e := &entry[V]{}
	e.p.Store(&i)
	return e
}

泛型方法中我们把Receiver的类型改成Map[K, V],最好是把(m *Map)全部替换一遍。返回值readonly也改成泛型类型的定义:

func (m *Map[K, V]) loadReadOnly() readOnly[K, V] {
	if p := m.read.Load(); p != nil {
		return *p
	}
	return readOnly[K, V]{}
}

接下来改造Load方法， 将其中的Key和Value类型从any类型改成K,V类型参数。
这里一个小机器就是源标准库中有可能返回nil值，改成返回V类型的零值，V类型的零值有三种方式创建出来，我们使用var zero V创建零值：

支持泛型的sync.Map

func (m *Map[K, V]) Load(key K) (value V, ok bool) {
	read := m.loadReadOnly()
	e, ok := read.m[key]
	if !ok && read.amended {
		m.mu.Lock()
		read = m.loadReadOnly()
		e, ok = read.m[key]
		if !ok && read.amended {
			e, ok = m.dirty[key]
			// Regardless of whether the entry was present, record a miss: this key
			// will take the slow path until the dirty map is promoted to the read
			// map.
			m.missLocked()
		}
		m.mu.Unlock()
	}
	if !ok {
		var zero V
		return zero, false
	}
	return e.load()
}

标准库中的sync.Map

func (m *Map) Load(key any) (value any, ok bool) {
	read := m.loadReadOnly()
	e, ok := read.m[key]
	if !ok && read.amended {
		m.mu.Lock()
		// Avoid reporting a spurious miss if m.dirty got promoted while we were
		// blocked on m.mu. (If further loads of the same key will not miss, it's
		// not worth copying the dirty map for this key.)
		read = m.loadReadOnly()
		e, ok = read.m[key]
		if !ok && read.amended {
			e, ok = m.dirty[key]
			// Regardless of whether the entry was present, record a miss: this key
			// will take the slow path until the dirty map is promoted to the read
			// map.
			m.missLocked()
		}
		m.mu.Unlock()
	}
	if !ok {
		return nil, false
	}
	return e.load()
}

接下来我再介绍一个棘手的改造情况。
load方法中会判断p == expunged, expunged是一个特殊的值，用来标记为被删除还没有被清理掉。但是这个时候p的类型是*V, p的类型是unsafe.Pointer(new(interface{})),类型不一样，指向的值也不一样。

改如何改造呢？

func (e *entry) load() (value any, ok bool) {
	p := e.p.Load()
	if p == nil || p == expunged {
		return nil, false
	}
	return *p, true
}

我们可以使用unsafe.Pointer(p) == expunged转为相同的类型，首先类型相同，另外我，们还是使用unsafe.Pointer(new(interface{}))来标记被删除的键:

func (e *entry[V]) load() (value V, ok bool) {
	p := e.p.Load()
	if p == nil || unsafe.Pointer(p) == expunged {
		var zero V
		return zero, false
	}
	return *p, true
}

在unexpungeLocked方法中，我们强制把expunged转换成*V,尽管它底层指向的是一个any,反正部门也不会更改expunged的值，这种强制转换也没有啥风险。

func (e *entry[V]) unexpungeLocked() (wasExpunged bool) {
	return e.p.CompareAndSwap((*V)(expunged), nil)
}

剩下的方法就用同样的方式进行改造，甚至可以通过批量替换的方式处理，最终我们把整个sync.Map改造成支持泛型了。

最终形成的代码如下：
generic sync map

把官方的测试代码复制过来测试，没有问题。

性能

最后大家比较关注的就是这一顿骚操作能不能带来性能的提升呢？功能的提升是显而易见的，我们输入的参数和返回值都是一个特定的类型，不需要做type assert了。

还是改造官方标准库对sync.Map的bench代码，我们增加一个这个泛型类型的支持，为了公平期间我们使用的key和value都是int类型。相关的代码可以看map_bench_test.go。

测试结果如下,可以看到绝大部分场景下我们的泛型Map要比标准库中的sync.Map要快，只有在BenchmarkMapCompareAndSwapNoExistingKey和BenchmarkMapCompareAndSwapValueNotEqual两个场景下比sync.Map慢，并且慢很多，待我找找原因，如果有分析结果会在另一篇文章中介绍。

BenchmarkMapLoadMostlyHits/*syncx.SyncMap[int,int]-8         	319121388	         3.725 ns/op
BenchmarkMapLoadMostlyHits/*syncx.Map[int,int]-8             	897178652	         2.491 ns/op
BenchmarkMapLoadMostlyMisses/*syncx.SyncMap[int,int]-8       	478179013	         2.426 ns/op
BenchmarkMapLoadMostlyMisses/*syncx.Map[int,int]-8           	866405142	         2.315 ns/op
BenchmarkMapLoadOrStoreBalanced/*syncx.SyncMap[int,int]-8    	 5727007	       223.4 ns/op
BenchmarkMapLoadOrStoreBalanced/*syncx.Map[int,int]-8        	 8712914	       171.7 ns/op
BenchmarkMapLoadOrStoreUnique/*syncx.SyncMap[int,int]-8      	 2843049	       405.7 ns/op
BenchmarkMapLoadOrStoreUnique/*syncx.Map[int,int]-8          	 4316347	       321.8 ns/op
BenchmarkMapLoadOrStoreCollision/*syncx.SyncMap[int,int]-8   	312317070	         3.799 ns/op
BenchmarkMapLoadOrStoreCollision/*syncx.Map[int,int]-8       	1000000000	         1.800 ns/op
BenchmarkMapLoadAndDeleteBalanced/*syncx.SyncMap[int,int]-8  	242927908	         4.978 ns/op
BenchmarkMapLoadAndDeleteBalanced/*syncx.Map[int,int]-8      	656539690	         3.040 ns/op
BenchmarkMapLoadAndDeleteUnique/*syncx.SyncMap[int,int]-8    	599912763	         1.958 ns/op
BenchmarkMapLoadAndDeleteUnique/*syncx.Map[int,int]-8        	1000000000	         1.186 ns/op
BenchmarkMapLoadAndDeleteCollision/*syncx.SyncMap[int,int]-8 	361118110	         3.283 ns/op
BenchmarkMapLoadAndDeleteCollision/*syncx.Map[int,int]-8     	669603936	         2.184 ns/op
BenchmarkMapRange/*syncx.SyncMap[int,int]-8                  	  744468	      1594 ns/op
BenchmarkMapRange/*syncx.Map[int,int]-8                      	  887029	      1378 ns/op
BenchmarkMapAdversarialAlloc/*syncx.SyncMap[int,int]-8       	 8579004	       139.1 ns/op
BenchmarkMapAdversarialAlloc/*syncx.Map[int,int]-8           	10151095	       117.4 ns/op
BenchmarkMapAdversarialDelete/*syncx.SyncMap[int,int]-8      	28653038	        42.42 ns/op
BenchmarkMapAdversarialDelete/*syncx.Map[int,int]-8          	37976685	        31.75 ns/op
BenchmarkMapDeleteCollision/*syncx.SyncMap[int,int]-8        	688516219	         1.797 ns/op
BenchmarkMapDeleteCollision/*syncx.Map[int,int]-8            	1000000000	         1.288 ns/op
BenchmarkMapSwapCollision/*syncx.SyncMap[int,int]-8          	 8715524	       137.1 ns/op
BenchmarkMapSwapCollision/*syncx.Map[int,int]-8              	10524654	       114.5 ns/op
BenchmarkMapSwapMostlyHits/*syncx.SyncMap[int,int]-8         	61979068	        20.24 ns/op
BenchmarkMapSwapMostlyHits/*syncx.Map[int,int]-8             	100000000	        14.85 ns/op
BenchmarkMapSwapMostlyMisses/*syncx.SyncMap[int,int]-8       	 2543487	       471.1 ns/op
BenchmarkMapSwapMostlyMisses/*syncx.Map[int,int]-8           	 3203072	       374.4 ns/op
BenchmarkMapCompareAndSwapCollision/*syncx.SyncMap[int,int]-8         	74106286	        17.56 ns/op
BenchmarkMapCompareAndSwapCollision/*syncx.Map[int,int]-8             	 2704873	       448.8 ns/op
BenchmarkMapCompareAndSwapNoExistingKey/*syncx.SyncMap[int,int]-8     	560624862	         2.116 ns/op
BenchmarkMapCompareAndSwapNoExistingKey/*syncx.Map[int,int]-8         	1000000000	         1.417 ns/op
BenchmarkMapCompareAndSwapValueNotEqual/*syncx.SyncMap[int,int]-8     	264941869	         4.297 ns/op
BenchmarkMapCompareAndSwapValueNotEqual/*syncx.Map[int,int]-8         	 5341881	       225.9 ns/op
BenchmarkMapCompareAndSwapMostlyHits/*syncx.SyncMap[int,int]-8        	74957446	        15.56 ns/op
BenchmarkMapCompareAndSwapMostlyHits/*syncx.Map[int,int]-8            	188375540	         8.393 ns/op
BenchmarkMapCompareAndSwapMostlyMisses/*syncx.SyncMap[int,int]-8      	215614034	         5.475 ns/op
BenchmarkMapCompareAndSwapMostlyMisses/*syncx.Map[int,int]-8          	514081917	         2.111 ns/op
BenchmarkMapCompareAndDeleteCollision/*syncx.SyncMap[int,int]-8       	100000000	        11.49 ns/op
BenchmarkMapCompareAndDeleteCollision/*syncx.Map[int,int]-8           	221624979	        10.07 ns/op
BenchmarkMapCompareAndDeleteMostlyHits/*syncx.SyncMap[int,int]-8      	55139349	        21.67 ns/op
BenchmarkMapCompareAndDeleteMostlyHits/*syncx.Map[int,int]-8          	129426009	        12.53 ns/op
BenchmarkMapCompareAndDeleteMostlyMisses/*syncx.SyncMap[int,int]-8    	500900576	         2.476 ns/op
BenchmarkMapCompareAndDeleteMostlyMisses/*syncx.Map[int,int]-8        	486840268	         2.307 ns/op