概要

小ネタ。結構便利なので個人的にはよく使っているパターンです。

C#には、複数のタスクをまとめてしまって、その内の全てが完了したら完了扱いにしてくれるタスクを返す

Task.WhenAll() というメソッドがあります。

Goの場合でも、非同期処理していると走らせている複数の処理の全部が完了したら、先に進みたいときがあります。

そのような場合に便利かもしれません。基本的に sync.WaitGroupがあるので、それをメイン側で作って渡しておいて wg.Wait() ってやるのと変わらないのですが、ちょっとしたプログラム作ってるときで一々wg作るの面倒なときとかによく利用してる関数です。

サンプル

package chans

import (
    "sync"
)

// WhenAll -- 指定した１つ以上のチャネルの全てが閉じられたら、閉じるチャネルを返します。
//
// チャネルを一つも渡さずに呼び出すと、既に close 済みのチャネルを返します。
func WhenAll(channels ...<-chan struct{}) <-chan struct{} {
    switch len(channels) {
    case 0:
        nilCh := make(chan struct{})
        close(nilCh)

        return nilCh
    case 1:
        return channels[0]
    }

    allDone := make(chan struct{})
    go func() {
        defer close(allDone)

        wg := sync.WaitGroup{}
        wg.Add(len(channels))

        for _, v := range channels {
            go func(ch <-chan struct{}) {
                defer wg.Done()
                <-ch
            }(v)
        }

        wg.Wait()
    }()

    return allDone
}

gomy/whenall.go at master · devlights/gomy · GitHub

以下、ユニットテストコードです。

package chans

import (
    "testing"
    "time"
)

func TestWhenAll(t *testing.T) {
    type (
        testin struct {
            makeChCount int
        }
        testout struct {
            limit time.Duration
        }
        testcase struct {
            in  testin
            out testout
        }
    )

    cases := []testcase{
        {
            in:  testin{makeChCount: 0},
            out: testout{100 * time.Millisecond},
        },
        {
            in:  testin{makeChCount: 1},
            out: testout{150 * time.Millisecond},
        },
        {
            in:  testin{makeChCount: 2},
            out: testout{250 * time.Millisecond},
        },
        {
            in:  testin{makeChCount: 3},
            out: testout{350 * time.Millisecond},
        },
        {
            in:  testin{makeChCount: 4},
            out: testout{450 * time.Millisecond},
        },
        {
            in:  testin{makeChCount: 5},
            out: testout{550 * time.Millisecond},
        },
        {
            in:  testin{makeChCount: 6},
            out: testout{650 * time.Millisecond},
        },
    }

    makeCh := func(closeAfter time.Duration) <-chan struct{} {
        ch := make(chan struct{})
        go func() {
            defer close(ch)
            time.Sleep(closeAfter)
        }()

        return ch
    }

    for _, c := range cases {
        func() {
            chList := make([]<-chan struct{}, 0, c.in.makeChCount)
            for i := 0; i < c.in.makeChCount; i++ {
                ch := makeCh(time.Duration((i+1)*100) * time.Millisecond)
                chList = append(chList, ch)
            }

            start := time.Now()
            if _, ok := <-WhenAll(chList...); ok {
                t.Errorf("want: false\tgot: %v", ok)
            }

            elapsed := time.Since(start)
            t.Logf("len(ch)=%d\telapsed=%v\n", len(chList), elapsed)

            if c.out.limit < elapsed {
                t.Errorf("want: within %v\tgot %v", c.out.limit, elapsed)
            }
        }()
    }
}

実行すると以下のようになります。

$ go test -v github.com/devlights/gomy/chans -run ^TestWhenAll$
=== RUN   TestWhenAll
    TestWhenAll: whenall_test.go:77: len(ch)=0  elapsed=0s
    TestWhenAll: whenall_test.go:77: len(ch)=1  elapsed=100.4317ms
    TestWhenAll: whenall_test.go:77: len(ch)=2  elapsed=200.8921ms
    TestWhenAll: whenall_test.go:77: len(ch)=3  elapsed=300.023ms
    TestWhenAll: whenall_test.go:77: len(ch)=4  elapsed=401.4599ms
    TestWhenAll: whenall_test.go:77: len(ch)=5  elapsed=501.9605ms
    TestWhenAll: whenall_test.go:77: len(ch)=6  elapsed=600.0772ms
--- PASS: TestWhenAll (2.10s)
PASS
ok      github.com/devlights/gomy/chans 2.187s