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 }