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

     1  package worker
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  )
     7  
     8  // Work will bring task handler for worker to call
     9  type Work struct {
    10  	Handler WorkHandler
    11  }
    12  
    13  // WorkHandler function that will be called by the worker to process incoming work data
    14  // will terminate all pending work if return false
    15  // example
    16  /*
    17  	handler := func (data interface{}) bool {
    18  		str, _ := data.(string)
    19  
    20  		fmt.Println(str)
    21  
    22  		return true
    23  	}
    24  */
    25  type WorkHandler func(data interface{}) bool
    26  
    27  // Worker ...
    28  type Worker struct {
    29  	id  int
    30  	ctx context.Context
    31  
    32  	workHandler func(data interface{}) bool
    33  	workChannel <-chan interface{}
    34  	wg          *sync.WaitGroup
    35  
    36  	stopChannel chan<- bool
    37  }
    38  
    39  // Start ...
    40  func (w *Worker) Start() {
    41  
    42  	go func() {
    43  		for {
    44  			var (
    45  				workData interface{}
    46  			)
    47  
    48  			select {
    49  			case <-w.ctx.Done():
    50  				// context canceled, stop worker
    51  				return
    52  			case work, ok := <-w.workChannel:
    53  				if !ok {
    54  					// channel closed, stop worker
    55  					return
    56  				}
    57  				workData = work
    58  			}
    59  
    60  			cont := true
    61  			if w.workChannel != nil {
    62  				cont = w.workHandler(workData)
    63  			}
    64  
    65  			if cont == false {
    66  				w.stopChannel <- true
    67  			}
    68  
    69  			w.wg.Done()
    70  		}
    71  	}()
    72  }