使用阻塞型channel控制协程并发数

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64

package main

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

type WaitGroup struct {
	workChan chan int
	wg       sync.WaitGroup
}

func NewPool(limit int) *WaitGroup {
	ch := make(chan int, limit)

	return &WaitGroup{
		workChan: ch,
		wg:       sync.WaitGroup{},
	}
}

func (wg *WaitGroup) add(num int) {
	for i := 0; i < num; i++ {
		wg.workChan <- i
		wg.wg.Add(1)
	}
}

func (wg *WaitGroup) done() {
LOOP:
	for {
		select {
		case <-wg.workChan:
			break LOOP
		}
	}
	wg.wg.Done()
}

func (wg *WaitGroup) wait() {
	wg.wg.Wait()
}

func main() {
	worker := NewPool(4)
	for i := 0; i < 64; i++ {
		worker.add(1)
		go run(i, worker)
	}

	fmt.Println("waiting...")
	worker.wait()
	fmt.Println("down...")
}

func run(i int, wg *WaitGroup) {
	defer wg.done()

	fmt.Println(time.Now().Format("2006-01-02 15:04:05"), "output: ", i)
	time.Sleep(time.Second * 1)
	fmt.Println(time.Now().Format("2006-01-02 15:04:05"), "output: ", i, "done")

}