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

package main

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

var wg = new(sync.WaitGroup)

var consumer = int32(0)

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

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

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

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

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

	wg.Add(1)
	go f1(ch)

	wg.Add(2)
	atomic.AddInt32(&consumer, 2)
	go f2(ch, ch2)
	go f2(ch, ch2)

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

	wg.Wait()
}