golang的协程池本应该是这样的

85 阅读1分钟

看过了一下github上 star比较高的协程池实现,还有字节开源的实现,完全是 java/c++之类的外行实现思路

协程/线程池,最基本的元件 就是 队列 + 协程/线程,M:N模型

这两个组件在go里边天生就有啊,为什么再搞一套task queue呢?

控制队列容量:make(chan, cap) 第二参数就可以

想要控制协程/线程数量,再辅助一个chan 就可以了,

代码实现如下,几十行搞定:

我把它放到github上 gopool 喜欢的老铁可以给个star

type GoPool struct {
    queueCap     int
    queueLen     atomic.Int32
    panicHandler func(any)

    workerSem chan struct{}
    queue     chan func()
}

// NewGoPool provite fixed number of goroutines, reusable. M:N model
//
// M: the number of reusable goroutines,
// N: the capacity for asynchronous task queue.
func NewGoPool(sizeM, preSpawn, queueCap int, panicHandler func(any)) *GoPool {
    if preSpawn <= 0 && queueCap > 0 {
        panic("GoPool: dead queue")
    }
    if preSpawn > sizeM {
        preSpawn = sizeM
    }
    p := &GoPool{
        queueCap:     queueCap,
        panicHandler: panicHandler,
        workerSem:    make(chan struct{}, sizeM),
        queue:        make(chan func(), queueCap),
    }
    for i := 0; i < preSpawn; i++ { // pre spawn
        p.workerSem <- struct{}{}
        go p.worker(func() {})
    }
    return p
}
// QueueFree returns (capacity of task-queue - length of task-queue)
func (p *GoPool) QueueFree() int {
    return p.queueCap - int(p.queueLen.Load())
}

// Go submits a task to this pool.
func (p *GoPool) Go(task func()) {
    select {
    case p.queue <- task:
        p.queueLen.Add(1)
    case p.workerSem <- struct{}{}:
        go p.worker(task)
    }
}

func (p *GoPool) worker(task func()) {
    defer func() {
        <-p.workerSem
        if e := recover(); e != nil {
            if p.panicHandler != nil {
                p.panicHandler(e)
            }
        }
    }()

    for {
        task()
        task = <-p.queue
        p.queueLen.Add(-1)
    }
}