实现一个简单的线程池

goroutine-pool.go

package main

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

var jobs = make(chan int, 10)
var results = make(chan int, 10)

func task(number int) int {
	sum := 0
	for i := 0; i < number; i++ {
		sum += i
	}
	// 用 sleep 来模拟耗时的计算
	time.Sleep(10 * time.Second)
	return sum
}

func worker(wg *sync.WaitGroup) {
	for job := range jobs {
		results <- task(job)
	}
	wg.Done()
}

// 创建 numberOfWorker 个 goroutine 的线程池
// 创建 goroutine 之前调用 wg.Add(1) 来增加计数器，并将 wg 的地址传给 worker
// 使用 wg.Wait() 等待所有的 goroutine 执行完毕，关闭 results channel，保证不再有 goroutine 能写入数据
func createWorkerPool(numberOfWorker int) {
	var wg sync.WaitGroup
	for i := 0; i < numberOfWorker; i++ {
		wg.Add(1)
		go worker(&wg)
	}
	wg.Wait()
	close(results)
}

// 向线程池分配任务
func allocate(numberOfJobs int) {
	for i := 0; i < numberOfJobs; i++ {
		jobs <- i
	}
	close(jobs)
}

func result(done chan bool) {
	for re := range results {
		fmt.Printf("Result is %d\n", re)
	}
	done <- true
}

func main() {
	startTime := time.Now()
	done := make(chan bool)

	go allocate(10)
	go result(done)
	createWorkerPool(10)

	endTime := time.Now()
	runtime := endTime.Sub(startTime)

	fmt.Println("total time: ", runtime)
}