go实现简单协程池

代码

package main

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

type Pool interface {
	Run(task *Task) error
}

type Task struct {
	// 任务名
	Name string
	// 任务回调函数
	Handler func(v ...interface{})
	// 任务回调函数的参数
	Params []interface{}
}

type TaskPool struct {
	// 任务
	Work chan *Task
	// 协程池数量
	Capacity chan struct{}
}

func NewTaskPool(capacity int) (*TaskPool, error) {
	if capacity <= 0 {
		return nil, errors.New("capacity less than 0")
	}
	return &TaskPool{
		Work:     make(chan *Task),              // 无缓冲管道
		Capacity: make(chan struct{}, capacity), // 缓冲管道，协程池大小
	}, nil
}

var _ Pool = (*TaskPool)(nil)

/*
	第一次调用添加任务，一定会走第二个case 分之
	当协程池大于2时，此时第二次传入task，如果第一个协程还在运行中，就一定会走第二个case 重新创建一个协程执行task
	如果传入的任务数大于设定的协程数，并且此时所有的任务都还在运行中，那此时传入task,这两个case都不会命中,会一直
	阻塞直到有任务执行完成
*/

func (tp *TaskPool) Run(task *Task) error {
	select {
	case tp.Work <- task:
	case tp.Capacity <- struct{}{}:
		go tp.worker(task)
	}
	return nil
}

func (tp *TaskPool) worker(task *Task) {
	// todo 由gc 回收？
	//defer func() { <-tp.Capacity }()
	// for 协程复用
	for {
		task.Handler(task.Params...)
		task = <-tp.Work
	}
}

func (tp *TaskPool) Close() {
	for {
		if len(tp.Work) == 0 {
			close(tp.Work)
			return
		}
	}
}

func main() {
	// 新建协程池
	taskPool, err := NewTaskPool(6)
	if err != nil {
		fmt.Println(err)
		return
	}
	// 提交20 个 任务
	wg := &sync.WaitGroup{}

	for i := 0; i < 20; i++ {
		wg.Add(1)
		_ = taskPool.Run(&Task{
			Name: fmt.Sprintf("Demo-%d", i),
			Handler: func(v ...interface{}) {
				defer wg.Done()
				// 测试分批
				time.Sleep(500 * time.Millisecond)
				fmt.Printf("Hello, %s worker\n", v[0])
			},
			Params: []interface{}{fmt.Sprintf("name-%d", i)},
		})
	}
	wg.Wait()
	defer taskPool.Close() // 安全关闭协程池
}

问题

无法知道 worker 与 pool 的状态

worker 数量不足无法动态扩增

worker 数量过多无法自动缩减
....