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 數量過多無法自動縮減
....