go项目源码分析与使用--协程池pool

1.项目介绍

地址: https://github.com/gobwas/ws-examples/blob/master/src/gopool/pool.go
只介绍一个代码文件，它简洁优雅的实现了一个goroutine的pool

2.关键源码

pool.go

// Package gopool contains tools for goroutine reuse.
// It is implemented only for examples of github.com/gobwas/ws usage.
package gopool

import (
	"fmt"
	"time"
)

// ErrScheduleTimeout returned by Pool to indicate that there no free
// goroutines during some period of time.
var ErrScheduleTimeout = fmt.Errorf("schedule error: timed out")

// Pool contains logic of goroutine reuse.
type Pool struct {
	sem  chan struct{}  //可以理解为坑位或者令牌
	work chan func()   //任务队列
}

// NewPool creates new goroutine pool with given size. It also creates a work
// queue of given size. Finally, it spawns given amount of goroutines
// immediately.
func NewPool(size, queue, spawn int) *Pool {
	if spawn <= 0 && queue > 0 {
		panic("dead queue configuration detected")
	}
	if spawn > size {
		panic("spawn > workers")
	}
	p := &Pool{
		sem:  make(chan struct{}, size),   //有size个坑位
		work: make(chan func(), queue),
	}
	//每个工作goroutine都要先抢占一个坑位才能工作，开始先启动spawn个工作goroutine
	for i := 0; i < spawn; i++ {
		p.sem <- struct{}{}  
		go p.worker(func() {})
	}

	return p
}

// Schedule schedules task to be executed over pool's workers.
//任务调度
func (p *Pool) Schedule(task func()) {
	p.schedule(task, nil)
}

// ScheduleTimeout schedules task to be executed over pool's workers.
// It returns ErrScheduleTimeout when no free workers met during given timeout.
func (p *Pool) ScheduleTimeout(timeout time.Duration, task func()) error {
	return p.schedule(task, time.After(timeout))
}

//分四种情况讨论
//p.work没满，p.sem没满:  任务可能投放到p.work中，也可能起一个新的goroutine处理
//p.work没满，p.sem已满:  任务投放p.work中
//p.work已满，p.sem没满:  任务被新的goroutine处理
//p.work已满，p.sem已满:  任务投放到p.work被阻塞，如果有timeout，则可能超时返回
func (p *Pool) schedule(task func(), timeout <-chan time.Time) error {
	select {
	case <-timeout:
		return ErrScheduleTimeout
	case p.work <- task:
		return nil
	case p.sem <- struct{}{}:
		go p.worker(task)
		return nil
	}
}

func (p *Pool) worker(task func()) {
	//退出才会让出坑位
	defer func() { <-p.sem }()

	task()

	for task := range p.work {
		task()
	}
}

3.例子

package main

import (
	"fmt"
	"net/http"
	"runtime"
	"strconv"
	"sync"
	"time"

	"github.com/gobwas/ws-examples/src/gopool"
)

func main() {
	//pool大小10个，任务队列50
	workers, queue := 10, 50
	pool := gopool.NewPool(workers, queue, 1)
	go monitor()
	var wg sync.WaitGroup
	time.Sleep(5 * time.Second)
	//在这里已经有7个goroutine，有1个是pool预启动的工作goroutine
	//后面不断发送任务到队列里去，goroutine的个数会逐渐增长到16个就停止增长

	for i := 0; i < 100; i++ {
		time.Sleep(1 * time.Second) //方便查看goroutine个数增加的过程
		n := i + 1
		wg.Add(1)
		fn := func() {
			defer wg.Done()
			fmt.Println("task number:", n)
			time.Sleep(5 * time.Second)
		}
		pool.Schedule(fn)
	}
	wg.Wait()
}

//访问http://localhost:6060/goroutines 可实时监控goroutine的个数
func monitor() {
	http.HandleFunc("/goroutines", func(w http.ResponseWriter, r *http.Request) {
		num := strconv.FormatInt(int64(runtime.NumGoroutine()), 10)
		w.Write([]byte(num))
	})
	http.ListenAndServe("localhost:6060", nil)
	fmt.Println("goroutine stats and pprof listen on 6060")
}