动手实现一个协程池(goroutine pool)

403 阅读2分钟

实现一个协程池(goroutine pool)可以有效管理goroutine的数量,防止因创建过多goroutine而耗尽系统资源。协程池的基本思路是预先创建一定数量的goroutine,这些goroutine从任务队列中取出任务执行,执行完毕后不退出而是继续等待新的任务。这样可以复用goroutine,避免频繁创建和销毁的开销。下面是一个简单的协程池实现示例:

定义协程池结构

首先定义协程池的结构,包括任务队列、控制信号等。

type Task func()

type GoroutinePool struct {
    TaskChannel   chan Task
    MaxWorkers    int
    StopChannel   chan bool
}

func NewGoroutinePool(maxWorkers int) *GoroutinePool {
    pool := &GoroutinePool{
        TaskChannel:   make(chan Task),
        MaxWorkers:    maxWorkers,
        StopChannel:   make(chan bool),
    }
    return pool
}

启动和停止协程池

接下来实现启动和停止协程池的方法。启动时,创建指定数量的goroutine等待任务;停止时,关闭所有goroutine。

func (p *GoroutinePool) Start() {
    for i := 0; i < p.MaxWorkers; i++ {
        go p.worker()
    }
}

func (p *GoroutinePool) worker() {
    for {
        select {
        case task := <-p.TaskChannel:
            task()
        case <-p.StopChannel:
            return
        }
    }
}

func (p *GoroutinePool) Stop() {
    for i := 0; i < p.MaxWorkers; i++ {
        p.StopChannel <- true
    }
}

提交任务到协程池

定义一个方法,允许外部提交任务到协程池。

func (p *GoroutinePool) Submit(task Task) {
    p.TaskChannel <- task
}

使用示例

以下是如何使用上述协程池的简单示例:

package main

import (
    "fmt"
    "time"
)

func main() {
    pool := NewGoroutinePool(5) // 创建一个最大容量为5的协程池
    pool.Start()

    for i := 0; i < 10; i++ {
        count := i
        pool.Submit(func() {
            fmt.Println("Executing task", count)
            time.Sleep(2 * time.Second) // 假设每个任务需要2秒钟
        })
    }

    time.Sleep(5 * time.Second) // 等待一些任务完成
    pool.Stop() // 停止所有协程
}

在这个例子中,我们创建了一个最大容量为5的协程池,并提交了10个任务。每个任务简单地打印一条消息并休眠2秒。通过调用Start()启动协程池,它会启动5个goroutine去执行提交的任务。Submit()方法用于提交新的任务到协程池。最后,通过调用Stop()停止所有的goroutine。

这个简单的协程池实现演示了协程池的基本概念,实际使用时可能需要根据具体需求进行调整和优化,例如增加错误处理、支持动态调整协程数量等。