pipeline_by_call.go

package main

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

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

	out := make(chan int)
	go func() {
		for o := range out {
			p("output:", o)
		}
	}()

	wg := new(sync.WaitGroup)

	input(
		stage1(wg, 2,
			stage2(wg, 3,
				func(o int) {
					out <- o
				})))

	wg.Wait()
	close(out)
}

func input(next func(int)) {
	for i := 0; i < 10; i++ {
		next(i)
	}
}

func stage1(wg *sync.WaitGroup, limit int, next func(int)) func(int) {
	g := NewSyncGroup(limit, wg)
	return func(v int) {
		g.Enter()
		go func(i int) {
			defer g.Done()
			p("start stage 1 worker with input", i)
			time.Sleep(time.Second / 2)
			p("exit stage 1 worker with input", i)
			next(i)
		}(v)
	}
}

func stage2(wg *sync.WaitGroup, limit int, next func(int)) func(int) {
	g := NewSyncGroup(limit, wg)
	return func(v int) {
		g.Enter()
		go func(i int) {
			defer g.Done()
			p("start stage 2 worker with input", i)
			time.Sleep(time.Second / 2)
			p("exit stage 2 worker with input", i)
			next(i)
		}(v)
	}
}

type SyncGroup struct {
	l chan struct{}
	w *sync.WaitGroup
}

func NewSyncGroup(limit int, wg *sync.WaitGroup) *SyncGroup {
	return &SyncGroup{make(chan struct{}, limit), wg}
}

func (g *SyncGroup) Enter() {
	g.l <- struct{}{}
	g.w.Add(1)
}

func (g *SyncGroup) Done() {
	g.w.Done()
	<-g.l
}

func (g *SyncGroup) Wait() {
	g.w.Wait()
}

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