Go 语言实现多路归并排序

第一个版本

基于 goroutine ,我们很容易将串行的归并排序改为并行的归并排序,只需要将每个 mergeSort 多个 goroutine,再使用 sync.WaitGroup 等待所有 goroutine 都完成后,再进行 merge 操作。于是可以得到第一个版本的代码如下:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
func MergeSort(src []int64) {
	if len(src) > 1 {
		mid := len(src) / 2

		var wg sync.WaitGroup
		wg.Add(2)

		go func() {
			defer wg.Done()
			MergeSort(src[mid:])
		}()

		go func() {
			defer wg.Done()
			MergeSort(src[:mid])
		}()

		wg.Wait()
		merge(src, mid)
	}
}

执行 make test,程序能够通过测试,得出正确的结果。 接着执行 make bench,发现程序运行得非常慢,得到的测试数据如下:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
go test -bench Benchmark -run xx -count 5 -benchmem
goos: darwin
goarch: amd64
pkg: pingcap/talentplan/tidb/mergesort
BenchmarkMergeSort-4    	       1	23657391787 ns/op	4152707152 B/op	36435339 allocs/op
BenchmarkMergeSort-4    	       1	22347034716 ns/op	3647074960 B/op	35318703 allocs/op
BenchmarkMergeSort-4    	       1	20944485290 ns/op	3613508272 B/op	35102621 allocs/op
BenchmarkMergeSort-4    	       1	23675804874 ns/op	3605332784 B/op	35000365 allocs/op
BenchmarkMergeSort-4    	       1	19576917879 ns/op	3632500240 B/op	35247621 allocs/op
BenchmarkNormalSort-4   	       1	3341777012 ns/op	      64 B/op	       2 allocs/op
BenchmarkNormalSort-4   	       1	3343564257 ns/op	      64 B/op	       2 allocs/op
BenchmarkNormalSort-4   	       1	3347931663 ns/op	      64 B/op	       2 allocs/op
BenchmarkNormalSort-4   	       1	3339944569 ns/op	      64 B/op	       2 allocs/op
BenchmarkNormalSort-4   	       1	3347761882 ns/op	      64 B/op	       2 allocs/op
PASS
ok  	pingcap/talentplan/tidb/mergesort	133.733s

发现这样子的并行程序,比串行的还要慢上一个数量级,慢了快 10 倍。

我们使用 PProf 进行分析其非常慢的原因,对测试用例执行以下命令:

go test -bench=. -cpuprofile=cpu.prof

接着执行:

go tool pprof -http=:8080 cpu.prof

启动 PProf 可视化界面

pprof

可以看到 CPU 占用最高的三个函数分别是 runtime.memmove, runtime.usleepruntime.mallocgc ,它们都是 gc 相关的,因为我们创建的 goroutine 给 gc 造成了很大的负担。

究其原因,可以假设我们要排序的数组元素有 1024 个,一次分割后 512 个,然后 256 个,直到 1 个元素,这意味着即便最后只有一个元素,我们还是要为之创建一个独立的 goroutine. 虽然 goroutine 非常轻量,但为了一个元素也创建一个 goroutine,调度的代价也是非常昂贵的,所以我们的代码速度非常慢——甚至不如线性的排序。

第二个版本

既然 goroutine 创建得太多,那么优化思路就是减少 goroutine 的创建。对于元素很少的数组,可以没有必要创建 goroutine 来进行运算。 对此我们定义一个阈值,经过测试我们可以选择 2048,当元素个数小于 2048 个的时候,直接使用串行的排序进行计算。由此我们可以得到第二版的代码:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
func MergeSort(src []int64) {
	if len(src) > 1 {
		if len(src) <= 1<<11 {
			sort.Slice(src, func(i, j int) bool {
				return src[i] < src[j]
			})
		} else {
			mid := len(src) / 2

			var wg sync.WaitGroup
			wg.Add(2)

			go func() {
				defer wg.Done()
				MergeSort(src[mid:])
			}()

			go func() {
				defer wg.Done()
				MergeSort(src[:mid])
			}()

			wg.Wait()
			merge(src, mid)
		}
	}
}

对此代码执行 make bench 得到数据:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
go test -bench Benchmark -run xx -count 5 -benchmem
goos: darwin
goarch: amd64
pkg: pingcap/talentplan/tidb/mergesort
BenchmarkMergeSort-4    	       1	1354622676 ns/op	1747587872 B/op	   43228 allocs/op
BenchmarkMergeSort-4    	       1	1356477088 ns/op	1745991776 B/op	   39274 allocs/op
BenchmarkMergeSort-4    	       1	1211078503 ns/op	1745841472 B/op	   37324 allocs/op
BenchmarkMergeSort-4    	       1	1244470076 ns/op	1745950240 B/op	   38694 allocs/op
BenchmarkMergeSort-4    	       1	1263144506 ns/op	1745927104 B/op	   38524 allocs/op
BenchmarkNormalSort-4   	       1	3370295794 ns/op	      64 B/op	       2 allocs/op
BenchmarkNormalSort-4   	       1	3393553236 ns/op	      64 B/op	       2 allocs/op
BenchmarkNormalSort-4   	       1	3420058771 ns/op	      64 B/op	       2 allocs/op
BenchmarkNormalSort-4   	       1	3371727069 ns/op	      64 B/op	       2 allocs/op
BenchmarkNormalSort-4   	       1	3471338096 ns/op	      64 B/op	       2 allocs/op
PASS
ok  	pingcap/talentplan/tidb/mergesort	27.173s

可以看到,经过这次的优化,程序的速度已经比线性的快上不少了,整个 benchmark 花费时间减少到了 27s.

那么,沿着减少多余 goroutine 的这个思路,我们还可不可以再继续优化呢?

第三个版本

在上一个版本中,其实我们并不需要为每个 MergeSort 创建一个 goroutine. 我们其实只需要对第一个 MergeSort 创建一个新的 goroutine,而第二个 MergeSort 可以在当前的 goroutine 中运行,这样可以得到第三个版本的代码:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
go test -bench Benchmark -run xx -count 5 -benchmem
goos: darwin
goarch: amd64
pkg: pingcap/talentplan/tidb/mergesort
BenchmarkMergeSort-4    	       1	1249984680 ns/op	1746230432 B/op	   36274 allocs/op
BenchmarkMergeSort-4    	       1	1227243541 ns/op	1745575456 B/op	   33964 allocs/op
BenchmarkMergeSort-4    	       1	1291195594 ns/op	1745596800 B/op	   34208 allocs/op
BenchmarkMergeSort-4    	       1	1227169037 ns/op	1745558768 B/op	   33735 allocs/op
BenchmarkMergeSort-4    	       1	1212597819 ns/op	1745663040 B/op	   34843 allocs/op
BenchmarkNormalSort-4   	       1	3430862175 ns/op	      64 B/op	       2 allocs/op
BenchmarkNormalSort-4   	       1	3478333443 ns/op	      64 B/op	       2 allocs/op
BenchmarkNormalSort-4   	       1	3362342652 ns/op	      64 B/op	       2 allocs/op
BenchmarkNormalSort-4   	       1	3412419430 ns/op	      64 B/op	       2 allocs/op
BenchmarkNormalSort-4   	       1	3462869950 ns/op	      64 B/op	       2 allocs/op
PASS
ok  	pingcap/talentplan/tidb/mergesort	27.030s

最后的 benchmark 显示性能得到微小的提升,但怎么说也算是提升了。