go语言如何通过channel或context来实现协程等待？

Question

想实现主协程等待n个子协程的效果，类似于sync.waitGroup。

Answer 1

示例如下：

// hello.go
package main

import (
    "fmt"
    "math/rand"
    "time"
)

func main() {
    num := 10
    ch := make(chan int, num)

    for i := 0; i < num; i++ {
        go func(ch chan int, key int) {
            fmt.Println("goroutine start: ", key)
            t := rand.Intn(5) + 1
            time.Sleep(time.Duration(t) * time.Second)
            fmt.Println("goroutine end: ", key)
            // 为了不使 goroutine 阻塞，chan 创建的时候加了缓冲量
            ch <- key
        }(ch, i)
    }

    fmt.Println("main start")
    for i := 0; i < num; i++ {
        key := <-ch
        fmt.Println("main get: ", key)
    }
    fmt.Println("main end")
}

$ go run ./hello.go 
main start
goroutine start:  9
goroutine start:  4
goroutine start:  1
goroutine start:  2
goroutine start:  8
goroutine start:  0
goroutine start:  3
goroutine start:  7
goroutine start:  5
goroutine start:  6
goroutine end:  6
main get:  6
goroutine end:  7
main get:  7
goroutine end:  3
main get:  3
goroutine end:  5
main get:  5
goroutine end:  9
main get:  9
goroutine end:  8
main get:  8
goroutine end:  1
main get:  1
goroutine end:  4
main get:  4
goroutine end:  0
main get:  0
goroutine end:  2
main get:  2
main end

简单封装

package main

import (
    "fmt"
    "math/rand"
    "time"
)

type CustomWaitGroup struct {
    ch  chan int
    num int
}

func NewCustomWaitGroup(num int) *CustomWaitGroup {
    return &CustomWaitGroup{
        ch:  make(chan int, num),
        num: num,
    }
}

func (r *CustomWaitGroup) Done() {
    r.ch <- 1
}

func (r *CustomWaitGroup) Wait() {
    for i := 0; i < r.num; i++ {
        <-r.ch
    }
}

func main() {

    wg := NewCustomWaitGroup(3)

    f := func(key int) {
        fmt.Println("goroutine start: ", key)
        t := rand.Intn(5) + 1
        time.Sleep(time.Duration(t) * time.Second)
        fmt.Println("goroutine end: ", key)
        wg.Done()
    }
    go f(1)
    go f(2)
    go f(3)

    fmt.Println("main start")
    wg.Wait()
    fmt.Println("main end")
}

$ go run ./hello.go 
main start
goroutine start:  3
goroutine start:  1
goroutine start:  2
goroutine end:  3
goroutine end:  1
goroutine end:  2
main end

Answer 2

n个子协程, 读取n次chan就可以了. 子协程完成时, 记得发送chan数据.