github.com/MeteorsLiu/simpleMQ@v1.0.3/worker/worker.go (about)

     1  package worker
     2  
     3  import (
     4  	"log"
     5  	"sync"
     6  	"sync/atomic"
     7  	"time"
     8  
     9  	"github.com/MeteorsLiu/simpleMQ/queue"
    10  )
    11  
    12  type Worker struct {
    13  	workerQueue queue.Queue
    14  	pollMap     *queue.PollTask
    15  	sem         chan struct{}
    16  	working     atomic.Int64
    17  	enablePoll  bool
    18  	unlimited   bool
    19  }
    20  
    21  func NewWorker(n, spwan int, q queue.Queue, enablePoll ...bool) *Worker {
    22  	w := &Worker{
    23  		workerQueue: q,
    24  		unlimited:   spwan <= 0 || n <= 0,
    25  	}
    26  	if n > 0 {
    27  		w.sem = make(chan struct{}, n)
    28  	}
    29  	if len(enablePoll) > 0 && enablePoll[0] {
    30  		w.enablePoll = true
    31  		w.pollMap = queue.NewPoll()
    32  	}
    33  	for i := 0; i < spwan; i++ {
    34  		w.sem <- struct{}{}
    35  		go w.Run()
    36  	}
    37  
    38  	return w
    39  }
    40  
    41  func (w *Worker) handleTask(task queue.Task) {
    42  	idx := w.working.Add(1)
    43  	defer w.working.Add(-1)
    44  	var err error
    45  	for {
    46  		err = task.Do()
    47  		if err == nil ||
    48  			err == queue.ErrTaskStopped ||
    49  			err == queue.ErrFailReachLimits ||
    50  			err == queue.ErrQueueClosed ||
    51  			!task.IsRunUntilSuccess() {
    52  			break
    53  		}
    54  		log.Printf("Worker ID: %d Execute Task: %s Fail: %v",
    55  			idx,
    56  			task.ID(),
    57  			err,
    58  		)
    59  		// interrupted by fail fast
    60  		if task.IsDone() {
    61  			return
    62  		}
    63  		log.Printf("Task: %s is going to re-run", task.ID())
    64  		// workerQueue can be nil when worker is in unlimited mode.
    65  		if !w.unlimited && !w.workerQueue.IsClosed() {
    66  			w.workerQueue.Publish(task)
    67  			return
    68  		}
    69  	}
    70  
    71  	if err != nil && err != queue.ErrTaskStopped {
    72  		if err == queue.ErrQueueClosed {
    73  			task.Interrupt()
    74  		} else {
    75  			task.Stop()
    76  		}
    77  	}
    78  }
    79  
    80  func (w *Worker) Run() {
    81  	defer func() { <-w.sem }()
    82  
    83  	for {
    84  		task, err := w.workerQueue.Pop()
    85  		if err != nil {
    86  			return
    87  		}
    88  		if task.IsDone() {
    89  			continue
    90  		}
    91  		w.handleTask(task)
    92  	}
    93  }
    94  
    95  func (w *Worker) PublishSync(task queue.Task, callback ...queue.Finalizer) error {
    96  	if w.enablePoll {
    97  		w.pollMap.Register(task)
    98  	}
    99  	task.OnDone(callback...)
   100  	w.handleTask(task)
   101  	return nil
   102  }
   103  
   104  func (w *Worker) PublishSyncTimeout(task queue.Task, timeout time.Duration, callback ...queue.Finalizer) error {
   105  	if w.enablePoll {
   106  		w.pollMap.Register(task)
   107  	}
   108  	task.OnDone(callback...)
   109  	timer := time.AfterFunc(timeout, func() {
   110  		task.Stop()
   111  	})
   112  	defer timer.Stop()
   113  	go w.handleTask(task)
   114  
   115  	task.Wait()
   116  
   117  	return nil
   118  }
   119  
   120  func (w *Worker) Publish(task queue.Task, callback ...queue.Finalizer) bool {
   121  	if w.enablePoll {
   122  		w.pollMap.Register(task)
   123  	}
   124  
   125  	if w.unlimited {
   126  		task.OnDone(callback...)
   127  		go w.handleTask(task)
   128  		return true
   129  	}
   130  
   131  	if w.IsBusy() {
   132  		select {
   133  		case w.sem <- struct{}{}:
   134  			go w.Run()
   135  		default:
   136  		}
   137  	}
   138  	// all the callback function should be registered
   139  	// before it starts to run!
   140  
   141  	if len(callback) > 0 {
   142  		task.OnDone(callback...)
   143  	}
   144  	return w.workerQueue.Publish(task)
   145  }
   146  
   147  func (w *Worker) Stop() {
   148  	w.Wait(5 * time.Second)
   149  	if w.workerQueue != nil {
   150  		w.workerQueue.Close()
   151  	}
   152  }
   153  
   154  func (w *Worker) KillTask(id string) error {
   155  	return w.pollMap.Kill(id)
   156  }
   157  
   158  func (w *Worker) Working() int {
   159  	return int(w.working.Load())
   160  }
   161  
   162  func (w *Worker) SetQueue(q queue.Queue) {
   163  	w.workerQueue.Close()
   164  	w.workerQueue = q
   165  }
   166  
   167  func (w *Worker) IsBusy() bool {
   168  	return w.workerQueue.Free() < cap(w.sem) || w.workerQueue.Len() > w.Working()
   169  }
   170  
   171  func (w *Worker) Wait(timeout ...time.Duration) {
   172  	if !w.enablePoll {
   173  		return
   174  	}
   175  	var wg sync.WaitGroup
   176  	w.pollMap.ForEach(func(_ string, p queue.Pollable) bool {
   177  		wg.Add(1)
   178  		go func(task queue.Task) {
   179  			defer wg.Done()
   180  			if len(timeout) > 0 {
   181  				timer := time.AfterFunc(timeout[0], func() {
   182  					task.Interrupt()
   183  				})
   184  				defer timer.Stop()
   185  			}
   186  			task.Wait()
   187  		}(p)
   188  		return true
   189  	})
   190  	wg.Wait()
   191  }