GOLANG使用Context管理關聯goroutine

一般一個業務很少不用到 goroutine 的，因為很多方法是需要等待的，例如http.Server.ListenAndServe這個就是等待的，除非關閉了 Server 或 Listener，否則是不會返回的。除非是一個 API 伺服器，否則肯定需要另外起 goroutine 發起其他的服務，而且對於 API 伺服器來說，在http.Handler的處理函式中一般也需要起 goroutine，如何管理這些 goroutine，在 GOLANG1.7 提供context.Context。

先看一個簡單的，如果啟動兩個 goroutine，一個是 HTTP，還有個訊號處理的收到退出訊號做清理：

wg := sync.WaitGroup{}
defer wg.Wait()

wg.Add(1)
go func() {
    defer wg.Done()

    ss := make(os.Signal, 0)
    signal.Notify(ss, syscall.SIGINT, syscall.SIGTERM)
    for s := ss {
        fmt.Println("Got signal", s)
        break
    }
}()

wg.Add(1)
go func() {
    defer wg.Done()

    svr := &http.Server{ Addr:":8080", Handler:nil, }
    fmt.Println(svr.ListenAndServe())
}

很清楚，起了兩個 goroutine，然後用 WaitGroup 等待它們退出。如果它們之間沒有互動，不互相影響，那真的是蠻簡單的，可惜這樣是不行的，因為訊號的 goroutine 收到退出訊號後，應該通知 server 退出。暴力一點的是直接呼叫svr.Close()，但是如果有些請求還需要取消怎麼辦呢？最好用 Context 了：

wg := sync.WaitGroup{}
defer wg.Wait()

ctx,cancel := context.WithCancel(context.Background())

wg.Add(1)
go func() {
    defer wg.Done()

    ss := make(chan os.Signal, 0)
    signal.Notify(ss, syscall.SIGINT, syscall.SIGTERM)
    select {
    case <- ctx.Done():
        return
    case s := <- ss:
        fmt.Println("Got signal", s)
        cancel() // 取消請求，通知用到ctx的所有goroutine
        return
    }
}()

wg.Add(1)
go func() {
    defer wg.Done()
    defer cancel()

    svr := &http.Server{ Addr:":8080", Handler:nil, }

    go func(){
        select {
        case <- ctx.Done():
            svr.Close()
        }
    }

    fmt.Println(svr.ListenAndServe())
}

這個方式可以在新開 goroutine 時繼續使用，譬如新加一個 goroutine，裡面讀寫了 UDPConn：

wg.Add(1)
go func() {
    defer wg.Done()
    defer cancel()

    var conn *net.UDPConn
    if conn,err = net.Dial("udp", "127.0.0.1:1935"); err != nil {
        fmt.Println("Dial UDP server failed, err is", err)
        return
    }

    fmt.Println(UDPRead(ctx, conn))
}()

UDPRead = func(ctx context.Context, conn *net.UDPConn) (err error) {
    wg := sync.WaitGroup{}
    defer wg.Wait()

    ctx, cancel := context.WithCancel(ctx)

    wg.Add(1)
    go func() {
        defer wg.Done()
        defer cancel()

        for {
            b := make([]byte, core.MTUSize)
            size, _, err := conn.ReadFromUDP(b)
            // 處理UDP包 b[:size]
        }
    }()

    select {
    case <-ctx.Done():
        conn.Close()
    }
    return
}

如果只是用到 HTTP Server，可以這麼寫：

func run(ctx contex.Context) {
    server := &http.Server{Addr: addr, Handler: nil}
    go func() {
        select {
        case <-ctx.Done():
            server.Close()
        }
    }()

    http.HandleFunc("/api", func(w http.ResponseWriter, r *http.Request) {
    })

    fmt.Println(server.ListenAndServe())
}

如果需要提供一個 API 來讓伺服器退出，可以這麼寫：

func run(ctx contex.Context) {
    server := &http.Server{Addr: addr, Handler: nil}

    ctx, cancel := context.WithCancel(ctx)
    http.HandleFunc("/quit", func(w http.ResponseWriter, r *http.Request) {
        cancel() // 使用區域性的ctx和cancel
    })

    go func() {
        select {
        case <-ctx.Done():
            server.Close()
        }
    }()

    fmt.Println(server.ListenAndServe())
}

使用區域性的 ctx 和 cancel，可以避免 cancel 傳入的 ctx，只是影響當前的 ctx。