base-golang/pkg/ants/pool.go

package ants

import (
	"go.uber.org/zap"
)

var Pool GoroutinePool

var _ GoroutinePool = (*goroutinePool)(nil)

type GoroutinePool interface {
	Submit(task func())
	Stop()
	Size() int
}

type goroutinePool struct {
	pool *ants.Pool
	zap  *zap.Logger
	size int
}

type poolLogger struct {
	zap *zap.Logger
}

func (l *poolLogger) Printf(format string, args ...interface{}) {
	l.zap.Sugar().Infof(format, args)
}

func Init(zap *zap.Logger, size int) error {
	l := &poolLogger{zap: zap}
	options := ants.Options{
		Nonblocking: true,
		PanicHandler: func(i interface{}) {
			zap.Sugar().Panic(i)
		},
		Logger: l,
	}

	p, err := ants.NewPool(size, ants.WithOptions(options))
	if err != nil {
		return err
	}

	Pool = &goroutinePool{
		pool: p,
		zap:  zap,
		size: size,
	}

	return nil
}

func (p *goroutinePool) Submit(task func()) {
	if p.pool.IsClosed() {
		panic("协程池异常关闭")
	}
	err := p.pool.Submit(task)
	if err == ants.ErrPoolOverload {
		p.size = p.size + p.size/10
		p.pool.Tune(p.size)
		p.Submit(task)
	}
}

func (p *goroutinePool) Size() int {
	return p.size
}

func (p *goroutinePool) Stop() {
	p.pool.Release()
}