Go Context 并发编程简明教程

1 为什么需要 Context

WaitGroup 和信道(channel)是常见的 2 种并发控制的方式。

如果并发启动了多个子协程，需要等待所有的子协程完成任务，WaitGroup 非常适合于这类场景，例如下面的例子：

1var wg sync.WaitGroup
2
3func doTask(n int) {
4	time.Sleep(time.Duration(n))
5	fmt.Printf("Task %d Done\n", n)
6	wg.Done()
7}
8
9func main() {
10	for i := 0; i < 3; i++ {
11		wg.Add(1)
12		go doTask(i + 1)
13	}
14	wg.Wait()
15	fmt.Println("All Task Done")
16}

wg.Wait() 会等待所有的子协程任务全部完成，所有子协程结束后，才会执行 wg.Wait() 后面的代码。

1Task 3 Done
2Task 1 Done
3Task 2 Done
4All Task Done

WaitGroup 只是傻傻地等待子协程结束，但是并不能主动通知子协程退出。假如开启了一个定时轮询的子协程，有没有什么办法，通知该子协程退出呢？这种场景下，可以使用 select + chan 的机制。

1var stop chan bool
2
3func reqTask(name string) {
4	for {
5		select {
6		case <-stop:
7			fmt.Println("stop", name)
8			return
9		default:
10			fmt.Println(name, "send request")
11			time.Sleep(1 * time.Second)
12		}
13	}
14}
15
16func main() {
17	stop = make(chan bool)
18	go reqTask("worker1")
19	time.Sleep(3 * time.Second)
20	stop <- true
21	time.Sleep(3 * time.Second)
22}

子协程使用 for 循环定时轮询，如果 stop 信道有值，则退出，否则继续轮询。

1worker1 send request
2worker1 send request
3worker1 send request
4stop worker1

更复杂的场景如何做并发控制呢？比如子协程中开启了新的子协程，或者需要同时控制多个子协程。这种场景下，select + chan的方式就显得力不从心了。

Go 语言提供了 Context 标准库可以解决这类场景的问题，Context 的作用和它的名字很像，上下文，即子协程的下上文。Context 有两个主要的功能：

• 通知子协程退出（正常退出，超时退出等）；
• 传递必要的参数。

2 context.WithCancel

context.WithCancel() 创建可取消的 Context 对象，即可以主动通知子协程退出。

2.1 控制单个协程

使用 Context 改写上述的例子，效果与 select + chan 相同。

1func reqTask(ctx context.Context, name string) {
2	for {
3		select {
4		case <-ctx.Done():
5			fmt.Println("stop", name)
6			return
7		default:
8			fmt.Println(name, "send request")
9			time.Sleep(1 * time.Second)
10		}
11	}
12}
13
14func main() {
15	ctx, cancel := context.WithCancel(context.Background())
16	go reqTask(ctx, "worker1")
17	time.Sleep(3 * time.Second)
18	cancel()
19	time.Sleep(3 * time.Second)
20}

• context.Backgroud() 创建根 Context，通常在 main 函数、初始化和测试代码中创建，作为顶层 Context。
• context.WithCancel(parent) 创建可取消的子 Context，同时返回函数 cancel。
• 在子协程中，使用 select 调用 <-ctx.Done() 判断是否需要退出。
• 主协程中，调用cancel()函数通知子协程退出。

2.2 控制多个协程

1func main() {
2	ctx, cancel := context.WithCancel(context.Background())
3	
4	go reqTask(ctx, "worker1")
5	go reqTask(ctx, "worker2")
6	
7	time.Sleep(3 * time.Second)
8	cancel()
9	time.Sleep(3 * time.Second)
10}

为每个子协程传递相同的上下文 ctx 即可，调用 cancel() 函数后该 Context 控制的所有子协程都会退出。

1worker1 send request
2worker2 send request
3worker1 send request
4worker2 send request
5worker1 send request
6worker2 send request
7stop worker1
8stop worker2

3 context.WithValue

如果需要往子协程中传递参数，可以使用 context.WithValue()。

1type Options struct{ Interval time.Duration }
2
3func reqTask(ctx context.Context, name string) {
4	for {
5		select {
6		case <-ctx.Done():
7			fmt.Println("stop", name)
8			return
9		default:
10			fmt.Println(name, "send request")
11			op := ctx.Value("options").(*Options)
12			time.Sleep(op.Interval * time.Second)
13		}
14	}
15}
16
17func main() {
18	ctx, cancel := context.WithCancel(context.Background())
19	vCtx := context.WithValue(ctx, "options", &Options{1})
20
21	go reqTask(vCtx, "worker1")
22	go reqTask(vCtx, "worker2")
23
24	time.Sleep(3 * time.Second)
25	cancel()
26	time.Sleep(3 * time.Second)
27}

• context.WithValue() 创建了一个基于 ctx 的子 Context，并携带了值 options。
• 在子协程中，使用 ctx.Value("options") 获取到传递的值，读取/修改该值。

4 context.WithTimeout

如果需要控制子协程的执行时间，可以使用 context.WithTimeout 创建具有超时通知机制的 Context 对象。

1func main() {
2	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
3	go reqTask(ctx, "worker1")
4	go reqTask(ctx, "worker2")
5
6	time.Sleep(3 * time.Second)
7	fmt.Println("before cancel")
8	cancel()
9	time.Sleep(3 * time.Second)
10}

WithTimeout()的使用与 WithCancel() 类似，多了一个参数，用于设置超时时间。执行结果如下：

1worker2 send request
2worker1 send request
3worker1 send request
4worker2 send request
5stop worker2
6stop worker1
7before cancel

因为超时时间设置为 2s，但是 main 函数中，3s 后才会调用 cancel()，因此，在调用 cancel() 函数前，子协程因为超时已经退出了。

5 context.WithDeadline

超时退出可以控制子协程的最长执行时间，那context.WithDeadline()则可以控制子协程的最迟退出时间。

1func reqTask(ctx context.Context, name string) {
2	for {
3		select {
4		case <-ctx.Done():
5			fmt.Println("stop", name, ctx.Err())
6			return
7		default:
8			fmt.Println(name, "send request")
9			time.Sleep(1 * time.Second)
10		}
11	}
12}
13
14func main() {
15	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(1*time.Second))
16	go reqTask(ctx, "worker1")
17	go reqTask(ctx, "worker2")
18
19	time.Sleep(3 * time.Second)
20	fmt.Println("before cancel")
21	cancel()
22	time.Sleep(3 * time.Second)
23}

• WithDeadline 用于设置截止时间。在这个例子中，将截止时间设置为1s后，cancel() 函数在 3s 后调用，因此子协程将在调用 cancel() 函数前结束。

• 在子协程中，可以通过 ctx.Err() 获取到子协程退出的错误原因。

运行结果如下：

1worker2 send request
2worker1 send request
3stop worker2 context deadline exceeded
4stop worker1 context deadline exceeded
5before cancel

可以看到，子协程 worker1 和 worker2 均是因为截止时间到了而退出。