golang-study/base/channel/example/main.go

package main

import (
	"fmt"
	"runtime"
	"sync"
	"sync/atomic"
)

var wg = new(sync.WaitGroup)

var consumer = int32(0)

func f1(ch chan int) {
	defer wg.Done()
	for i := 0; i < 1_000_000; i++ {
		ch <- i
	}

	// 使用完必须 close 不然将导致死锁
	close(ch)
}

func f2(in chan int, out chan int) {
	defer wg.Done()
	for i := range in {
		out <- i
	}

	if atomic.AddInt32(&consumer, -1) <= 0 {
		// 使用完必须 close 不然将导致死锁
		close(out)
	}
}

func main() {
	ch := make(chan int, 50)
	ch2 := make(chan int, 5)

	wg.Add(1)
	go f1(ch)

	wg.Add(runtime.NumCPU())
	atomic.AddInt32(&consumer, int32(runtime.NumCPU()))

	for i := 0; i < runtime.NumCPU(); i++ {
		go f2(ch, ch2)
	}

	for i := range ch2 {
		fmt.Println(i)
	}

	wg.Wait()
}