/
sqs_worker.go
129 lines (102 loc) · 3.01 KB
/
sqs_worker.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
package main
import "os"
import "sync"
import "flag"
import "fmt"
import "strconv"
import "time"
import "./heartbeat"
import "./work_order"
import "github.com/AdRoll/goamz/sqs"
import "github.com/ianneub/logger"
const (
VERSION = "2.0.2"
)
func init() {
print_version := flag.Bool("v", false, "display version and exit")
debug := flag.Bool("d", false, "enable debug mode")
// parse command line options
flag.Parse()
// display version and exit
if *print_version {
fmt.Println("SQS worker version:", VERSION)
os.Exit(0)
}
// set debug
if *debug {
logger.SetDebug(true)
}
}
func main() {
logger.Info("Starting worker v%s", VERSION)
// get worker count
workers, err := strconv.Atoi(os.Getenv("WORKER_COUNT"))
if err != nil {
workers = 10
}
logger.Info("Worker count: %d", workers)
// access key, secret key, receive queue and report queue should be in ENV variables
logger.Info("SQS queue: %s", os.Getenv("SQS_WORKER_QUEUE"))
// create sqs client
client, err := sqs.NewFrom(os.Getenv("SQS_WORKER_ACCESS_KEY"), os.Getenv("SQS_WORKER_SECRET_KEY"), "us-east-1")
if err != nil {
logger.Fatal("CLIENT ERROR: %v", err)
}
// get the SQS queue
queue, err := client.GetQueue(os.Getenv("SQS_WORKER_QUEUE"))
if err != nil {
logger.Fatal("QUEUE ERROR: %v", err)
}
logger.Info("Worker started.")
// create the wait group
var wg sync.WaitGroup
for {
// get some messages from the sqs queue
logger.Debug("Checking for messages on the queue...")
resp, err := queue.ReceiveMessageWithVisibilityTimeout(workers, 60)
if err != nil {
logger.Error("Could not receive messages: %v", err)
time.Sleep(10 * time.Second)
}
if cap(resp.Messages) == 0 {
logger.Debug("Did not find any messages on the queue.")
}
// for each message
for _, message := range resp.Messages {
// get the message details
wo, err := work_order.NewFromJson(message.Body)
if err != nil {
logger.Error("Could not process SQS message: %s with JSON ERROR: %v", message.MessageId, err)
} else {
// process the message in a goroutine
wg.Add(1)
go processMessage(queue, message, wo, &wg)
}
}
// wait for each goroutine to exit
wg.Wait()
}
}
// process a message from the SQS queue. This should be run inside a goroutine.
func processMessage(q *sqs.Queue, m sqs.Message, wo work_order.WorkOrder, wg *sync.WaitGroup) {
logger.Debug("Starting process on %d from '%s'", wo.Id, m.MessageId)
// start heartbeat
beat := heartbeat.Start(q, &m)
// execute the work
err := wo.Execute()
if err != nil {
logger.Error("Error executing: %d - %v", wo.Id, err)
}
// send response back to devops-web
wo.Report()
// stop the heartbeat
beat.Stop()
// delete message
logger.Debug("Deleting message: %s", m.MessageId)
_, err = q.DeleteMessage(&m)
if err != nil {
logger.Error("ERROR: Couldn't delete message: %s - %v", m.MessageId, err)
}
// exit this goroutine
wg.Done()
}