go实现协程池

时间:2024-03-25 18:41:58
package main
 
import (
    "fmt"
    "sync"
    "sync/atomic"
    "time"
)
 
type WorkerPool struct {
    workers int
    jobs    chan func()
    wg      sync.WaitGroup
}
 
func NewWorkerPool(workers int) *WorkerPool {
    pool := &WorkerPool{
        workers: workers,
        jobs:    make(chan func()),
    }
    pool.startWorkers()
    return pool
}
 
func (p *WorkerPool) startWorkers() {
    for i := 0; i < p.workers; i++ {
        go func() {
            for job := range p.jobs {
                job()
            }
        }()
    }
}
 
func (p *WorkerPool) Do(job func()) {
    p.wg.Add(1)
    go func() {
        defer p.wg.Done()
        p.jobs <- job
    }()
}
 
func (p *WorkerPool) Wait() {
    p.wg.Wait()
    close(p.jobs)
}
 
func main() {
    pool := NewWorkerPool(10)
    var count int32
 
    for i := 0; i < 100; i++ {
        pool.Do(func() {
            time.Sleep(10 * time.Millisecond)
            atomic.AddInt32(&count, 1)
        })
    }
 
    pool.Wait()
    fmt.Printf("Count is %d\n", count)
}