Go 协程池 Gopool VS ants 原理解析

写过高并发的都知道，控制协程数量是问题的关键，如何高效利用协程，本文将介绍gopool和ants两个广泛应用的协程池，通过本文你可以了解到：

1. 实现原理
2. 使用方法
3. 区别

背景

虽然通过go func()即可轻量级实现并发，但如果通过for循环创建go
func()，会出现成千上万个协程（1个协程2KB，过多协程会达到GB级别）导致系统资源耗尽，为了保护系统资源，往往需要控制协程数量。

协程池的核心思想是限制并发执行的 goroutine 数量，从而有效管理资源使用和避免过度并发带来的问题。

使用原生Channel实现协程池

Channel 支持多个协程之间的数据通信，并且可以通过有限长度的channel当做协程池，最大长度即为协程池最大容量。

核心原理

关键点1：定义协程池结构和构造函数

首先，定义一个协程池的结构体，它包含一个 channel 用于控制并发数量，以及一个 WaitGroup 用于等待所有任务完成。

package main

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

type GoroutinePool struct {
   
    maxGoroutines int
    pool          chan struct{
   }
    wg            sync.WaitGroup
}

func NewGoroutinePool(maxGoroutines int) *GoroutinePool {
   
    return &GoroutinePool{
   
        maxGoroutines: maxGoroutines,
        pool:          make(chan struct{
   }, maxGoroutines),
    }
}

关键点2：实现任务提交和执行的方法

为协程池添加方法以提交和执行任务。使用 channel 来控制同时运行的 goroutine 数量。

func (p *GoroutinePool) Submit(task func()) {
   
    p.wg.Add(1)
    p.pool <- struct{
   }{
   } // 获取令牌，如果池已满，这里会阻塞

    go func() {
   
        defer p.wg.Done()
        defer func() {
    <-p.pool }() // 释放令牌

        task() // 执行任务
    }()
}

func (p *GoroutinePool) Wait() {
   
    p.wg.Wait()
}

举个🌰

func main() {
   
    pool := NewGoroutinePool(5) // 创建一个大小为5的协程池

    for i := 0; i < 10; i++ {
   
        count := i
        pool.Submit(func() {
   
            fmt.Printf("Running task %d, timeNow: %v\n", count, time.Now())
            time.Sleep(1 * time.Second) // 模拟耗时任务
        })
    }

    pool.Wait() // 等待所有任务完成
    fmt.Println("All tasks completed.")
}

// 输出
Running task 1, timeNow: 2024-08-14 15:46:19.937473 +0800 CST m=+0.000572251
Running task 2, timeNow: 2024-08-14 15:46:19.937135 +0800 CST m=+0.000234542
Running task 3, timeNow: 2024-08-14 15:46:19.937344 +0800 CST m=+0.000443501
Running task 0, timeNow: 2024-08-14 15:46:19.937135 +0800 CST m=+0.000234167
Running task 4, timeNow: 2024-08-14 15:46:19.937104 +0800 CST m=+0.000203709
Running task 9, timeNow: 2024-08-14 15:46:20.938929 +0800 CST m=+1.002042292
Running task 6, timeNow: 2024-08-14 15:46:20.938935 +0800 CST m=+1.002048334
Running task 8, timeNow: 2024-08-14 15:46:20.938965 +0800 CST m=+1.002078501
Running task 7, timeNow: 2024-08-14 15:46:20.938963 +0800 CST m=+1.002076542
Running task 5, timeNow: 2024-08-14 15:46:20.938948 +0800 CST m=+1.002061167
All tasks completed.

可以看到，因为协程池大小为5，所以前5个任务都在0s执行，后5个在1s执行。

说实话，原生的Channel实现，个人感觉已经能够完全覆盖日常使用，但是公司和业界还是推出了很多协程池的实现，存在即合理，我们来学习一下~

使用gopkg/gopool实现协程池

这个包在字节内部已经广泛应用，咨询作者后，才发现这个包已经开源啦
https://github.com/bytedance/gopkg/tree/main/util/gopool

核心原理

这三个“池”纯属个人理解，抽象了两个池的概念
三个“池”的关系

实现流程

黄色 为新任务到来业务处理流程

协程池结构体

给用户侧暴露的协程池结构体为pool，存储了名称，最大容量，任务链表，正在进行的任务数量。通过NewPool来进行创建。通过pool.Go()方法进行添加事件处理

type pool struct {
   
    // pool 的名字，打 metrics 和打 log 时用到
    name string
    // pool 的容量，也就是最大的真正在工作的 goroutine 的数量
    // 为了性能考虑，可能会有小误差
    cap int32
    // 配置信息
    config *Config
    // 任务链表
    taskHead  *task
    taskTail  *task
    taskLock  sync.Mutex
    taskCount int32

    // 记录正在运行的 worker 数量
    workerCount int32

    // 用来标记是否关闭
    closed int32

    // worker panic 的时候会调用这个方法
    panicHandler func(context.Context, interface{
   })
}

// name 必须是不含有空格的，只能含有字母、数字和下划线，否则 metrics 会失败
func NewPool(name string, cap int32, config *Config) Pool {
   
    p := &pool{
   
       name:   name,
       cap:    cap,
       config: config,
    }
    return p
}

// 添加事件处理
func (p *pool) Go(f func()) {
   
    p.CtxGo(context.Background(), f)  // 在下文分析
}

关键点1：任务队列池

taskPool，存储对象为*task

用于存放待执行任务的队列，如果当前任务数<=最大容量，直接新建工作协程执行；如果当前任务数量>最大容量，新的任务会被链表链接到上一个任务尾部，不新建工作协程。这样所有的任务通过链表链接在一起，形成了任务队列池。

var taskPool sync.Pool

func init() {
   
    taskPool.New = newTask
}

type task struct {
   
    ctx context.Context
    f   func()
    next *task
}

func newTask() interface{
   } {
   
    return &task{
   }
}

// 核心代码-将新任务放到任务队列池并执行
func (p *pool) CtxGo(ctx context.Context, f func()) {
   
    t := taskPool.Get().(*task)
    t.ctx = ctx
    t.f = f
    p.taskLock.Lock()
    // 链接节点
    if p.taskHead == nil {
   
       p.taskHead = t
       p.taskTail = t
    } else {
   
       p.taskTail.next = t
       p.taskTail = t
    }
    p.taskLock.Unlock()
    atomic.AddInt32(&p.taskCount, 1)
    // 如果 pool 已经被关闭了，就 panic
    if atomic.LoadInt32(&p.closed) == 1 {
   
       panic("use closed pool")
    }
    // 判断条件满足以下两个条件：
    // 1. task 数量大于阈值
    // 2. 目前的 worker 数量小于上限 p.cap
    // 或者目前没有 worker
    if (atomic.LoadInt32(&p.taskCount) >= p.config.ScaleThreshold && p.WorkerCount() < atomic.LoadInt32(&p.cap)) || p.WorkerCount() == 0 {
   
       p.