forked from travis-ci/worker
-
Notifications
You must be signed in to change notification settings - Fork 0
/
amqp_job_queue.go
128 lines (107 loc) · 3.69 KB
/
amqp_job_queue.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
package worker
import (
"encoding/json"
"github.com/bitly/go-simplejson"
"github.com/streadway/amqp"
"github.com/travis-ci/worker/backend"
"github.com/travis-ci/worker/context"
gocontext "golang.org/x/net/context"
)
// AMQPJobQueue is a JobQueue that uses AMQP
type AMQPJobQueue struct {
conn *amqp.Connection
queue string
DefaultLanguage, DefaultDist, DefaultGroup, DefaultOS string
}
// NewAMQPJobQueue creates a AMQPJobQueue backed by the given AMQP connections and
// connects to the AMQP queue with the given name. The queue will be declared
// in AMQP when this function is called, so an error could be raised if the
// queue already exists, but with different attributes than we expect.
func NewAMQPJobQueue(conn *amqp.Connection, queue string) (*AMQPJobQueue, error) {
channel, err := conn.Channel()
if err != nil {
return nil, err
}
_, err = channel.QueueDeclare(queue, true, false, false, false, nil)
if err != nil {
return nil, err
}
err = channel.Close()
if err != nil {
return nil, err
}
return &AMQPJobQueue{
conn: conn,
queue: queue,
}, nil
}
// Jobs creates a new consumer on the queue, and returns three channels. The
// first channel gets sent every BuildJob that we receive from AMQP. The
// stopChan is a channel that can be closed in order to stop the consumer.
func (q *AMQPJobQueue) Jobs(ctx gocontext.Context) (outChan <-chan Job, err error) {
channel, err := q.conn.Channel()
if err != nil {
return
}
err = channel.Qos(1, 0, false)
if err != nil {
return
}
deliveries, err := channel.Consume(q.queue, "build-job-consumer", false, false, false, false, nil)
if err != nil {
return
}
buildJobChan := make(chan Job)
outChan = buildJobChan
go func() {
for delivery := range deliveries {
buildJob := &amqpJob{
payload: &JobPayload{},
startAttributes: &backend.StartAttributes{},
}
startAttrs := &jobPayloadStartAttrs{Config: &backend.StartAttributes{}}
err := json.Unmarshal(delivery.Body, buildJob.payload)
if err != nil {
context.LoggerFromContext(ctx).WithField("err", err).Error("payload JSON parse error, attempting to nack delivery")
err := delivery.Ack(false)
if err != nil {
context.LoggerFromContext(ctx).WithField("err", err).WithField("delivery", delivery).Error("couldn't nack delivery")
}
continue
}
err = json.Unmarshal(delivery.Body, &startAttrs)
if err != nil {
context.LoggerFromContext(ctx).WithField("err", err).Error("start attributes JSON parse error, attempting to nack delivery")
err := delivery.Ack(false)
if err != nil {
context.LoggerFromContext(ctx).WithField("err", err).WithField("delivery", delivery).Error("couldn't nack delivery")
}
continue
}
buildJob.rawPayload, err = simplejson.NewJson(delivery.Body)
if err != nil {
context.LoggerFromContext(ctx).WithField("err", err).Error("raw payload JSON parse error, attempting to nack delivery")
err := delivery.Ack(false)
if err != nil {
context.LoggerFromContext(ctx).WithField("err", err).WithField("delivery", delivery).Error("couldn't nack delivery")
}
continue
}
buildJob.startAttributes = startAttrs.Config
buildJob.startAttributes.VMType = buildJob.payload.VMType
buildJob.startAttributes.SetDefaults(q.DefaultLanguage, q.DefaultDist, q.DefaultGroup, q.DefaultOS, VMTypeDefault)
buildJob.conn = q.conn
buildJob.delivery = delivery
buildJobChan <- buildJob
}
err := channel.Close()
if err != nil {
context.LoggerFromContext(ctx).WithField("err", err).WithField("channel", channel).Error("couldn't close channel")
}
}()
return
}
// Cleanup closes the underlying AMQP connection
func (q *AMQPJobQueue) Cleanup() error {
return q.conn.Close()
}