github.com/tamaxcode/worker@v1.0.0/collector.go (about)

     1  package worker
     2  
     3  import (
     4  	"context"
     5  	"runtime"
     6  	"sync"
     7  )
     8  
     9  // Config represent configuration setting for collector to send received task to worker
    10  type Config struct {
    11  	// Number of worker spawn to handle concurrent work. Default to "runtime.NumCPU"
    12  	// giving value <= 0 will fallback to default value
    13  	NoOfWorkers int
    14  	Handler     func(data interface{}) bool
    15  }
    16  
    17  // Collector will communicate with outer layer to receive handler for worker to process task
    18  type Collector struct {
    19  	workers []*Worker
    20  	work    chan<- interface{}
    21  	wg      *sync.WaitGroup
    22  	ctx     context.Context
    23  
    24  	stopChan <-chan bool
    25  	stopper  func()
    26  }
    27  
    28  // NewCollector ...
    29  func NewCollector(config Config) *Collector {
    30  
    31  	ctx, canc := context.WithCancel(context.Background())
    32  
    33  	noOfWorkers := config.NoOfWorkers
    34  	if noOfWorkers <= 0 {
    35  		noOfWorkers = runtime.NumCPU()
    36  	}
    37  
    38  	workers := make([]*Worker, noOfWorkers)
    39  
    40  	inputChannel := make(chan interface{}, noOfWorkers)
    41  	wg := sync.WaitGroup{}
    42  
    43  	stopChan := make(chan bool)
    44  
    45  	for i := 0; i < noOfWorkers; i++ {
    46  		w := Worker{
    47  			id:  i,
    48  			ctx: ctx,
    49  
    50  			workHandler: config.Handler,
    51  			workChannel: inputChannel,
    52  			wg:          &wg,
    53  
    54  			stopChannel: stopChan,
    55  		}
    56  		w.Start()
    57  
    58  		workers[i] = &w
    59  	}
    60  
    61  	return &Collector{
    62  		work:    inputChannel,
    63  		workers: workers,
    64  		wg:      &wg,
    65  
    66  		ctx:      ctx,
    67  		stopChan: stopChan,
    68  		stopper:  canc,
    69  	}
    70  }
    71  
    72  // Add ...
    73  // Example
    74  /*
    75  	Collector.Add(somedata)
    76  */
    77  func (c *Collector) Add(data interface{}) {
    78  	c.wg.Add(1)
    79  
    80  	go func() {
    81  		c.work <- data
    82  	}()
    83  }
    84  
    85  // Wait until all task completed
    86  func (c *Collector) Wait() {
    87  	done := make(chan bool)
    88  	go func() {
    89  		c.wg.Wait()
    90  		done <- true
    91  	}()
    92  
    93  	select {
    94  
    95  	case <-c.ctx.Done():
    96  		// context cancelled
    97  	case <-c.stopChan:
    98  		// worker abort
    99  		c.stopper()
   100  	case <-done:
   101  		// wait group complete
   102  	}
   103  }
   104  
   105  // Stop all ongoing task
   106  func (c *Collector) Stop() {
   107  	c.stopper()
   108  }