Golang协程池ants使用笔记

最近工程中遇到goroutine滥用导致的bug，采用了ants协程池来解决。
github-ants官方源码和使用说明。
记录一个例子

package main

import (
	"fmt"
	"runtime"
	"sync"
	"time"

	"github.com/panjf2000/ants/v2"
)

// 模拟一个任务：求数字x的平方
func square(x int, ch chan int) {
	time.Sleep(1 * time.Second) // 模拟任务
	fmt.Printf("Calculate: %d x %d = %d\n", x, x, x*x)
	ch <- x * x // 用channel存放返回值
	close(ch)
}

func test_ants() {
	fmt.Println("CPU cores: ", runtime.GOMAXPROCS(0))
	numTask := 10
	poolSize := 3

	// 申请一个协程池对象
	pool, _ := ants.NewPoolWithFunc(poolSize, func(i interface{}) {
		arr := i.([]interface{}) // 先转为数组，再挨个取出元素填入到任务函数的参数
		square(arr[0].(int), arr[1].(chan int))
	})
	// 关闭协程池。关闭后，继续提交任务会返回错误，但运行中的任务(goroutine)会一直执行完
	defer pool.Release()

	// 要执行n次任务，每个任务的返回值用channel接收
	chanReceivers := make([]chan int, numTask)
	for i := 0; i < numTask; i++ {
		chanReceivers[i] = make(chan int) // 注意这里用的是无缓冲区的channel
	}

	// 由于ants提交任务是阻塞的，所以放在routine中执行
	go func(numTask, poolSize int) {
		for i := 0; i < numTask; i++ {
			err := pool.Invoke([]interface{}{i, chanReceivers[i]}) // 提交任务，超过容量时会阻塞在这
			if err != nil {
				fmt.Println("[Error] Failed to invoke task", i, err)
				break // 如果需求是遇到一次提交任务失败，就不再发新任务了，那这里直接break就好了
			} else {
				fmt.Println("Submitted task ", i)
			}
		}
	}(numTask, poolSize)

	// 读取channel中的返回值
	for _, ch := range chanReceivers {
		ret := <-ch
		fmt.Println("Received:", ret)
		// 如果通过ch拿到的值发现这个任务出了错，可以直接在这里return。
		// 随后会触发关闭pool，致使子routine不能提交新任务。
		// 但已经提交了的任务无法终止（除非通过context等方式给个信号）
		if ret == 4 { // 故意遇到一个错误
			fmt.Println("test_ants occurred error ret==4, so terminate.")
			return
		}
	}

	fmt.Println("Done")
}

func main() {
	// 练习使用 wait group 优雅地等待一个 go routine 执行完毕
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		test_ants()
		wg.Done()
	}()
	wg.Wait()
}