goroutineと仲良くする（2）

この記事はgoroutineと仲良くする（1）の続きです。
前回はgoroutineの実行方法を今回はgoroutineの止め方についてです。

sync.WaitGroup

syncパッケージのWaitGroup関数を使用したgoroutineの止め方です。
前回の記事で使用していた以下のコードですが、
プログラムの処理が非同期処理も早く終了してしまうために、結果が出力されないので0.1秒遅延させてます。

package main

import (
    "fmt"
    "time"
)

func printNumbers(begin, end int) {
    for i := begin; i < end; i++ {
        fmt.Println(i)
    }
}

func main() {
    fmt.Println("--- start ---")

    printNumbers(1, 5)
    go printNumbers(6, 10)  // goroutine1
    go printNumbers(11, 15) // goroutine2

    time.Sleep(100 * time.Millisecond) // 0.1秒遅延
    fmt.Println("--- finish ---")
}

本来であれば、goroutineの処理が終了するのを待ってからプログラムの処理を終了させたいところです。
そんな時にsync.WaitGroupを使うことができます。
以下のようにコードになります。

package main

import (
	"fmt"
	"sync"
)

var wg sync.WaitGroup

func printNumbers(begin, end int) {
	for i := begin; i < end; i++ {
		fmt.Println(i)
	}
}

func goPrintNumbers(begin, end int) {
	for i := begin; i < end; i++ {
		fmt.Println(i)
	}
	wg.Done() // wgの数を1つ減らす
}

func main() {
	fmt.Println("--- start ---")

	printNumbers(1, 5)

	wg.Add(2)                 // goroutineの数を2つ増やす
	go goPrintNumbers(6, 10)  // goroutine1
	go goPrintNumbers(11, 15) // goroutine2
	wg.Wait()                 // goroutineが完了するまで待つ

	fmt.Println("--- finish ---")
}

WaitGroupの値を作成し、Addメソッドを使用して、使用するgoroutineの数分追加をする。
Doneメソッドを使用して、1つのgoroutineの処理の終了時にWaitGroupの値を1つ減らす。
WaitメソッドはWaitGroupの値が0になるまで待つ処理です。

context.WithCancel

さらに途中でgoroutineの処理を停止させる必要がある場合があると思います
そんな場合は、go1.7以上であればcontextパッケージを使用することでgoroutineの処理を停止させることができます。

package main

import (
	"context"
	"fmt"
	"sync"
)

var wg sync.WaitGroup

func goPrintNumbers(ctx context.Context, ch chan int) {
	for {
		select {
		case <-ctx.Done():
			fmt.Println("goroutine finish")
			wg.Done() // wgの数を一つ減らす
			return
		case num := <-ch:
			for i := num; i < num+5; i++ {
				fmt.Println(i)
			}
		}
	}
}

func main() {
	fmt.Println("--- start ---")

	// contextを生成
	ctx, cancel := context.WithCancel(context.Background())
	ch := make(chan int)
	for i := 0; i < 2; i++ {
		wg.Add(1) // goroutineの数を追加
		go goPrintNumbers(ctx, ch)
	}

	for number := 1; ; {
		// numberが16になったらgoroutineを停止させる
		if number == 16 {
			cancel() // ctxを終了させる
			break
		}

		ch <- number
		number += 5
	}
	wg.Wait() // goroutineが完了するまで待つ

	fmt.Println("--- finish ---")
}

ctx.WithTimeoutを使用することでタイムアウト処理とキャンセル処理を行うこともできます

参考

みんなのGo言語【現場で使える実践テクニック】