GoでGeneratorとWorkerのサンプル

大量のJobを捌くこんなケースのGoのサンプルです。

• （Jobの）Generator
  • 生成処理が軽くて1つ
  • 途中でコケたら全体を中断したい（それ以上Jobを生成しない）
• Worker
  • それなりに重たい処理
  • 横に複数並べてJobを捌きたい
  • エラーが発生しても次のJobを継続（最後にエラーをハンドルする）

処理結果とエラーを持つResult構造体

// packageとimportは割愛

type Result struct {
	s string
	e error
}

GeneratorはJobの数だけ値を入れてくチャネルを作って渡します。
時々コケると、途中で止めます。

func Generator(num int) <-chan int {
	ch := make(chan int)
	go func() {
		i := 0
		defer close(ch)

		for num > 0 {
			num--

			// sometime happen error, quit to generate
			if rand.Int()%10 == 0 {
				log.Printf("quit to generate\n")
				return
			}

			ch <- i
			log.Printf("gen: %d\n", i)
			i++
		}
	}()
	return ch
}

workerは入力と出力のチャネルを受けます。
Jobの実行処理の成否をResultとして詰めていきます。
上位のWaitGroupを受けてDoneするのはお決まりのパターンですね。

func worker(wg *sync.WaitGroup, chIn <-chan int, chResult chan<- *Result) {
	defer wg.Done()

	for {
		i, ok := <-chIn
		if !ok {
			return
		}
		log.Printf("work: %d\n", i)

		time.Sleep(2 * time.Second)

		// sometime happen error, continue to work
		if rand.Int()%10 == 0 {
			log.Printf("continue to work\n")
			chResult <- &Result{
				s: strconv.Itoa(i),
				e: fmt.Errorf("error %d", i),
			}
			continue
		}

		chResult <- &Result{
			s: strconv.Itoa(i),
		}
	}
}

本体は上記のGeneratorとworkerとチャネルを取り持つだけです。
Resultの順番はバラけてしまうので、私は配列に詰め直してソートしたりしました。

func main() {
	log.Printf("---start\n")
	rand.Seed(time.Now().UnixNano())

	job := 10
	chIn := Generator(job)
	chResult := make(chan *Result, job)

	var wg sync.WaitGroup
	concurrency := 3
	for i := 0; i < concurrency; i++ {
		wg.Add(1)
		go worker(&wg, chIn, chResult)
	}
	wg.Wait()
	close(chResult)
	log.Printf("---worker done\n")

	for r := range chResult {
		if r.e == nil {
			log.Printf("result %s ok\n", r.s)
		} else {
			log.Printf("result %s err: %s\n", r.s, r.e)
		}
	}
}