github.com/fawick/restic@v0.1.1-0.20171126184616-c02923fbfc79/internal/worker/pool.go (about)

     1  package worker
     2  
     3  import "context"
     4  
     5  // Job is one unit of work. It is given to a Func, and the returned result and
     6  // error are stored in Result and Error.
     7  type Job struct {
     8  	Data   interface{}
     9  	Result interface{}
    10  	Error  error
    11  }
    12  
    13  // Func does the actual work within a Pool.
    14  type Func func(ctx context.Context, job Job) (result interface{}, err error)
    15  
    16  // Pool implements a worker pool.
    17  type Pool struct {
    18  	f     Func
    19  	jobCh <-chan Job
    20  	resCh chan<- Job
    21  
    22  	numWorkers     int
    23  	workersExit    chan struct{}
    24  	allWorkersDone chan struct{}
    25  }
    26  
    27  // New returns a new worker pool with n goroutines, each running the function
    28  // f. The workers are started immediately.
    29  func New(ctx context.Context, n int, f Func, jobChan <-chan Job, resultChan chan<- Job) *Pool {
    30  	p := &Pool{
    31  		f:              f,
    32  		workersExit:    make(chan struct{}),
    33  		allWorkersDone: make(chan struct{}),
    34  		numWorkers:     n,
    35  		jobCh:          jobChan,
    36  		resCh:          resultChan,
    37  	}
    38  
    39  	for i := 0; i < n; i++ {
    40  		go p.runWorker(ctx, i)
    41  	}
    42  
    43  	go p.waitForExit()
    44  
    45  	return p
    46  }
    47  
    48  // waitForExit receives from p.workersExit until all worker functions have
    49  // exited, then closes the result channel.
    50  func (p *Pool) waitForExit() {
    51  	n := p.numWorkers
    52  	for n > 0 {
    53  		<-p.workersExit
    54  		n--
    55  	}
    56  	close(p.allWorkersDone)
    57  	close(p.resCh)
    58  }
    59  
    60  // runWorker runs a worker function.
    61  func (p *Pool) runWorker(ctx context.Context, numWorker int) {
    62  	defer func() {
    63  		p.workersExit <- struct{}{}
    64  	}()
    65  
    66  	var (
    67  		// enable the input channel when starting up a new goroutine
    68  		inCh = p.jobCh
    69  		// but do not enable the output channel until we have a result
    70  		outCh chan<- Job
    71  
    72  		job Job
    73  		ok  bool
    74  	)
    75  
    76  	for {
    77  		select {
    78  		case <-ctx.Done():
    79  			return
    80  
    81  		case job, ok = <-inCh:
    82  			if !ok {
    83  				return
    84  			}
    85  
    86  			job.Result, job.Error = p.f(ctx, job)
    87  			inCh = nil
    88  			outCh = p.resCh
    89  
    90  		case outCh <- job:
    91  			outCh = nil
    92  			inCh = p.jobCh
    93  		}
    94  	}
    95  }
    96  
    97  // Wait waits for all worker goroutines to terminate, afterwards the output
    98  // channel is closed.
    99  func (p *Pool) Wait() {
   100  	<-p.allWorkersDone
   101  }