ants 协程池的WaitGroup简单封装

283 阅读1分钟

概述

ants的Submit方法是会返回error的,此时如果传入了WaitGroup,就有死锁的风险

代码

package utils

import (
	"sync"

	"github.com/panjf2000/ants/v2"
	"k8s.io/klog"
)

type RoutinePool interface {
	Submit(task func()) error
	SubmitWithWg(task func(), wg *sync.WaitGroup) error
}

var (
	routinePoolSingle RoutinePool
	routinePoolOnce   sync.Once
	routinePoolSize   int = 10000
)

func GetRoutinePool() RoutinePool {
	routinePoolOnce.Do(func() {
		antsPool, err := ants.NewPool(routinePoolSize, ants.WithNonblocking(true))
		if err != nil {
			klog.Fatal(err)
			return
		}
		routinePoolSingle = newRoutinePool(antsPool)
	})
	return routinePoolSingle
}

type routinePool struct {
	pool *ants.Pool
}

func newRoutinePool(pool *ants.Pool) RoutinePool {
	return &routinePool{pool: pool}
}

func (p *routinePool) Submit(task func()) error {
	return p.pool.Submit(task)
}

func (p *routinePool) SubmitWithWg(task func(), wg *sync.WaitGroup) error {
	err := p.pool.Submit(task)
	if err != nil {
		wg.Done()
	}
	return err
}