自己动手写了一个协程池

Go语言虽然有着高效的GMP调度模型，理论上支持成千上万的goroutine，但是goroutine过多，对调度，GC以及系统内存都会造成压力，这样会使我们的服务性能不升反降。常用做法可以用池化技术，构造一个协程池，把进程中的协程控制在一定的数量，防止系统中goroutine过多，影响服务性能

1. 协程池模型

协程池简单理解就是有一个池子一样的东西，里面装着固定数量的goroutine，当有一个任务到来的时候，会将这个任务交给池子里的一个空闲的goroutine去处理，如果池子里没有空闲的goroutine了，任务就会阻塞等待。所以协程池有三个角色Worker，Task，Pool

1.1 属性定义

Worker ：用于执行任务的goroutine
Task : 具体的任务
Pool : 池子

下面看一下各个角色的定义

Task ：内部有一个函数成员，表示这个task具体的执行逻辑，Task代码定义如下：

go 复制代码

type Task struct {
    f func() error  // 具体的执行逻辑
}

Pool ：有两个成员，Capacity表示池子里的worker的数量，即工作的goroutine的数量，JobCh 表示任务队列用于存放任务，goroutine从这个JobCh 获取任务并执行任务逻辑，Pool代码定义如下：

go 复制代码

type Pool struct {
    RunningWorkers int64
    Capacity       int64     // goroutine数量
    JobCh          chan *Task // 用于worker取任务
    sync.Mutex
}

worker ：执行任务单元，简单理解就是干活的goroutine，这个worker其实只做一件事情，就是不断的从任务队列里面取任务执行，而worker的数量就是协程池里协程的数量，由Pool的参数WorkerNum指定

go 复制代码

// p为Pool对象指针
for task := range p.JobCh{
    do ...      
}

1.2 方法定义

go 复制代码

func NewTask(funcArg func() error) *Task

NewTask用于创建一个任务，参数是一个函数，返回值是一个Task类型

go 复制代码

func NewPool(Capacity int, taskNum int) *Pool

NewPool返回一个协程数量固定为workerNum的协程池对象指针，其任务队列的长度为taskNum

接下来主要介绍协程池的各个方法

scss 复制代码

func (p *Pool) AddTask(task *Task)

AddTask方法是往协程池添加任务，如果当前运行着的worker数量小于协程池worker容量，则立即启动一个协程worker来处理任务，否则将任务添加到任务队列

scss 复制代码

func (p *Pool) Run()

将协程池跑起来，启动一个worker来处理任务

协程池处理任务流程图如下：

2. 协程池实现

go 复制代码

package main

import (
    "fmt"
    "sync"
    "sync/atomic"
    "time"
)

type Task struct {
    f func() error // 具体的任务逻辑
}

func NewTask(funcArg func() error) *Task {
    return &Task{
       f: funcArg,
    }
}

type Pool struct {
    RunningWorkers int64      // 运行着的worker数量
    Capacity       int64      // 协程池worker容量
    JobCh          chan *Task // 用于worker取任务
    sync.Mutex
}

func NewPool(capacity int64, taskNum int) *Pool {
    return &Pool{
       Capacity: capacity,
       JobCh:    make(chan *Task, taskNum),
    }
}

func (p *Pool) GetCap() int64 {
    return p.Capacity
}

func (p *Pool) incRunning() { // runningWorkers + 1
    atomic.AddInt64(&p.RunningWorkers, 1)
}

func (p *Pool) decRunning() { // runningWorkers - 1
    atomic.AddInt64(&p.RunningWorkers, -1)
}

func (p *Pool) GetRunningWorkers() int64 {
    return atomic.LoadInt64(&p.RunningWorkers)
}

func (p *Pool) run() {
    p.incRunning()
    go func() {
       defer func() {
          p.decRunning()
       }()
       for task := range p.JobCh {
          task.f()
       }
    }()
}

// AddTask 往协程池添加任务
func (p *Pool) AddTask(task *Task) {

    // 加锁防止启动多个 worker
    p.Lock()
    defer p.Unlock()

    if p.GetRunningWorkers() < p.GetCap() { // 如果任务池满, 则不再创建 worker
       // 创建启动一个 worker
       p.run()
    }

    // 将任务推入队列, 等待消费
    p.JobCh <- task
}

func main() {
    // 创建任务池
    pool := NewPool(3, 10)

    for i := 0; i < 20; i++ {
       // 任务放入池中
       pool.AddTask(NewTask(func() error {
          fmt.Printf("I am Task\n")
          return nil
       }))
    }

    time.Sleep(1e9) // 等待执行
}

运行结果：

css 复制代码

I am Task
I am Task
I am Task
I am Task
I am Task
I am Task
I am Task
I am Task
I am Task
I am Task
I am Task
I am Task
I am Task
I am Task
I am Task
I am Task
I am Task
I am Task
I am Task
I am Task

程序创建了一个WorkerNum为3，任务队列长度为10的协程池，往里面添加了20个任务，一直只有3个worker在做任务，起到了控制goroutine数量的作用

交流学习

如果您觉得文章有帮助，请帮忙转发给更多好友，或关注公众号：IT杨秀才，持续更新更多硬核文章，一起聊聊互联网那些事儿!