Golang 协程池

67 阅读1分钟

初始版本

func PoolDemo() {
    numTasks := 100
    concurrency := 5
    var wg sync.WaitGroup
    taskCh := make(chan int)

    for i := 0; i < concurrency; i++ {
       wg.Add(1)
       go func() {
          defer wg.Done()
          for taskId := range taskCh {
             PrintHello(taskId)
          }
       }()
    }

    for i := 0; i <= numTasks; i++ {
       taskCh <- i
    }
    close(taskCh)
    wg.Wait()
}

完善版本

package main

import "sync"
# 需要发布执行的任务
type ITask interface {
    Run()
}
# 开始执行任务 发布任务 和 等待执行关闭通道
type pool interface {
    Start()
    Schedule(task ITask)
    WaitAndStop()
}
# 任务数量 任务队列 任务计数器
type gPool struct {
    workers int
    tasks   chan ITask
    wg      sync.WaitGroup
}
# 开始执行任务
func (p *gPool) Start() {
    for i := 0; i < p.workers; i++ {
      # 任务计数
       p.wg.Add(1)
       go func() {
          defer p.wg.Done()
          for task := range p.tasks {
             # 执行任务
             task.Run()
          }
       }()
    }
}

// 任务发布
func (p *gPool) Schedule(task ITask) {
    p.tasks <- task
}
# 等待任务完成 和 关闭通道
func (p *gPool) WaitAndStop() {
    close(p.tasks)
   defer p.wg.Wait()
}
# 协程池构造函数
func NewPool(workers int) gPool {
    return gPool{
       workers: workers,
       tasks:   make(chan ITask, workers),
       wg:      sync.WaitGroup{},
    }
}