【小ネタ】Goで複数チャネルのCloseを待ち受ける

要約

Goでプログラムを作っていると、たまに複数チャネルのCloseを待ち受けるということがあります。 あまり馴染みがないのと、ググってもあまり出てこないのでメモ。

問題のコード

メンテをしているGoのプログラムで下記のコードがありました。

• channelが複数(ch1,ch2)あり、for selectで両方受信する
• ch1が先にcloseされる前提となっており、ch1のcloseではcontinueしている
• ch2がcloseされたらLoopを抜ける

package main
import (
    "log"
    "time"
)
func main() {
    ch1 := make(chan struct{})
    ch2 := make(chan struct{})
    go func() {
        defer close(ch1)
        for i := 0; i < 2; i++ {
            ch1 <- struct{}{}
            time.Sleep(time.Second)
        }
    }()
    go func() {
        defer close(ch2)
        for i := 0; i < 5; i++ {
            ch2 <- struct{}{}
            time.Sleep(time.Second)
        }
    }()
Loop:
    for {
        select {
        case _, ok := <-ch1:
            if !ok {
                // log.Println("ch1 is closed")
                continue
            }
            log.Println("ch1")
        case _, ok := <-ch2:
            if !ok {
                // log.Println("ch2 is closed")
                break Loop
            }
            log.Println("ch2")
        }
    }
}

実行すると下記のように出力されます。

$ go run main.go 
2021/03/18 22:57:01 ch2
2021/03/18 22:57:01 ch1
2021/03/18 22:57:02 ch1
2021/03/18 22:57:02 ch2
2021/03/18 22:57:03 ch2
2021/03/18 22:57:04 ch2
2021/03/18 22:57:05 ch2

問題ないように見えますが、ch1がcloseされたあとのログを出力すると2021/03/18 22:59:54 ch1 is closedのようなログが大量に出力されることがわかります。

Receiving operatorの言語仕様を確認します。

A receive operation on a closed channel can always proceed immediately, yielding the element type's zero value after any previously sent values have been received.

出典: Receiving operator - The Go Programming Language Specification

閉じたチャネルの受信操作は、常にすぐZero値を返します。

改善する

ch1がcloseされたら、次の繰り返しではch1を受信するcase文を通らないでほしいですね。 channel closeされたらnil channelにしてブロックするようにします。

Receiving from a nil channel blocks forever.

出典: Receiving operator - The Go Programming Language Specification

 Loop:
        for {
+               if ch1 == nil && ch2 == nil {
+                       break Loop
+               }
                select {
                case _, ok := <-ch1:
                        if !ok {
-                               // log.Println("ch1 is closed")
+                               log.Println("ch1 is closed")
+                               ch1 = nil
                                continue
                        }
                        log.Println("ch1")
                case _, ok := <-ch2:
                        if !ok {
-                               // log.Println("ch2 is closed")
-                               break Loop
+                               log.Println("ch2 is closed")
+                               ch2 = nil
+                               continue
                        }
                        log.Println("ch2")
                }

実行してみます。

$ go run main.go 
2021/03/18 23:03:52 ch2
2021/03/18 23:03:52 ch1
2021/03/18 23:03:53 ch2
2021/03/18 23:03:53 ch1
2021/03/18 23:03:54 ch1 is closed
2021/03/18 23:03:54 ch2
2021/03/18 23:03:55 ch2
2021/03/18 23:03:56 ch2
2021/03/18 23:03:57 ch2 is closed

これで正しく動きましたね。

まとめ

正しく動いていそうで、動いていないプログラムなどがあると厄介ですよね。 goroutineを使った並列処理などは特にデバッグが難しいので、できるだけ正しい知識を身に着けて実装・レビューしたいものです。