Golang协程池ants使用笔记

雪的期许2024-01-27 8:52

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

记录一个例子

go 复制代码

package main

import (
	"fmt"
	"runtime"

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

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

func main() {
	fmt.Println(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))
	})
	defer pool.Release() // 记得关闭协程池

	// 要执行n次任务，每个任务的返回值用channle接收
	chanReceivers := make([]chan int, numTask)
	for i := 0; i < numTask; i++ {
		chanReceivers[i] = make(chan int) // 注意这里用的是无缓冲区的channel
	}
	// 由于ants提交任务是阻塞的，所以放在routine中执行
	go func() {
		for i := 0; i < numTask; i++ {
			err := pool.Invoke([]interface{}{i, chanReceivers[i]}) // 提交任务，超过容量时会阻塞在这
			if err != nil {
				fmt.Println("[Error] Failed to invoke task: ", err)
			}
			fmt.Println("Submitted task ", i)
		}
	}()

	// 读取channel中的返回值
	for _, ch := range chanReceivers {
		// 如果通过ch发现某一任务出了错，可以直接在这里return，子协程会都被终止，避免继续做无用功。
		fmt.Println("Received:", <-ch)
	}

	fmt.Println("Done")
}