github.com/pingcap/br@v5.3.0-alpha.0.20220125034240-ec59c7b6ce30+incompatible/pkg/utils/worker.go (about)

     1  // Copyright 2020 PingCAP, Inc. Licensed under Apache-2.0.
     2  
     3  package utils
     4  
     5  import (
     6  	"github.com/pingcap/log"
     7  	"go.uber.org/zap"
     8  	"golang.org/x/sync/errgroup"
     9  )
    10  
    11  // WorkerPool contains a pool of workers.
    12  type WorkerPool struct {
    13  	limit   uint
    14  	workers chan *Worker
    15  	name    string
    16  }
    17  
    18  // Worker identified by ID.
    19  type Worker struct {
    20  	ID uint64
    21  }
    22  
    23  type taskFunc func()
    24  
    25  type identifiedTaskFunc func(uint64)
    26  
    27  // NewWorkerPool returns a WorkPool.
    28  func NewWorkerPool(limit uint, name string) *WorkerPool {
    29  	workers := make(chan *Worker, limit)
    30  	for i := uint(0); i < limit; i++ {
    31  		workers <- &Worker{ID: uint64(i + 1)}
    32  	}
    33  	return &WorkerPool{
    34  		limit:   limit,
    35  		workers: workers,
    36  		name:    name,
    37  	}
    38  }
    39  
    40  // Apply executes a task.
    41  func (pool *WorkerPool) Apply(fn taskFunc) {
    42  	worker := pool.ApplyWorker()
    43  	go func() {
    44  		defer pool.RecycleWorker(worker)
    45  		fn()
    46  	}()
    47  }
    48  
    49  // ApplyWithID execute a task and provides it with the worker ID.
    50  func (pool *WorkerPool) ApplyWithID(fn identifiedTaskFunc) {
    51  	worker := pool.ApplyWorker()
    52  	go func() {
    53  		defer pool.RecycleWorker(worker)
    54  		fn(worker.ID)
    55  	}()
    56  }
    57  
    58  // ApplyOnErrorGroup executes a task in an errorgroup.
    59  func (pool *WorkerPool) ApplyOnErrorGroup(eg *errgroup.Group, fn func() error) {
    60  	worker := pool.ApplyWorker()
    61  	eg.Go(func() error {
    62  		defer pool.RecycleWorker(worker)
    63  		return fn()
    64  	})
    65  }
    66  
    67  // ApplyWithIDInErrorGroup executes a task in an errorgroup and provides it with the worker ID.
    68  func (pool *WorkerPool) ApplyWithIDInErrorGroup(eg *errgroup.Group, fn func(id uint64) error) {
    69  	worker := pool.ApplyWorker()
    70  	eg.Go(func() error {
    71  		defer pool.RecycleWorker(worker)
    72  		return fn(worker.ID)
    73  	})
    74  }
    75  
    76  // ApplyWorker apply a worker.
    77  func (pool *WorkerPool) ApplyWorker() *Worker {
    78  	var worker *Worker
    79  	select {
    80  	case worker = <-pool.workers:
    81  	default:
    82  		log.Debug("wait for workers", zap.String("pool", pool.name))
    83  		worker = <-pool.workers
    84  	}
    85  	return worker
    86  }
    87  
    88  // RecycleWorker recycle a worker.
    89  func (pool *WorkerPool) RecycleWorker(worker *Worker) {
    90  	if worker == nil {
    91  		panic("invalid restore worker")
    92  	}
    93  	pool.workers <- worker
    94  }
    95  
    96  // HasWorker checks if the pool has unallocated workers.
    97  func (pool *WorkerPool) HasWorker() bool {
    98  	return len(pool.workers) > 0
    99  }