使用contexts來避免goroutines泄露

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