go协程池实现原理小结

type Worker interface {
	Task()
}

type Worker interface {
	Task()
}

type Pool struct {
	// 使用无缓冲通道实现协程池
	work chan Worker
	// 辅助计数器，用于协程池同步
	wg   sync.WaitGroup
}

// New 创建一个工作池
func New(maxGoroutine int) *Pool {
	p := Pool{
		work: make(chan Worker),
	}
	p.wg.Add(maxGoroutine)
	for i := 0; i < maxGoroutine; i++ {
		go func() {
			// 一直循环取任务，直到work被关闭为止并且通道中的任务执行完毕为止
			for w := range p.work {
				w.Task()
			}
			// 退出时候，减少一个计数器
			p.wg.Done()
		}()
	}
	return &p
}

// Run 将任务放入工作池
func (p *Pool) Run(w Worker) {
	p.work <- w
}

// Shutdown 等待所有goroutine完成工作
func (p *Pool) Shutdown() {
	// 关闭work所有协程完成任务之后会退出for循环
	close(p.work)
	p.wg.Wait()
}

package work

import "sync"

// Worker interface 必须满足worker的要求才能使用工作池
type Worker interface {
	Task()
}

// Pool 提供一个goroutine池，这个池可以完成任何已提交的woker任务
type Pool struct {
	work chan Worker
	wg   sync.WaitGroup
}

// New 创建一个工作池
func New(maxGoroutine int) *Pool {
	p := Pool{
		work: make(chan Worker),
	}
	p.wg.Add(maxGoroutine)
	for i := 0; i < maxGoroutine; i++ {
		go func() {
			// 一直循环取任务，直到work被关闭为止
			for w := range p.work {
				w.Task()
			}
			// 退出时候，减少一个计数器
			p.wg.Done()
		}()
	}
	return &p
}

// Run 将任务放入工作池
func (p *Pool) Run(w Worker) {
	p.work <- w
}

// Shutdown 等待所有goroutine完成工作
func (p *Pool) Shutdown() {
	// 关闭work所有协程完成任务之后会退出for循环
	close(p.work)
	p.wg.Wait()
}

// This sample program demonstrates how to use the work package
// to use a pool of goroutines to get work done.
package main

import (
	"log"
	"sync"
	"time"

	"work"
)

// names provides a set of names to display.
var names = []string{
	"steve",
	"bob",
	"mary",
	"therese",
	"jason",
}

// namePrinter provides special support for printing names.
type namePrinter struct {
	name string
}

// Task implements the Worker interface.
func (m *namePrinter) Task() {
	log.Println(m.name)
	time.Sleep(time.Second)
}

// main is the entry point for all Go programs.
func main() {
	// Create a work pool with 2 goroutines.
	p := work.New(2)

	var wg sync.WaitGroup
	wg.Add(100 * len(names))

	for i := 0; i < 100; i++ {
		// Iterate over the slice of names.
		for _, name := range names {
			// Create a namePrinter and provide the
			// specific name.
			np := namePrinter{
				name: name,
			}

			go func() {
				// Submit the task to be worked on. When RunTask
				// returns we know it is being handled.
				p.Run(&np)
				wg.Done()
			}()
		}
	}

	wg.Wait()

	// Shutdown the work pool and wait for all existing work
	// to be completed.
	p.Shutdown()
}