/
task_monitor.go
244 lines (207 loc) · 5.96 KB
/
task_monitor.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
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
package celeriac
import (
"encoding/json"
"fmt"
// Package dependencies
log "github.com/sirupsen/logrus"
"github.com/streadway/amqp"
)
/*
TaskMonitor is a Celery task event consumer
*/
type TaskMonitor struct {
connection *amqp.Connection
channel *amqp.Channel
consumerTag string
done chan error
processMessageChannel bool
deliveryChannel <-chan amqp.Delivery
monitorWorkerHeartbeatEvents bool
// Public channel on which events are piped
EventsChannel chan interface{}
}
/*
NewTaskMonitor is a factory function that creates a new Celery consumer
*/
func NewTaskMonitor(connection *amqp.Connection, channel *amqp.Channel,
exchangeName string, exchangeType string, queueName string, bindingKey string, ctag string) (*TaskMonitor, error) {
monitor := &TaskMonitor{
connection: connection,
channel: channel,
consumerTag: ctag,
done: make(chan error),
EventsChannel: make(chan interface{}),
monitorWorkerHeartbeatEvents: false,
}
var err error
log.Printf("Declaring exchange (%q)", exchangeName)
if err = monitor.channel.ExchangeDeclare(
exchangeName, // name of the exchange
exchangeType, // type
true, // durable
false, // delete when complete
false, // internal
false, // noWait
nil, // arguments
); err != nil {
return nil, fmt.Errorf("error declaring exchange: %s", err)
}
log.Printf("Declared exchange, declaring queue %q", queueName)
queue, err := monitor.channel.QueueDeclare(
queueName, // name of the queue
true, // durable
false, // delete when unused
false, // exclusive
false, // noWait
nil, // arguments
)
if err != nil {
return nil, fmt.Errorf("error declaring queue: %s", err)
}
log.Printf("Declared queue (%q %d messages, %d consumers), binding to Exchange (binding key %q)",
queue.Name, queue.Messages, queue.Consumers, bindingKey)
if err = monitor.channel.QueueBind(
queue.Name, // name of the queue
bindingKey, // binding key
exchangeName, // sourceExchange
false, // noWait
nil, // arguments
); err != nil {
return nil, fmt.Errorf("error binding to queue: %s", err)
}
log.Printf("Queue bound to exchange, starting monitoring (consumer tag %q)", monitor.consumerTag)
monitor.deliveryChannel, err = monitor.channel.Consume(
queue.Name, // name
monitor.consumerTag, // consumerTag,
false, // auto-ack
false, // exclusive
false, // noLocal
false, // noWait
nil, // arguments
)
if err != nil {
return nil, fmt.Errorf("error consuming queue: %s", err)
}
// Spawn a Go routine to handle event monitoring
go monitor.handle(monitor.deliveryChannel, monitor.done, monitor.EventsChannel)
return monitor, nil
}
/*
Shutdown stops all monitoring, cleaning up any open connections
*/
func (monitor *TaskMonitor) Shutdown() error {
// Close the events channel
close(monitor.EventsChannel)
// Will close() the deliveries channel
if err := monitor.channel.Cancel(monitor.consumerTag, true); err != nil {
return fmt.Errorf("failed to cancel monitor: %s", err)
}
if err := monitor.connection.Close(); err != nil {
return fmt.Errorf("error closing AMQP connection: %s", err)
}
defer log.Printf("shutdown AMQP OK")
// Wait for handle() to exit
return <-monitor.done
}
/*
SetMonitorWorkerHeartbeatEvents sets the property whether to process heartbeat events emitted
by workers.
NOTE: By default this is set to 'false' so as to minimize unnecessary "noisy heartbeat" events.
*/
func (monitor *TaskMonitor) SetMonitorWorkerHeartbeatEvents(processHeartbeatEvents bool) {
monitor.monitorWorkerHeartbeatEvents = processHeartbeatEvents
}
/*
Handles Celery event messages on the task queue
*/
func (monitor *TaskMonitor) handle(deliveries <-chan amqp.Delivery, done chan error, out chan interface{}) {
rawEvent := NewEvent()
for d := range deliveries {
d.Ack(false)
// Debug code for ALL messages - useful for extracting ones we are unsure of the json format!
/*
log.Printf("Received %d bytes: [%v] %q",
len(d.Body),
d.DeliveryTag,
d.Body,
)
*/
err := rawEvent.UnmarshalJSON(d.Body)
if err != nil {
log.Warnf("%v, checking if TaskEvent array...", err)
var taskEvents TaskEventsList
err = json.Unmarshal(d.Body, &taskEvents)
if err != nil {
log.Errorf("Error: %v, Data was %d bytes: [%v] %q",
err,
len(d.Body),
d.DeliveryTag,
d.Body,
)
continue
}
for _, taskEvent := range taskEvents {
out <- &taskEvent
}
} else {
switch rawEvent.Type {
case ConstEventTypeWorkerOnline,
ConstEventTypeWorkerOffline:
var t = NewWorkerEvent()
err := t.UnmarshalJSON(d.Body)
if err != nil {
log.Errorf("Error: %v, Data was %d bytes: [%v] %q",
err,
len(d.Body),
d.DeliveryTag,
d.Body,
)
continue
}
out <- t
break
case ConstEventTypeWorkerHeartbeat:
if monitor.monitorWorkerHeartbeatEvents {
var t = NewWorkerEvent()
err := t.UnmarshalJSON(d.Body)
if err != nil {
log.Errorf("Error: %v, Data was %d bytes: [%v] %q",
err,
len(d.Body),
d.DeliveryTag,
d.Body,
)
continue
}
out <- t
}
break
case ConstEventTypeTaskSent,
ConstEventTypeTaskReceived,
ConstEventTypeTaskStarted,
ConstEventTypeTaskSucceeded,
ConstEventTypeTaskFailed,
ConstEventTypeTaskRevoked,
ConstEventTypeTaskRetried:
var t = NewTaskEvent()
err := t.UnmarshalJSON(d.Body)
if err != nil {
log.Errorf("Error: %v, Data was %d bytes: [%v] %q",
err,
len(d.Body),
d.DeliveryTag,
d.Body,
)
continue
}
out <- t
break
default:
out <- rawEvent
break
}
}
}
log.Printf("deliveries channel closed")
done <- nil
}