github.com/artisanhe/tools@v1.0.1-0.20210607022958-19a8fef2eb04/mq/job_queue.go (about)

     1  package mq
     2  
     3  import (
     4  	"fmt"
     5  	"runtime/debug"
     6  	"sync"
     7  
     8  	"github.com/sirupsen/logrus"
     9  )
    10  
    11  func NewJobQueue(broker Broker, backend Backend) *JobQueue {
    12  	return &JobQueue{
    13  		broker:  broker,
    14  		backend: backend,
    15  		jobs:    sync.Map{},
    16  	}
    17  }
    18  
    19  type JobQueue struct {
    20  	broker  Broker
    21  	backend Backend
    22  
    23  	channel string
    24  	jobs    sync.Map
    25  	worker  *Worker
    26  
    27  	receiver       func(*TaskStatus) error
    28  	receiverWorker *Worker
    29  }
    30  
    31  func (mq *JobQueue) ListChannel() ([]string, error) {
    32  	return mq.backend.ListChannel()
    33  }
    34  
    35  func (mq *JobQueue) ListSubject(channel string) ([]string, error) {
    36  	return mq.backend.ListSubject(channel)
    37  }
    38  
    39  func (mq *JobQueue) RegisterReceiver(receiver func(*TaskStatus) error) {
    40  	mq.receiver = receiver
    41  }
    42  
    43  func (mq *JobQueue) StartReceiver(numWorkers int) {
    44  	mq.receiverWorker = NewWorker(mq.receiverProcess, numWorkers)
    45  	mq.receiverWorker.Start()
    46  }
    47  
    48  func (mq *JobQueue) StopReceiver() {
    49  	if mq.receiverWorker != nil {
    50  		mq.receiverWorker.Stop()
    51  	}
    52  }
    53  
    54  func (mq *JobQueue) receiverProcess() error {
    55  	taskStatus, err := mq.backend.GetFeedback()
    56  	if err != nil || taskStatus == nil {
    57  		return nil
    58  	}
    59  
    60  	if taskStatus.Status == StatusCancelled {
    61  		if err := mq.backend.ClearCancellation(taskStatus.ID); err != nil {
    62  			logrus.Warnf("clear cancellation %s failed %s", taskStatus.ID, err)
    63  		}
    64  	}
    65  
    66  	if mq.receiver != nil {
    67  		if err := mq.receiver(taskStatus); err != nil {
    68  			return err
    69  		}
    70  	}
    71  
    72  	return nil
    73  }
    74  
    75  type Job = func(task *Task) (result interface{}, err error)
    76  
    77  func (mq *JobQueue) Register(subject string, job Job) {
    78  	mq.jobs.Store(subject, job)
    79  }
    80  
    81  func (mq *JobQueue) Cancel(id string) error {
    82  	logrus.Debugf("cancelling %s", id)
    83  	return mq.backend.Cancel(id)
    84  }
    85  
    86  func (mq *JobQueue) SendTask(task *Task) error {
    87  	return mq.broker.SendTask(task)
    88  }
    89  
    90  func (mq *JobQueue) Next(channel string, subject string, data []byte) (*Task, error) {
    91  	task := NewTask(channel, subject, data)
    92  	return task, mq.backend.FeedBack(task.Pending())
    93  }
    94  
    95  func (mq *JobQueue) Publish(channel string, subject string, data []byte) (*Task, error) {
    96  	task := NewTask(channel, subject, data)
    97  	return task, mq.SendTask(task)
    98  }
    99  
   100  func (mq *JobQueue) StartWorker(channel string, numWorkers int) {
   101  	subjects := make([]string, 0)
   102  	mq.jobs.Range(func(key, value interface{}) bool {
   103  		subjects = append(subjects, key.(string))
   104  		return true
   105  	})
   106  
   107  	mq.channel = channel
   108  
   109  	if err := mq.backend.RegisterChannel(mq.channel, subjects); err != nil {
   110  		logrus.Panic(err)
   111  	}
   112  
   113  	mq.worker = NewWorker(mq.jobProcess, numWorkers)
   114  	mq.worker.Start()
   115  }
   116  
   117  func (mq *JobQueue) StopWorker() {
   118  	if mq.worker != nil {
   119  		mq.worker.Stop()
   120  	}
   121  }
   122  
   123  func (mq *JobQueue) jobProcess() error {
   124  	task, err := mq.broker.GetTask(mq.channel)
   125  	if err != nil || task == nil {
   126  		return nil
   127  	}
   128  
   129  	defer func() {
   130  		if e := recover(); e != nil {
   131  			jobErr := fmt.Errorf("panic: %s; calltrace:%s", fmt.Sprint(e), string(debug.Stack()))
   132  			logrus.Warnf("%s.%s failed: %s", task.Subject, task.ID, jobErr)
   133  			if err := mq.backend.FeedBack(task.Failed(jobErr)); err != nil {
   134  				logrus.Errorf("feed back FAILED failed %s/%s.%s: %s", task.Channel, task.Subject, task.ID, err)
   135  			}
   136  		}
   137  	}()
   138  
   139  	if isCancelled, err := mq.backend.IsCancelled(task.ID); isCancelled {
   140  		if err := mq.backend.FeedBack(task.Cancelled()); err != nil {
   141  			logrus.Errorf("feed back CANCELLED failed %s/%s.%s: %s", task.Channel, task.Subject, task.ID, err)
   142  		}
   143  		return err
   144  	}
   145  
   146  	job, ok := mq.jobs.Load(task.Subject)
   147  	if !ok {
   148  		jobErr := fmt.Errorf("missing subject %s", task)
   149  		if err := mq.backend.FeedBack(task.Failed(jobErr)); err != nil {
   150  			logrus.Errorf("feed back FAILED failed %s/%s.%s: %s", task.Channel, task.Subject, task.ID, err)
   151  		}
   152  		return jobErr
   153  	}
   154  
   155  	if err := mq.backend.FeedBack(task.Running()); err != nil {
   156  		logrus.Errorf("feed back RUNNING failed %s/%s.%s: %s", task.Channel, task.Subject, task.ID, err)
   157  	}
   158  
   159  	jobResult, jobErr := job.(Job)(task)
   160  
   161  	if jobErr != nil {
   162  		logrus.Warnf("%s.%s failed: %s", task.Subject, task.ID, jobErr)
   163  		if err := mq.backend.FeedBack(task.Failed(jobErr)); err != nil {
   164  			logrus.Errorf("feed back FAILED failed %s/%s.%s: %s", task.Channel, task.Subject, task.ID, err)
   165  		}
   166  		return jobErr
   167  	}
   168  
   169  	if err := mq.backend.FeedBack(task.Success(jobResult)); err != nil {
   170  		logrus.Errorf("feed back SUCCESS failed %s/%s.%s: %s", task.Channel, task.Subject, task.ID, err)
   171  	}
   172  
   173  	return nil
   174  }