使用contexts来避免goroutines泄露

context包通过context的Done通道(channel)使得管理在同一个调用路径下的链条式调用变成了可能。

在本文中，将审查怎么使用context包来避免goroutines的泄露。

假定有一个启用一个内部goroutine的函数。一旦调用此函数，调用者就可能无法终止这个函数启动的goroutine。

// gen is a broken generator that will leak a goroutine.
func gen() <-chan int {
    ch := make(chan int)
    go func() {
        var n int
        for {
            ch <- n
            n++
        }
    }()
    return ch
}

上面的生成器启动一个无限循环的goroutine，但调用者将在值达到5时销毁掉。

// The call site of gen doesn't have a 
for n := range gen() {
    fmt.Println(n)
    if n == 5 {
        break
    }
}

一旦调用者调用了这个生成器，goroutine将执行无限循环永远地执行下去。代码中将会泄露一个goroutine。

可以通过向一个停止通道中发送信号至内部goroutine来避免这个问题，但是这里有一个更好的解决方案：可取消的contexts。生成器通过select监听context的Done通道，一旦context的完成，内部goroutine将被取消。

// gen is a generator that can be cancellable by cancelling the ctx.
func gen(ctx context.Context) <-chan int {
    ch := make(chan int)
    go func() {
        var n int
        for {
            select {
            case <-ctx.Done():
                return // avoid leaking of this goroutine when ctx is done.
            case ch <- n:
                n++
            }
        }
    }()
    return ch
}

现在调用者在完成任务进行销毁时可以发生信号至生成器。一旦取消函数被调用，内部goroutine将被返回。

ctx,cancel := context.WithCancel(context.Background())
defer cancel() // make sure all paths cancel the context to avoid context leak

for n := range gen(ctx) {
    fmt.Println(n)
    if n == 5 {
        cancel()
        break
    }
}

// ...

完整的示例代码如下：

package main

import (
    "context"
    "fmt"
)

func gen(ctx context.Context) <-chan int {
    ch := make(chan int)
    go func() {
        var n int
        for {
            select {
            case <-ctx.Done():
                return
            case ch <- n:
                n++
            }
        }
    }()
    return ch
}

func main() {
    ctx,cancel := context.WithCancel(context.Background())
    defer cancel()

    for n := range gen(ctx) {
        fmt.Println(n)
        if n == 5 {
            cancel()
            break
        }
    }
}