goのWaitGroupでハマった話

goで並列処理をするときにsync.WautGroupを使って、Goroutineが終わるのを待つ処理を書いて、意図した通りに動かず30分ほど無駄にした話です。

結論

先に結論を書いておくと、Goroutineの中でAddするのではなく、Goroutineを呼び出す前にAddをしておけば問題ないです。

func something(wg *sync.WaitGroup) {
    fmt.Println("hello")
    defer wg.Done()
}

func main() {
    var wg sync.WaitGroup

    for i := 0; i < 2; i++ {
        // これをGorutineの中でやらない
        wg.Add(1) 
        go something(&wg)
    }
    fmt.Println("wait")
    wg.Wait()
    fmt.Println("end")
}

なにでハマったのか

最初に書いていたコードはこんな感じ。

func something(){
    fmt.Println("hello")
}

func main() {

    for i := 0; i < 2; i++ {
        something()
    }
    fmt.Println("end")
}

これはhelloが2回表示されたあとにendが表示されるだけの、なんの問題もないコード。

hello
hello
end

mainから呼んでいるsomethingをGoroutineとして実行するように、somethingの呼び出し時にgoを追加。

func something(){
    fmt.Println("hello")
}

func main() {

    for i := 0; i < 2; i++ {
        go something()
    }
    fmt.Println("end")
}

この状態だとmainはGoroutineの終了をまたずに終了してしまうので、実行すると下記のように出力された。

end

このままだとGoroutineの処理が最後まで終わらないので、mainがsomethingが終わってから終わるようにsync.WautGroupを使うように変更しました。

func something(wg *sync.WaitGroup) {
    fmt.Println("hello")
    wg.Add(1)
    defer wg.Done()
}

func main() {
    var wg sync.WaitGroup

    for i := 0; i < 2; i++ {
        go something(&wg)
    }
    fmt.Println("wait")
    wg.Wait()
    fmt.Println("end")
}

これで以下のような出力になることを想定してました。

wait
hello
hello
end

が、実際に出力されたのは下記。

wait
end

helloが出力されないままmainが終わってしまってるようです。

試しにforの中でsleepするように変更する。

func something(wg *sync.WaitGroup) {
    fmt.Println("hello")
    wg.Add(1)
    defer wg.Done()
}

func main() {
    var wg sync.WaitGroup

    for i := 0; i < 2; i++ {
        go something(&wg)
        time.Sleep(1 * time.Millisecond)
    }
    fmt.Println("wait")
    wg.Wait()
    fmt.Println("end")
}

helloが出力されるようになりました。
これで一応はGoroutineが終わるまでまってmainが終わるようになりました…
(sleepさせる時間は、少し試してみたところ1msほどsleepさせれば安定してhelloが出力された)

hello
hello
wait
end

sleepをforループの中で毎回するのは、ループの回数が増えてきた単純に遅くなっていくのでなるべく避けたいので、調べてみたら下記の記事でわかりやすく説明してくれてました。
https://qiita.com/ruiu/items/dba58f7b03a9a2ffad65
Goroutineがスケジューリングされるタイミングは即時ではなく、任意なのでGoroutineが動いてAddされる前にWaitまで到達することがあるということで、Goroutineの中ではなく、Goroutineの呼び出し前にAddをしておく必要があるということでした。
最初にも載せてありますが、直したコードがこちら。

func something(wg *sync.WaitGroup) {
    fmt.Println("hello")
    defer wg.Done()
}

func main() {
    var wg sync.WaitGroup

    for i := 0; i < 2; i++ {
        wg.Add(1)
        go something(&wg)
    }
    fmt.Println("wait")
    wg.Wait()
    fmt.Println("end")
}

今までは、Goroutineを生成する個数自体が多かったり、Goroutineを呼び出してからすぐにWaitをすることがなかったので、気づくことがなかっただけで、何回かGoroutineの中でAddをしているコード書いてしまっていました。しかもそれがなんの問題もなく動いていたので、原因に気づくのにずいぶん遠回りすることになりました。
今後は気をつけて書いていこうと思います。