[Go] 工作池

源自：juejin.cn/post/699169…

代码部分

// ExecutionFn job执行的方法
type ExecutionFn func(ctx context.Context, args interface{}) (interface{}, error)

// Job 工作任务
type Job struct {
   Descriptor JobDescriptor
   ExecFn     ExecutionFn
   Args       interface{}
}

// Result 结果
type Result struct {
   Value      interface{}
   Err        error
   Descriptor JobDescriptor
}

// execute 处理工作 返回结果使用result
func (j Job) execute(ctx context.Context) Result {
   value, err := j.ExecFn(ctx, j.Args)
   // 出错了
   if err != nil {
      return Result{
         Err:        err,
         Descriptor: j.Descriptor,
      }
   }
   return Result{
      Value:      value,
      Descriptor: j.Descriptor,
   }
}

// worker 执行任务
func worker(ctx context.Context, wg *sync.WaitGroup, jobs <-chan Job, results chan<- Result) {
   defer wg.Done()
   for {
      select {
      case job, ok := <-jobs:
         if !ok {
            return
         }
         results <- job.execute(ctx)
         //log.Println(job.Descriptor.ID)
      case <-ctx.Done():
         log.Printf("cancel worker,error detail: %v \n", ctx.Err())
         results <- Result{
            Err: ctx.Err(),
         }
         return
      }
   }
}

type WorkPool struct {
   workersCount int           // worker数量
   jobs         chan Job      // 工作集
   results      chan Result   // 结果集
   Done         chan struct{} // 结束标志
}

func New(wc int) WorkPool {
   return WorkPool{
      workersCount: wc,
      jobs:         make(chan Job, wc),
      results:      make(chan Result, wc),
      Done:         make(chan struct{}),
   }
}

// Run 运行
func (wp WorkPool) Run(ctx context.Context) {
   var wg sync.WaitGroup
   for i := 0; i < wp.workersCount; i++ {
      wg.Add(1)
      go worker(ctx, &wg, wp.jobs, wp.results)
   }
   wg.Wait()
   close(wp.Done)
   close(wp.results)
}

// ShutDown 关闭工作池
func (wp WorkPool) ShutDown() {
   wp.Done <- struct{}{}
   //资源释放
   defer close(wp.Done)
   defer close(wp.results)
   defer close(wp.jobs)
}

// Results 返回结果
func (wp WorkPool) Results() <-chan Result {
   return wp.results
}

// GenerateFromSlice 生成Job
func (wp WorkPool) GenerateFromSlice(jobsBulk []Job) {
   for i, _ := range jobsBulk {
      wp.jobs <- jobsBulk[i]
   }
   //close(wp.jobs)
   return
}

func (wp WorkPool) AddJob(job Job) {
   wp.jobs <- job
   return
}

测试，将结果集打印出来

func TestWorkPool(t *testing.T) {
   wp := New(10)
   ctx, cancel := context.WithCancel(context.TODO())
   defer cancel()

   go wp.GenerateFromSlice(testJob())
   //go wp.GenerateFromSlice(testJob())
   //go wp.GenerateFromSlice(testJob())
   go wp.AddJob(generateJob())

   go wp.Run(ctx)

   // 十秒钟后关闭
   go func() {
      time.AfterFunc(10*time.Second, wp.ShutDown)
   }()
   for {
      select {
      case r, ok := <-wp.Results():
         if !ok {
            continue
         }
         fmt.Println(r.Value)
      case <-wp.Done:
         return
      default:

      }
   }
}