pipeline_by_chan_quit.go

package main

import (
	"fmt"
	"runtime"
	"time"
)

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	lim1, lim2 := 2, 3
	quit1 := make(chan bool)
	quit2 := make(chan bool)
	quit_all := make(chan bool)

	in := input_gen(10, func() {
		for i := 0; i < lim1; i++ {
			quit1 <- true
		}
		for i := 0; i < lim2; i++ {
			quit2 <- true
		}
		quit_all <- true
	})

	out1 := stage1_gen(in, quit1, lim1)
	out2 := stage2_gen(out1, quit2, lim2)

	for {
		select {
		case o := <-out2:
			p("output: ", o)
		case <-quit_all:
			return
		}
	}
}

func input_gen(count int, quit func()) chan int {
	in := make(chan int)
	go func() {
		for i := 0; i < count; i++ {
			in <- i
		}
		quit()
	}()
	return in
}

func stage1_gen(in <-chan int, quit <-chan bool, limit int) chan int {
	out := make(chan int)

	for j := 0; j < limit; j++ {
		go func() {
			for {
				select {
				case i := <-in:
					p("start stage 1 worker with input", i)
					time.Sleep(time.Second / 2)
					p("exit stage 1 worker with input", i)
					out <- i
				case <-quit:
					return
				}
			}
		}()
	}
	return out
}

func stage2_gen(in <-chan int, quit <-chan bool, limit int) chan int {
	out := make(chan int)

	for j := 0; j < limit; j++ {
		go func() {
			for {
				select {
				case i := <-in:
					p("start stage 2 worker with input", i)
					time.Sleep(time.Second / 2)
					p("exit stage 2 worker with input", i)
					out <- i
				case <-quit:
					return
				}
			}
		}()
	}
	return out
}

func p(v ...interface{}) {
	fmt.Println(v...)
}