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

egonelbre/gophersによるGitHubからの画像

#要約

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を使った並列処理などは特にデバッグが難しいので、できるだけ正しい知識を身に着けて実装・レビューしたいものです。