github.com/godevsig/adaptiveservice@v0.9.23/workerpool.go (about)

     1  package adaptiveservice
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  )
     7  
     8  // workerPool is a pool of workers.
     9  type workerPool struct {
    10  	wg sync.WaitGroup
    11  	sync.Mutex
    12  	workers []*chan struct{}
    13  	cnt     int32
    14  }
    15  
    16  // newWorkerPool creates a new workerPool.
    17  func newWorkerPool() *workerPool {
    18  	return &workerPool{}
    19  }
    20  
    21  type status interface {
    22  	idle()
    23  	working()
    24  }
    25  
    26  func (wp *workerPool) idle() {
    27  	atomic.AddInt32(&wp.cnt, 1)
    28  }
    29  func (wp *workerPool) working() {
    30  	atomic.AddInt32(&wp.cnt, -1)
    31  }
    32  
    33  // worker is a function with a cancel channel it should check for exit.
    34  type worker func(done <-chan struct{}, st status)
    35  
    36  // addWorker adds a worker into the workerPool.
    37  func (wp *workerPool) addWorker(w worker) {
    38  	done := make(chan struct{})
    39  	wp.Lock()
    40  	defer wp.Unlock()
    41  	added := false
    42  	for i, pc := range wp.workers {
    43  		if *pc == nil {
    44  			wp.workers[i] = &done
    45  			added = true
    46  			break
    47  		}
    48  	}
    49  	if !added {
    50  		wp.workers = append(wp.workers, &done)
    51  	}
    52  	atomic.AddInt32(&wp.cnt, 1)
    53  	wp.wg.Add(1)
    54  	go func() {
    55  		defer func() {
    56  			if done != nil {
    57  				close(done)
    58  				done = nil
    59  			}
    60  			atomic.AddInt32(&wp.cnt, -1)
    61  			wp.wg.Done()
    62  		}()
    63  		w(done, wp)
    64  	}()
    65  }
    66  
    67  // rmWorker tries to remove a random running worker from the workerPool.
    68  // It may do nothing if all workers in the workerPool have been removed
    69  // or exited.
    70  func (wp *workerPool) rmWorker() {
    71  	wp.Lock()
    72  	defer wp.Unlock()
    73  	var pdone *chan struct{}
    74  	for _, pc := range wp.workers {
    75  		if *pc != nil {
    76  			pdone = pc
    77  			break
    78  		}
    79  	}
    80  	if pdone == nil {
    81  		return
    82  	}
    83  	done := *pdone
    84  	*pdone = nil
    85  	close(done)
    86  }
    87  
    88  // close closes and waits all the workers in the workerPool to exit.
    89  func (wp *workerPool) close() {
    90  	wp.Lock()
    91  	for _, pc := range wp.workers {
    92  		if *pc != nil {
    93  			done := *pc
    94  			*pc = nil
    95  			close(done)
    96  		}
    97  	}
    98  	wp.Unlock()
    99  	wp.wg.Wait()
   100  }
   101  
   102  // len returns the number of idle workers in the workerPool.
   103  func (wp *workerPool) len() int {
   104  	cnt := atomic.LoadInt32(&wp.cnt)
   105  	return int(cnt)
   106  }