go-study/routine/pool.go

package routine

import (
	"fmt"
	"runtime/debug"
	"sync"
)

type Pool[T any] struct {
	taskQueue chan T
	taskFn    func(T)
	workers   int
	wg        sync.WaitGroup
}

func NewPool[T any](workers, capacity int, taskFn func(T)) *Pool[T] {
	pool := &Pool[T]{
		taskQueue: make(chan T, capacity),
		taskFn: func(t T) {
			defer func() {
				// 处理协程运行中出现panic的情况
				if r := recover(); r != nil {
					fmt.Printf("Panic: %v\n %s", r, string(debug.Stack()))
				}
			}()

			taskFn(t)
		},
		workers: workers,
	}
	pool.wg.Add(workers)

	return pool
}

// Start 启动任务
func (p *Pool[T]) Start() {
	for i := 0; i < p.workers; i++ {
		go func() {
			defer p.wg.Done()

			for {
				task, ok := <-p.taskQueue
				if !ok {
					return
				}

				p.taskFn(task)
			}
		}()
	}
}

// Push 提交任务
func (p *Pool[T]) Push(task T) {
	p.taskQueue <- task
}

// Wait 挂起当前协程
func (p *Pool[T]) Wait() {
	close(p.taskQueue)
	p.wg.Wait()
}