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

package main

import (
	"fmt"
	"math/rand"
	"runtime"
	"sync"
	"time"

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

const totalTasks = 10

func main() {
	fmt.Println("Start to work. GOMAXPROCS =", runtime.GOMAXPROCS(0))

	ant, _ := ants.NewPool(3) // 协程池大小为3
	defer ant.Release()

	var wg sync.WaitGroup
	results := make([]int, totalTasks)
	for i := 0; i < totalTasks; i++ {
		idx := i // 不能在闭包内访问i,因为i一直在变,必须复制出一份idx
		input := i
		wg.Add(1)
		err := ant.Submit( // Submit 池满则阻塞
			func() {
				defer wg.Done()
				// 任务逻辑:
				rd := rand.New(rand.NewSource(time.Now().UnixNano()))
				time.Sleep(time.Duration(rd.Intn(1000)) * time.Millisecond)
				fmt.Println("task index:", input)
				results[idx] = input * input
			},
		)
		if err != nil {
			fmt.Println("Failed to submit task, err=", err)
			wg.Done() // 由于任务提交失败,需要手动还原wg
			break     // 或许大多数情况下是直接 return err
		}
	}

	wg.Wait()

	fmt.Println("Done. result =", results)
}
Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐