golangのチャンネルでセマフォ的なナニカ

mattnさんのエントリに触発されて、某所で使用したちょっと変わったgolangのチャンネルの使い方をご紹介します。

mattn.kaoriya.net

特定の処理の並列度をある程度までに抑えたい、みたいなコトありますよね？
例えばCPUヘビーな処理の並列数をたかだかコア数くらいまでに抑えたい、とか。

そんなときはバッファ付きチャンネルを用意しておいて、当該処理の前後でそのチャンネルにwrite/readをすることで、セマフォ的な制御ができます。

以下のようなカンジです。

package main

import (
    "fmt"
    "sync"
    "time"
)

var ch chan int = make(chan int, 4) // 並列度を4に制限

func heavyFunc(i int) {
    ch <- 1 // チャンネルのバッファがイッパイになっていたら、ブロックする
    defer func() { <-ch }()

    fmt.Println("start:", i)
    time.Sleep(time.Second)
    fmt.Println("end:", i)
}

func main() {
    wg := sync.WaitGroup{}
    for i := 0; i < 20; i++ {
        wg.Add(1)
        go func(i int) {
            heavyFunc(i)
            wg.Done()
        }(i)
    }
    wg.Wait()
}

並列数が上限に達している際にブロックさせたくない場合は、ちょっと変えて以下のようなカンジにすると良いです。

func heavyFunc(i int) error {
    if len(ch) == cap(ch) {
        // バッファがイッパイだったらさっさとerrorを返す
        return errors.New("too busy!")

    }
    ch <- 1
    defer func() { <-ch }()

    fmt.Println("start:", i)
    time.Sleep(time.Second)
    fmt.Println("end:", i)

    return nil
}