Golanggoroutine池实现与资源管理技巧

package main

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

// WorkerPool 定义了goroutine池的结构
type WorkerPool struct {
    taskQueue chan func() // 任务队列
    workerNum int         // 工作者数量
    wg        sync.WaitGroup // 用于等待所有任务完成
    quit      chan struct{} // 退出信号
}

// NewWorkerPool 创建一个新的goroutine池
func NewWorkerPool(workerNum int) *WorkerPool {
    if workerNum <= 0 {
        workerNum = 1 // 至少一个工作者
    }
    return &WorkerPool{
        taskQueue: make(chan func()),
        workerNum: workerNum,
        quit:      make(chan struct{}),
    }
}

// Start 启动goroutine池
func (p *WorkerPool) Start() {
    for i := 0; i < p.workerNum; i++ {
        p.wg.Add(1)
        go p.worker(i)
    }
}

// worker 是实际执行任务的goroutine
func (p *WorkerPool) worker(id int) {
    defer p.wg.Done()
    for {
        select {
        case task, ok := <-p.taskQueue:
            if !ok { // 任务队列已关闭
                fmt.Printf("Worker %d: Task queue closed, exiting.\n", id)
                return
            }
            fmt.Printf("Worker %d: Starting task.\n", id)
            task() // 执行任务
            fmt.Printf("Worker %d: Finished task.\n", id)
        case <-p.quit: // 收到退出信号
            fmt.Printf("Worker %d: Received quit signal, exiting.\n", id)
            return
        }
    }
}

// Submit 提交一个任务到goroutine池
func (p *WorkerPool) Submit(task func()) {
    p.taskQueue <- task
}

// Shutdown 关闭goroutine池，等待所有任务完成
func (p *WorkerPool) Shutdown() {
    close(p.taskQueue) // 关闭任务队列，通知所有worker不再接收新任务
    // 发送退出信号给所有worker，这在某些情况下可能需要，但通常关闭taskQueue就足够了
    // for i := 0; i < p.workerNum; i++ {
    //  p.quit <- struct{}{}
    // }
    p.wg.Wait() // 等待所有worker完成
    close(p.quit) // 关闭退出信号通道
    fmt.Println("Worker pool shutdown complete.")
}

func main() {
    pool := NewWorkerPool(3) // 创建一个包含3个worker的goroutine池
    pool.Start()

    // 提交一些任务
    for i := 0; i < 10; i++ {
        taskID := i
        pool.Submit(func() {
            time.Sleep(time.Duration(taskID%3+1) * time.Second) // 模拟耗时任务
            fmt.Printf("Task %d processed.\n", taskID)
        })
    }

    time.Sleep(2 * time.Second) // 给一些任务处理时间

    pool.Shutdown() // 关闭池
    fmt.Println("Main goroutine finished.")
}

Submit

Submit

Shutdown

WaitGroup

// 一个更健壮的WorkerPool结构示例，包含结果和错误处理
type Result struct {
    Value interface{}
    Err   error
}

type Task func(ctx context.Context) Result

type RobustWorkerPool struct {
    taskQueue   chan Task
    resultsChan chan Result // 用于收集任务结果
    workerNum   int
    wg          sync.WaitGroup
    ctx         context.Context
    cancel      context.CancelFunc
}

func NewRobustWorkerPool(workerNum int, resultBufferSize int) *RobustWorkerPool {
    ctx, cancel := context.WithCancel(context.Background())
    if workerNum <= 0 {
        workerNum = 1
    }
    if resultBufferSize < workerNum {
        resultBufferSize = workerNum // 至少能缓冲与worker数量相同的任务结果
    }
    return &RobustWorkerPool{
        taskQueue:   make(chan Task),
        resultsChan: make(chan Result, resultBufferSize),
        workerNum:   workerNum,
        ctx:         ctx,
        cancel:      cancel,
    }
}

func (p *RobustWorkerPool) Start() {
    for i := 0; i < p.workerNum; i++ {
        p.wg.Add(1)
        go p.worker(i)
    }
}

func (p *RobustWorkerPool) worker(id int) {
    defer p.wg.Done()
    for {
        select {
        case task, ok := <-p.taskQueue:
            if !ok {
                return // 任务队列已关闭
            }
            res := task(p.ctx) // 执行任务，传递上下文
            select {
            case p.resultsChan <- res: // 将结果发送到结果通道
            case <-p.ctx.Done(): // 如果池已关闭，则放弃结果
                fmt.Printf("Worker %d: Pool shutting down, discarding result.\n", id)
                return
            }
        case <-p.ctx.Done(): // 收到取消信号
            return
        }
    }
}

func (p *RobustWorkerPool) Submit(task Task) error {
    select {
    case p.taskQueue <- task:
        return nil
    case <-p.ctx.Done():
        return p.ctx.Err() // 池已关闭
    default: // 非阻塞提交，如果通道满则报错
        return fmt.Errorf("task queue is full")
    }
}

func (p *RobustWorkerPool) GetResults() <-chan Result {
    return p.resultsChan
}

func (p *RobustWorkerPool) Shutdown() {
    p.cancel()          // 发送取消信号给所有worker
    close(p.taskQueue) // 关闭任务队列，确保所有待处理任务被取出
    p.wg.Wait()         // 等待所有worker完成
    close(p.resultsChan) // 关闭结果通道
    fmt.Println("Robust Worker pool shutdown complete.")
}

// 示例用法
func mainRobustPool() {
    pool := NewRobustWorkerPool(2, 5) // 2个worker，结果通道缓冲5个
    pool.Start()

    // 提交一些任务
    for i := 0; i < 7; i++ { // 提交7个任务，但池只有2个worker
        taskID := i
        err := pool.Submit(func(ctx context.Context) Result {
            select {
            case <-ctx.Done():
                return Result{nil, fmt.Errorf("task %d cancelled", taskID)}
            case <-time.After(time.Duration(taskID%3+1) * time.Second): // 模拟耗时
                return Result{fmt.Sprintf("Processed %d", taskID), nil}
            }
        })
        if err != nil {
            fmt.Printf("Failed to submit task %d: %v\n", taskID, err)
        }
    }

    // 收集结果
    go func() {
        for res := range pool.GetResults() {
            if res.Err != nil {
                fmt.Printf("Task error: %v\n", res.Err)
            } else {
                fmt.Printf("Task result: %v\n", res.Value)
            }
        }
        fmt.Println("Result collector finished.")
    }()

    time.Sleep(5 * time.Second)
    pool.Shutdown()
    fmt.Println("Main robust pool goroutine finished.")
}

RobustWorkerPool

context.Context

resultsChan

Submit