github.com/iron-io/functions@v0.0.0-20180820112432-d59d7d1c40b2/api/mqs/ironmq.go (about) 1 package mqs 2 3 import ( 4 "context" 5 "encoding/json" 6 "fmt" 7 "net/url" 8 "strconv" 9 "strings" 10 "sync" 11 12 "github.com/Sirupsen/logrus" 13 "github.com/iron-io/functions/api/models" 14 mq_config "github.com/iron-io/iron_go3/config" 15 ironmq "github.com/iron-io/iron_go3/mq" 16 ) 17 18 type assoc struct { 19 msgId string 20 reservationId string 21 } 22 23 type IronMQ struct { 24 queues []ironmq.Queue 25 // Protects the map 26 sync.Mutex 27 // job id to {msgid, reservationid} 28 msgAssoc map[string]*assoc 29 } 30 31 type IronMQConfig struct { 32 Token string `mapstructure:"token"` 33 ProjectId string `mapstructure:"project_id"` 34 Host string `mapstructure:"host"` 35 Scheme string `mapstructure:"scheme"` 36 Port uint16 `mapstructure:"port"` 37 QueuePrefix string `mapstructure:"queue_prefix"` 38 } 39 40 func NewIronMQ(url *url.URL) *IronMQ { 41 42 if url.User == nil || url.User.Username() == "" { 43 logrus.Fatal("IronMQ requires PROJECT_ID and TOKEN") 44 } 45 p, ok := url.User.Password() 46 if !ok { 47 logrus.Fatal("IronMQ requires PROJECT_ID and TOKEN") 48 } 49 settings := &mq_config.Settings{ 50 Token: p, 51 ProjectId: url.User.Username(), 52 Host: url.Host, 53 Scheme: "https", 54 } 55 56 if url.Scheme == "ironmq+http" { 57 settings.Scheme = "http" 58 } 59 60 parts := strings.Split(url.Host, ":") 61 if len(parts) > 1 { 62 settings.Host = parts[0] 63 p, err := strconv.Atoi(parts[1]) 64 if err != nil { 65 logrus.WithFields(logrus.Fields{"host_port": url.Host}).Fatal("Invalid host+port combination") 66 } 67 settings.Port = uint16(p) 68 } 69 70 var queueName string 71 if url.Path != "" { 72 queueName = url.Path 73 } else { 74 queueName = "titan" 75 } 76 mq := &IronMQ{ 77 queues: make([]ironmq.Queue, 3), 78 msgAssoc: make(map[string]*assoc), 79 } 80 81 // Check we can connect by trying to create one of the queues. Create is 82 // idempotent, so this is fine. 83 _, err := ironmq.ConfigCreateQueue(ironmq.QueueInfo{Name: fmt.Sprintf("%s_%d", queueName, 0)}, settings) 84 if err != nil { 85 logrus.WithError(err).Fatal("Could not connect to IronMQ") 86 } 87 88 for i := 0; i < 3; i++ { 89 mq.queues[i] = ironmq.ConfigNew(fmt.Sprintf("%s_%d", queueName, i), settings) 90 } 91 92 logrus.WithFields(logrus.Fields{"base_queue": queueName}).Info("IronMQ initialized") 93 return mq 94 } 95 96 func (mq *IronMQ) Push(ctx context.Context, job *models.Task) (*models.Task, error) { 97 if job.Priority == nil || *job.Priority < 0 || *job.Priority > 2 { 98 return nil, fmt.Errorf("IronMQ Push job %s: Bad priority", job.ID) 99 } 100 101 // Push the work onto the queue. 102 buf, err := json.Marshal(job) 103 if err != nil { 104 return nil, err 105 } 106 _, err = mq.queues[*job.Priority].PushMessage(ironmq.Message{Body: string(buf), Delay: int64(job.Delay)}) 107 return job, err 108 } 109 110 func (mq *IronMQ) Reserve(ctx context.Context) (*models.Task, error) { 111 var job models.Task 112 113 var messages []ironmq.Message 114 var err error 115 for i := 2; i >= 0; i-- { 116 messages, err = mq.queues[i].LongPoll(1, 60, 0 /* wait */, false /* delete */) 117 if err != nil { 118 // It is OK if the queue does not exist, it will be created when a message is queued. 119 if !strings.Contains(err.Error(), "404 Not Found") { 120 return nil, err 121 } 122 } 123 124 if len(messages) == 0 { 125 // Try next priority. 126 if i == 0 { 127 return nil, nil 128 } 129 continue 130 } 131 132 // Found a message! 133 break 134 } 135 136 message := messages[0] 137 if message.Body == "" { 138 return nil, nil 139 } 140 141 err = json.Unmarshal([]byte(message.Body), &job) 142 if err != nil { 143 return nil, err 144 } 145 mq.Lock() 146 mq.msgAssoc[job.ID] = &assoc{message.Id, message.ReservationId} 147 mq.Unlock() 148 return &job, nil 149 } 150 151 func (mq *IronMQ) Delete(ctx context.Context, job *models.Task) error { 152 if job.Priority == nil || *job.Priority < 0 || *job.Priority > 2 { 153 return fmt.Errorf("IronMQ Delete job %s: Bad priority", job.ID) 154 } 155 mq.Lock() 156 assoc, exists := mq.msgAssoc[job.ID] 157 delete(mq.msgAssoc, job.ID) 158 mq.Unlock() 159 160 if exists { 161 return mq.queues[*job.Priority].DeleteMessage(assoc.msgId, assoc.reservationId) 162 } 163 return nil 164 }