面试题：Go固定worker工作池弹性伸缩策略之中等难度：实现基本的工作池

Question

请用Go语言实现一个固定worker数量的工作池，其中每个worker可以处理来自任务队列的任务。假设任务是简单的函数调用，例如 `func(task int) int`。要求代码能够正确处理任务的分发与执行，并展示如何启动和停止工作池。

Answer 1

package main

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

// 定义任务类型
type Task func(int) int

// 工作池结构体
type WorkerPool struct {
    Workers    int
    TaskQueue  chan Task
    WaitGroup  sync.WaitGroup
    StopSignal chan struct{}
}

// 创建新的工作池
func NewWorkerPool(workers int, taskQueueSize int) *WorkerPool {
    return &WorkerPool{
        Workers:    workers,
        TaskQueue:  make(chan Task, taskQueueSize),
        StopSignal: make(chan struct{}),
    }
}

// 启动工作池
func (wp *WorkerPool) Start() {
    for i := 0; i < wp.Workers; i++ {
        wp.WaitGroup.Add(1)
        go func() {
            defer wp.WaitGroup.Done()
            for {
                select {
                case task, ok := <-wp.TaskQueue:
                    if!ok {
                        return
                    }
                    result := task(1) // 这里简单传入1作为参数示例
                    fmt.Printf("Worker processed task, result: %d\n", result)
                case <-wp.StopSignal:
                    return
                }
            }
        }()
    }
}

// 停止工作池
func (wp *WorkerPool) Stop() {
    close(wp.StopSignal)
    close(wp.TaskQueue)
    wp.WaitGroup.Wait()
}

func main() {
    // 创建工作池，2个worker，任务队列大小为10
    pool := NewWorkerPool(2, 10)
    // 启动工作池
    pool.Start()

    // 添加任务
    task1 := func(task int) int {
        return task * 2
    }
    pool.TaskQueue <- task1

    task2 := func(task int) int {
        return task + 3
    }
    pool.TaskQueue <- task2

    // 等待一段时间，确保任务处理完成
    time.Sleep(2 * time.Second)

    // 停止工作池
    pool.Stop()
}

上述代码实现了一个固定worker数量的工作池：

1. 定义任务类型：Task 是一个函数类型，接受一个 int 参数并返回一个 int。
2. 工作池结构体：WorkerPool 包含了worker数量、任务队列、等待组和停止信号通道。
3. 创建工作池：NewWorkerPool 函数初始化工作池。
4. 启动工作池：Start 方法启动指定数量的worker，每个worker从任务队列中取出任务并执行。
5. 停止工作池：Stop 方法发送停止信号并等待所有worker完成任务。
6. 主函数：在 main 函数中，创建工作池，启动它，添加任务，等待一段时间后停止工作池。