├── Dockerfile ├── .gitignore ├── go.mod ├── LICENSE ├── .github └── workflows │ └── build.yaml ├── README.md ├── cron_worker └── worker.go ├── supervisor ├── supervisor.go └── supervisor_test.go ├── cmd └── simplesqsd │ └── simplesqsd.go └── go.sum /Dockerfile: -------------------------------------------------------------------------------- 1 | FROM golang:1.16-alpine as builder 2 | 3 | WORKDIR /app 4 | COPY . . 5 | 6 | RUN apk --no-cache add git alpine-sdk build-base gcc 7 | 8 | RUN go build cmd/simplesqsd/simplesqsd.go 9 | 10 | FROM alpine:latest 11 | RUN apk --no-cache add ca-certificates 12 | WORKDIR /root/ 13 | COPY --from=builder /app/simplesqsd . 14 | CMD ["./simplesqsd"] 15 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Binaries for programs and plugins 2 | *.exe 3 | *.dll 4 | *.so 5 | *.dylib 6 | 7 | # Test binary, build with `go test -c` 8 | *.test 9 | 10 | # Output of the go coverage tool, specifically when used with LiteIDE 11 | *.out 12 | 13 | # Project-local glide cache, RE: https://github.com/Masterminds/glide/issues/736 14 | .glide/ 15 | 16 | vendor/ 17 | 18 | # GoLand IDE config file 19 | .idea 20 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/fterrag/simple-sqsd 2 | 3 | go 1.12 4 | 5 | require ( 6 | github.com/aws/aws-sdk-go v1.36.18 7 | github.com/fsnotify/fsnotify v1.5.4 // indirect 8 | github.com/onsi/ginkgo v1.15.1 // indirect 9 | github.com/onsi/gomega v1.11.0 // indirect 10 | github.com/robfig/cron/v3 v3.0.0 // indirect 11 | github.com/sirupsen/logrus v1.0.4 12 | github.com/stretchr/testify v1.2.2 13 | gopkg.in/airbrake/gobrake.v2 v2.0.9 // indirect 14 | gopkg.in/gemnasium/logrus-airbrake-hook.v2 v2.1.2 // indirect 15 | gopkg.in/yaml.v2 v2.4.0 // indirect 16 | ) 17 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Frank Terragna 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /.github/workflows/build.yaml: -------------------------------------------------------------------------------- 1 | name: build 2 | 3 | on: 4 | pull_request: 5 | branches: 6 | - 'master' 7 | push: 8 | branches: 9 | - '**' 10 | tags: 11 | - 'v*.*.*' 12 | release: 13 | types: 14 | - published 15 | 16 | jobs: 17 | docker: 18 | runs-on: ubuntu-latest 19 | steps: 20 | - name: Checkout 21 | uses: actions/checkout@v2 22 | 23 | - name: Docker meta 24 | id: docker_meta 25 | uses: crazy-max/ghaction-docker-meta@v1 26 | with: 27 | # list of Docker images to use as base name for tags 28 | images: | 29 | ghcr.io/fterrag/simple-sqsd 30 | # add git short SHA as Docker tag 31 | tag-sha: true 32 | 33 | - name: Set up QEMU 34 | uses: docker/setup-qemu-action@v1 35 | 36 | - name: Set up Docker Buildx 37 | uses: docker/setup-buildx-action@v1 38 | 39 | - name: Login to GitHub Container Registry 40 | uses: docker/login-action@v1 41 | with: 42 | registry: ghcr.io 43 | username: ${{ github.repository_owner }} 44 | password: ${{ secrets.GHCR_TOKEN }} 45 | 46 | - name: Build and push 47 | uses: docker/build-push-action@v2 48 | with: 49 | context: . 50 | platforms: linux/amd64,linux/arm64 51 | push: ${{ github.event_name != 'pull_request' }} 52 | tags: ${{ steps.docker_meta.outputs.tags }} 53 | labels: ${{ steps.docker_meta.outputs.labels }} 54 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | [![Go Report Card](https://goreportcard.com/badge/github.com/fterrag/simple-sqsd)](https://goreportcard.com/report/github.com/fterrag/simple-sqsd) 2 | 3 | # simple-sqsd 4 | 5 | A simple version of the AWS Elastic Beanstalk Worker Environment SQS daemon ([sqsd](https://docs.aws.amazon.com/elasticbeanstalk/latest/dg/using-features-managing-env-tiers.html#worker-daemon)). 6 | 7 | ## Getting Started 8 | 9 | ```bash 10 | $ SQSD_QUEUE_REGION=us-east-1 SQSD_QUEUE_URL=http://queue.url SQSD_HTTP_URL=http://service.url/endpoint go run cmd/simplesqsd/simplesqsd.go 11 | ``` 12 | 13 | Docker (uses a GitHub Container Registry): 14 | ```bash 15 | $ docker run -e AWS_ACCESS_KEY_ID=your-access-id -e AWS_SECRET_ACCESS_KEY=your-secret-key -e SQSD_QUEUE_REGION=us-east-1 -e SQSD_QUEUE_URL=http://queue.url -e SQSD_HTTP_URL=http://service.url/endpoint ghcr.io/fterrag/simple-sqsd:latest 16 | ``` 17 | 18 | ## Configuration 19 | 20 | |**Environment Variable**|**Default Value**|**Required**|**Description**| 21 | |-|-|-|-| 22 | |`SQSD_QUEUE_REGION`||yes|The region of the SQS queue.| 23 | |`SQSD_QUEUE_URL`||yes|The URL of the SQS queue.| 24 | |`SQSD_QUEUE_MAX_MSGS`|`10`|no|Max number of messages a worker should try to receive from the SQS queue.| 25 | |`SQSD_QUEUE_WAIT_TIME`|`10`|no|The duration (in seconds) for which the call waits for a message to arrive in the queue before returning. Setting this to `0` disables long polling. Maximum of `20` seconds.| 26 | |`SQSD_HTTP_MAX_CONNS`|`25`|no|Maximum number of concurrent HTTP requests to make to SQSD_HTTP_URL.| 27 | |`SQSD_HTTP_URL`||yes|The URL of your service to make a request to.| 28 | |`SQSD_HTTP_CONTENT_TYPE` ||no|The value to send for the HTTP header `Content-Type` when making a request to your service.| 29 | |`SQSD_HTTP_USER_AGENT`||no|The value to send for the HTTP header `User-Agent` when making a request to your service.| 30 | |`SQSD_AWS_ENDPOINT` ||no|Sets the AWS endpoint.| 31 | |`SQSD_HTTP_HMAC_HEADER`||no|The name of the HTTP header to send the HMAC hash with.| 32 | |`SQSD_HMAC_SECRET_KEY`||no|Secret key to use when generating HMAC hash send to `SQSD_HTTP_URL`.| 33 | |`SQSD_HTTP_HEALTH_PATH`||no|The path to a health check endpoint of your service. When provided, messages will not be processed until the health check returns a 200 for `HTTPHealthInterval` times | 34 | |`SQSD_HTTP_HEALTH_WAIT`|`5`|no|How long to wait before starting health checks| 35 | |`SQSD_HTTP_HEALTH_INTERVAL`|`5`|no|How often to wait between health checks| 36 | |`SQSD_HTTP_HEALTH_SUCCESS_COUNT`|`1`|no|How many successful health checks required in a row| 37 | |`SQSD_HTTP_TIMEOUT`|`15`|no|Number of seconds to wait for a response from the worker| 38 | |`SQSD_SQS_HTTP_TIMEOUT`|`15`|no|Number of seconds to wait for a response from sqs| 39 | |`SQSD_HTTP_SSL_VERIFY`|`true`|no|Enable SSL Verification on the URL of your service to make a request to (if you're using self-signed certificate)| 40 | |`SQSD_HTTP_AUTHORIZATION_HEADER`||no|A simple feature to add a jwt/simple token to Authorization header for basic auth on SQSD_HTTP_URL | 41 | |`SQSD_HTTP_AUTHORIZATION_HEADER_NAME`||no|override the http header name (defaults to Authorization) in SQSD_HTTP_AUTHORIZATION_HEADER | 42 | |`SQSD_CRON_FILE`||no|The elastic beanstalk cron.yaml file to load| 43 | |`SQSD_CRON_ENDPOINT`|`SQSD_HTTP_URL` without path/query|yes if SQSD_CRON_FILE|The base URL to call (e.g. http://localhost:3000). cron.yaml url will be appended to this| 44 | |`SQSD_CRON_TIMEOUT`|`15`|no|Duration (in seconds) To wait for the cron endpoint to response| 45 | 46 | ## HMAC 47 | 48 | *Optionally* (when SQSD_HTTP_HMAC_HEADER and SQSD_HMAC_SECRET_KEY are set), HMAC hashes are generated using SHA-256 with the signature made up of the following: 49 | ``` 50 | POST {SQSD_HTTP_URL}\n 51 | 52 | ``` 53 | 54 | ## Support 429 Status codes with Retry-After 55 | 56 | * SQSD will attempt to change the message visibility when the service responds with [429 status code](https://tools.ietf.org/html/rfc6585#section-4). 57 | * `Retry-After` response header should contain an integer with the amount of senconds to wait. 58 | 59 | ## Todo 60 | - [ ] More Tests 61 | - [ ] Documentation 62 | 63 | ## Contributing 64 | 65 | * Submit a PR 66 | * Add or improve documentation 67 | * Report issues 68 | * Suggest new features or enhancements 69 | -------------------------------------------------------------------------------- /cron_worker/worker.go: -------------------------------------------------------------------------------- 1 | package cron_worker 2 | 3 | import ( 4 | "errors" 5 | "github.com/fsnotify/fsnotify" 6 | "github.com/robfig/cron/v3" 7 | log "github.com/sirupsen/logrus" 8 | "gopkg.in/yaml.v2" 9 | "io" 10 | "net/http" 11 | "os" 12 | "sync" 13 | "time" 14 | ) 15 | 16 | type ( 17 | Config struct { 18 | File string 19 | EndPoint string 20 | Timeout time.Duration 21 | UserAgent string 22 | HTTPAUTHORIZATIONHeader string 23 | HTTPAUTHORIZATIONHeaderName string 24 | HTTPContentType string 25 | } 26 | Worker struct { 27 | config *Config 28 | crontab *sqsCron 29 | 30 | cron *cron.Cron 31 | 32 | fsDoneChan chan bool 33 | 34 | mu sync.Mutex 35 | } 36 | sqsCronItem struct { 37 | Name string `yaml:"name"` 38 | Url string `yaml:"url"` 39 | Schedule string `yaml:"schedule"` 40 | cronEntryId cron.EntryID 41 | } 42 | sqsCron struct { 43 | Version int `yaml:"version"` 44 | Cron []sqsCronItem 45 | } 46 | ) 47 | 48 | func New(c *Config) *Worker { 49 | if "" == c.File { 50 | log.Info("No Cron File Supplied") 51 | return nil 52 | } 53 | 54 | if c.Timeout.Seconds() < 1 { 55 | c.Timeout = 30 * time.Second 56 | } 57 | 58 | wkr := Worker{ 59 | config: c, 60 | crontab: nil, 61 | } 62 | wkr.fsDoneChan = make(chan bool, 10) 63 | wkr.loadCronTab() 64 | 65 | return &wkr 66 | } 67 | 68 | // Run handles the running of the cron tab and the watching of the cron yaml file and reloading when required 69 | func (w *Worker) Run() { 70 | if nil == w.cron { 71 | log.Errorf("Cannot run cron") 72 | return 73 | } 74 | w.cron.Start() 75 | 76 | for _, entry := range w.crontab.Cron { 77 | log. 78 | WithField("what", "cron"). 79 | WithField("next", w.cron.Entry(entry.cronEntryId).Next). 80 | WithField("name", entry.Name). 81 | Debug("Next Occurrence") 82 | } 83 | 84 | watcher, err := fsnotify.NewWatcher() 85 | if nil != err { 86 | log.WithError(err).WithField("file", w.config.File).Error("Unable to watch cron file for changes") 87 | return 88 | } 89 | err = watcher.Add(w.config.File) 90 | if nil != err { 91 | log.WithError(err).WithField("file", w.config.File).Error("Unable to watch cron file for changes") 92 | return 93 | } 94 | 95 | go func() { 96 | defer func() { 97 | if err = watcher.Close(); nil != err { 98 | log.WithError(err).Error("error closing fs watcher") 99 | } 100 | }() 101 | for { 102 | select { 103 | case <-w.fsDoneChan: 104 | return 105 | case event, ok := <-watcher.Events: 106 | if !ok { 107 | return 108 | } 109 | if event.Op&fsnotify.Write == fsnotify.Write { 110 | log.WithField("file", w.config.File).Info("cron file changed. reloading") 111 | w.loadCronTab() 112 | w.Run() 113 | return 114 | } 115 | case err, ok := <-watcher.Errors: 116 | if !ok { 117 | return 118 | } 119 | log.WithError(err).Error("FSNotify Error") 120 | } 121 | } 122 | }() 123 | } 124 | 125 | // Stop handles shutting down the cron worker safely 126 | func (w *Worker) Stop() { 127 | w.fsDoneChan <- true 128 | if nil != w.cron { 129 | w.cron.Stop() 130 | } 131 | } 132 | 133 | // loadCronTab is the parent method that reads, parses and then loads the crontab 134 | func (w *Worker) loadCronTab() { 135 | w.mu.Lock() 136 | defer w.mu.Unlock() 137 | if nil != w.cron { 138 | w.cron.Stop() 139 | w.cron = nil 140 | } 141 | 142 | contents, err := w.readCronTab(w.config.File) 143 | if nil != err { 144 | log.WithError(err).Error("Failed to load crontab") 145 | return 146 | } 147 | 148 | if err = w.parseCronTab(contents); nil != err { 149 | log.WithError(err).Error("Failed to parse crontab") 150 | return 151 | } 152 | 153 | // log some info about our crontab 154 | log. 155 | WithField("file", w.config.File). 156 | WithField("version", w.crontab.Version). 157 | WithField("num-entries", len(w.crontab.Cron)). 158 | Info("EBS Crontab Info") 159 | 160 | if err = w.loadCronEntries(); nil != err { 161 | log.WithError(err).Error("Failed to start crontab") 162 | return 163 | } 164 | } 165 | 166 | func (w *Worker) readCronTab(path string) ([]byte, error) { 167 | fh, err := os.Open(path) 168 | if nil != err { 169 | return nil, err 170 | } 171 | 172 | return io.ReadAll(fh) 173 | } 174 | 175 | func (w *Worker) parseCronTab(contents []byte) error { 176 | crontab := sqsCron{} 177 | 178 | err := yaml.Unmarshal(contents, &crontab) 179 | if nil != err { 180 | return err 181 | } 182 | 183 | w.crontab = &crontab 184 | return nil 185 | } 186 | 187 | func (w *Worker) loadCronEntries() error { 188 | if nil == w.crontab { 189 | return errors.New("please parse a crontab before loading it") 190 | } 191 | 192 | w.cron = cron.New() 193 | 194 | for idx, entry := range w.crontab.Cron { 195 | entryId, err := w.cron.AddFunc(entry.Schedule, w.makeCronRequestFunc(entry)) 196 | if err != nil { 197 | log. 198 | WithField("what", "cron"). 199 | WithError(err). 200 | WithField("entry", entry.Name). 201 | Error("Failed to load cron entry") 202 | continue 203 | } 204 | w.crontab.Cron[idx].cronEntryId = entryId 205 | } 206 | 207 | return nil 208 | } 209 | 210 | func (w *Worker) makeCronRequestFunc(entry sqsCronItem) func() { 211 | cronUrl := w.config.EndPoint + entry.Url 212 | return func() { 213 | t1 := time.Now() 214 | rqLog := log. 215 | WithField("what", "cron"). 216 | WithField("entry", entry.Name). 217 | WithField("url", cronUrl). 218 | WithField("start", t1) 219 | 220 | rqLog.Debug("Requesting Cron URL") 221 | 222 | client := &http.Client{} 223 | client.Timeout = w.config.Timeout 224 | 225 | //resp, err := client.Post(cronUrl, "application/json", nil) 226 | req, err := http.NewRequest("POST", cronUrl, nil) 227 | req.Header.Add("X-Aws-Sqsd-Taskname", entry.Name) 228 | 229 | if len(w.config.HTTPAUTHORIZATIONHeader) > 0 { 230 | headerName := w.config.HTTPAUTHORIZATIONHeaderName 231 | if len(headerName) == 0 { 232 | headerName = "Authorization" 233 | } 234 | req.Header.Set(headerName, w.config.HTTPAUTHORIZATIONHeader) 235 | } 236 | 237 | if len(w.config.HTTPContentType) > 0 { 238 | req.Header.Set("Content-Type", w.config.HTTPContentType) 239 | } 240 | 241 | if len(w.config.UserAgent) > 0 { 242 | req.Header.Set("User-Agent", w.config.UserAgent) 243 | } 244 | 245 | res, err := client.Do(req) 246 | 247 | t2 := time.Now() 248 | dur := t2.Sub(t1) 249 | rqLog = rqLog.WithField("duration", dur.String()) 250 | if err != nil { 251 | rqLog. 252 | WithError(err). 253 | Error("Failed Requesting Endpoint") 254 | return 255 | } 256 | rqLog = rqLog.WithField("http-status", res.StatusCode) 257 | 258 | if res.StatusCode < 200 || res.StatusCode > 299 { 259 | rqLog. 260 | Error("Requesting cron endpoint resulted in non 2XX Status Code") 261 | } else { 262 | rqLog.Info("Cron Success") 263 | } 264 | 265 | res.Body.Close() 266 | 267 | } 268 | } 269 | -------------------------------------------------------------------------------- /supervisor/supervisor.go: -------------------------------------------------------------------------------- 1 | package supervisor 2 | 3 | import ( 4 | "bytes" 5 | "crypto/hmac" 6 | "crypto/sha256" 7 | "encoding/hex" 8 | "errors" 9 | "fmt" 10 | "net/http" 11 | "strconv" 12 | "strings" 13 | "sync" 14 | 15 | "github.com/aws/aws-sdk-go/aws" 16 | "github.com/aws/aws-sdk-go/service/sqs" 17 | "github.com/aws/aws-sdk-go/service/sqs/sqsiface" 18 | log "github.com/sirupsen/logrus" 19 | ) 20 | 21 | type Supervisor struct { 22 | sync.Mutex 23 | 24 | logger *log.Entry 25 | sqs sqsiface.SQSAPI 26 | httpClient httpClient 27 | workerConfig WorkerConfig 28 | hmacSignature string 29 | 30 | startOnce sync.Once 31 | wg sync.WaitGroup 32 | 33 | shutdown bool 34 | } 35 | 36 | type WorkerConfig struct { 37 | QueueURL string 38 | QueueMaxMessages int 39 | QueueWaitTime int 40 | 41 | HTTPURL string 42 | HTTPContentType string 43 | 44 | HTTPAUTHORIZATIONHeader string 45 | HTTPAUTHORIZATIONHeaderName string 46 | 47 | HTTPHMACHeader string 48 | HMACSecretKey []byte 49 | 50 | UserAgent string 51 | } 52 | 53 | type httpClient interface { 54 | Do(req *http.Request) (*http.Response, error) 55 | } 56 | 57 | func NewSupervisor(logger *log.Entry, sqs sqsiface.SQSAPI, httpClient httpClient, config WorkerConfig) *Supervisor { 58 | return &Supervisor{ 59 | logger: logger, 60 | sqs: sqs, 61 | httpClient: httpClient, 62 | workerConfig: config, 63 | hmacSignature: fmt.Sprintf("POST %s\n", config.HTTPURL), 64 | } 65 | } 66 | 67 | func (s *Supervisor) Start(numWorkers int) { 68 | s.startOnce.Do(func() { 69 | s.wg.Add(numWorkers) 70 | 71 | for i := 0; i < numWorkers; i++ { 72 | go s.worker() 73 | } 74 | }) 75 | } 76 | 77 | func (s *Supervisor) Wait() { 78 | s.wg.Wait() 79 | } 80 | 81 | func (s *Supervisor) Shutdown() { 82 | defer s.Unlock() 83 | s.Lock() 84 | 85 | s.shutdown = true 86 | } 87 | 88 | func (s *Supervisor) worker() { 89 | defer s.wg.Done() 90 | 91 | s.logger.Info("Starting worker") 92 | 93 | for { 94 | if s.shutdown { 95 | return 96 | } 97 | 98 | recInput := &sqs.ReceiveMessageInput{ 99 | MaxNumberOfMessages: aws.Int64(int64(s.workerConfig.QueueMaxMessages)), 100 | QueueUrl: aws.String(s.workerConfig.QueueURL), 101 | WaitTimeSeconds: aws.Int64(int64(s.workerConfig.QueueWaitTime)), 102 | MessageAttributeNames: aws.StringSlice([]string{"All"}), 103 | } 104 | 105 | output, err := s.sqs.ReceiveMessage(recInput) 106 | if err != nil { 107 | s.logger.Errorf("Error while receiving messages from the queue: %s", err) 108 | continue 109 | } 110 | 111 | if len(output.Messages) == 0 { 112 | continue 113 | } 114 | 115 | deleteEntries := make([]*sqs.DeleteMessageBatchRequestEntry, 0) 116 | changeVisibilityEntries := make([]*sqs.ChangeMessageVisibilityBatchRequestEntry, 0) 117 | 118 | for _, msg := range output.Messages { 119 | res, err := s.httpRequest(msg) 120 | if err != nil { 121 | s.logger.Errorf("Error making HTTP request: %s", err) 122 | continue 123 | } 124 | 125 | if res.StatusCode < http.StatusOK || res.StatusCode > http.StatusIMUsed { 126 | 127 | if res.StatusCode == http.StatusTooManyRequests { 128 | sec, err := getRetryAfterFromResponse(res) 129 | if err != nil { 130 | s.logger.Errorf("Error getting retry after value from HTTP response: %s", err) 131 | continue 132 | } 133 | 134 | changeVisibilityEntries = append(changeVisibilityEntries, &sqs.ChangeMessageVisibilityBatchRequestEntry{ 135 | Id: msg.MessageId, 136 | ReceiptHandle: msg.ReceiptHandle, 137 | VisibilityTimeout: aws.Int64(sec), 138 | }) 139 | } 140 | 141 | s.logger.Errorf("Non-successful status code: %d", res.StatusCode) 142 | 143 | continue 144 | 145 | } 146 | 147 | deleteEntries = append(deleteEntries, &sqs.DeleteMessageBatchRequestEntry{ 148 | Id: msg.MessageId, 149 | ReceiptHandle: msg.ReceiptHandle, 150 | }) 151 | 152 | s.logger.Debugf("Message %s successfully processed", *msg.MessageId) 153 | } 154 | 155 | if len(deleteEntries) > 0 { 156 | delInput := &sqs.DeleteMessageBatchInput{ 157 | Entries: deleteEntries, 158 | QueueUrl: aws.String(s.workerConfig.QueueURL), 159 | } 160 | 161 | _, err = s.sqs.DeleteMessageBatch(delInput) 162 | if err != nil { 163 | s.logger.Errorf("Error while deleting messages from SQS: %s", err) 164 | } 165 | } 166 | 167 | if len(changeVisibilityEntries) > 0 { 168 | changeVisibilityInput := &sqs.ChangeMessageVisibilityBatchInput{ 169 | Entries: changeVisibilityEntries, 170 | QueueUrl: aws.String(s.workerConfig.QueueURL), 171 | } 172 | 173 | _, err = s.sqs.ChangeMessageVisibilityBatch(changeVisibilityInput) 174 | if err != nil { 175 | s.logger.Errorf("Error while changing visibility on messages from SQS: %s", err) 176 | } 177 | } 178 | } 179 | } 180 | 181 | func (s *Supervisor) httpRequest(msg *sqs.Message) (*http.Response, error) { 182 | body := *msg.Body 183 | req, err := http.NewRequest("POST", s.workerConfig.HTTPURL, bytes.NewBufferString(body)) 184 | req.Header.Add("X-Aws-Sqsd-Msgid", *msg.MessageId) 185 | s.addMessageAttributesToHeader(msg.MessageAttributes, req.Header) 186 | if err != nil { 187 | return nil, fmt.Errorf("Error while creating HTTP request: %s", err) 188 | } 189 | 190 | if len(s.workerConfig.HMACSecretKey) > 0 { 191 | hmac, err := makeHMAC(strings.Join([]string{s.hmacSignature, body}, ""), s.workerConfig.HMACSecretKey) 192 | if err != nil { 193 | return nil, err 194 | } 195 | 196 | req.Header.Set(s.workerConfig.HTTPHMACHeader, hmac) 197 | } 198 | 199 | if len(s.workerConfig.HTTPAUTHORIZATIONHeader) > 0 { 200 | headerName := s.workerConfig.HTTPAUTHORIZATIONHeaderName 201 | if len(headerName) == 0 { 202 | headerName = "Authorization" 203 | } 204 | req.Header.Set(headerName, s.workerConfig.HTTPAUTHORIZATIONHeader) 205 | } 206 | 207 | if len(s.workerConfig.HTTPContentType) > 0 { 208 | req.Header.Set("Content-Type", s.workerConfig.HTTPContentType) 209 | } 210 | 211 | if len(s.workerConfig.UserAgent) > 0 { 212 | req.Header.Set("User-Agent", s.workerConfig.UserAgent) 213 | } 214 | 215 | res, err := s.httpClient.Do(req) 216 | if err != nil { 217 | return res, err 218 | } 219 | 220 | res.Body.Close() 221 | 222 | return res, nil 223 | } 224 | 225 | func (s *Supervisor) addMessageAttributesToHeader(attrs map[string]*sqs.MessageAttributeValue, header http.Header) { 226 | for k, v := range attrs { 227 | header.Add("X-Aws-Sqsd-Attr-" + k, *v.StringValue) 228 | } 229 | } 230 | 231 | func makeHMAC(signature string, secretKey []byte) (string, error) { 232 | mac := hmac.New(sha256.New, secretKey) 233 | 234 | _, err := mac.Write([]byte(signature)) 235 | if err != nil { 236 | return "", fmt.Errorf("Error while writing HMAC: %s", err) 237 | } 238 | 239 | return hex.EncodeToString(mac.Sum(nil)), nil 240 | } 241 | 242 | func getRetryAfterFromResponse(res *http.Response) (int64, error) { 243 | retryAfter := res.Header.Get("Retry-After") 244 | if len(retryAfter) == 0 { 245 | return 0, errors.New("Retry-After header value is empty") 246 | } 247 | 248 | seconds, err := strconv.ParseInt(retryAfter, 10, 0) 249 | if err != nil { 250 | return 0, err 251 | } 252 | 253 | return seconds, nil 254 | } 255 | -------------------------------------------------------------------------------- /cmd/simplesqsd/simplesqsd.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "crypto/tls" 5 | "errors" 6 | "fmt" 7 | "github.com/fterrag/simple-sqsd/cron_worker" 8 | "net/http" 9 | "net/url" 10 | "os" 11 | "strconv" 12 | "time" 13 | 14 | "github.com/aws/aws-sdk-go/aws" 15 | "github.com/aws/aws-sdk-go/aws/session" 16 | "github.com/aws/aws-sdk-go/service/sqs" 17 | "github.com/fterrag/simple-sqsd/supervisor" 18 | log "github.com/sirupsen/logrus" 19 | ) 20 | 21 | type config struct { 22 | QueueRegion string 23 | QueueURL string 24 | QueueMaxMessages int 25 | QueueWaitTime int 26 | 27 | HTTPMaxConns int 28 | HTTPURL string 29 | HTTPContentType string 30 | HTTPTimeout int 31 | 32 | AWSEndpoint string 33 | HTTPHMACHeader string 34 | HTTPAUTHORIZATIONHeader string 35 | HTTPAUTHORIZATIONHeaderName string 36 | HMACSecretKey []byte 37 | 38 | HTTPHealthPath string 39 | HTTPHealthWait int 40 | HTTPHealthInterval int 41 | HTTPHealthSucessCount int 42 | 43 | SQSHTTPTimeout int 44 | SSLVerify bool 45 | 46 | CronFile string 47 | CronEndPoint string 48 | CronTimeout int 49 | 50 | UserAgent string 51 | } 52 | 53 | func main() { 54 | 55 | c := &config{} 56 | 57 | c.QueueRegion = os.Getenv("SQSD_QUEUE_REGION") 58 | c.QueueURL = os.Getenv("SQSD_QUEUE_URL") 59 | c.QueueMaxMessages = getEnvInt("SQSD_QUEUE_MAX_MSGS", 10) 60 | c.QueueWaitTime = getEnvInt("SQSD_QUEUE_WAIT_TIME", 10) 61 | 62 | c.HTTPMaxConns = getEnvInt("SQSD_HTTP_MAX_CONNS", 25) 63 | c.HTTPURL = os.Getenv("SQSD_HTTP_URL") 64 | c.HTTPContentType = os.Getenv("SQSD_HTTP_CONTENT_TYPE") 65 | c.UserAgent = os.Getenv("SQSD_HTTP_USER_AGENT") 66 | 67 | c.HTTPHealthPath = os.Getenv("SQSD_HTTP_HEALTH_PATH") 68 | c.HTTPHealthWait = getEnvInt("SQSD_HTTP_HEALTH_WAIT", 5) 69 | c.HTTPHealthInterval = getEnvInt("SQSD_HTTP_HEALTH_INTERVAL", 5) 70 | c.HTTPHealthSucessCount = getEnvInt("SQSD_HTTP_HEALTH_SUCCESS_COUNT", 1) 71 | c.HTTPTimeout = getEnvInt("SQSD_HTTP_TIMEOUT", 15) 72 | 73 | c.AWSEndpoint = os.Getenv("SQSD_AWS_ENDPOINT") 74 | c.HTTPHMACHeader = os.Getenv("SQSD_HTTP_HMAC_HEADER") 75 | c.HTTPAUTHORIZATIONHeader = os.Getenv("SQSD_HTTP_AUTHORIZATION_HEADER") 76 | c.HTTPAUTHORIZATIONHeaderName = os.Getenv("SQSD_HTTP_AUTHORIZATION_HEADER_NAME") 77 | c.HMACSecretKey = []byte(os.Getenv("SQSD_HMAC_SECRET_KEY")) 78 | 79 | c.SQSHTTPTimeout = getEnvInt("SQSD_SQS_HTTP_TIMEOUT", 15) 80 | c.SSLVerify = getenvBool("SQSD_HTTP_SSL_VERIFY", true) 81 | 82 | c.CronFile = os.Getenv("SQSD_CRON_FILE") 83 | c.CronEndPoint = os.Getenv("SQSD_CRON_ENDPOINT") 84 | c.CronTimeout = getEnvInt("SQSD_CRON_TIMEOUT", 15) 85 | 86 | 87 | if len(c.QueueRegion) == 0 { 88 | log.Fatal("SQSD_QUEUE_REGION cannot be empty") 89 | } 90 | 91 | if len(c.QueueURL) == 0 { 92 | log.Fatal("SQSD_QUEUE_URL cannot be empty") 93 | } 94 | 95 | if len(c.HTTPURL) == 0 { 96 | log.Fatal("SQSD_HTTP_URL cannot be empty") 97 | } 98 | 99 | log.SetFormatter(&log.JSONFormatter{}) 100 | 101 | logLevel := os.Getenv("LOG_LEVEL") 102 | if len(logLevel) == 0 { 103 | logLevel = "info" 104 | } 105 | if parsedLevel, err := log.ParseLevel(logLevel); err == nil { 106 | log.SetLevel(parsedLevel) 107 | } else { 108 | log.Fatal(err) 109 | } 110 | 111 | logger := log.WithFields(log.Fields{ 112 | "queueRegion": c.QueueRegion, 113 | "queueUrl": c.QueueURL, 114 | "httpMaxConns": c.HTTPMaxConns, 115 | "httpPath": c.HTTPURL, 116 | }) 117 | 118 | if len(c.HTTPHealthPath) != 0 { 119 | numSuccesses := 0 120 | healthURL := fmt.Sprintf("%s%s", c.HTTPURL, c.HTTPHealthPath) 121 | log.Infof("Waiting %d seconds before staring health check at '%s'", c.HTTPHealthWait, healthURL) 122 | time.Sleep(time.Duration(c.HTTPHealthWait) * time.Second) 123 | for { 124 | if resp, err := http.Get(healthURL); err == nil { 125 | log.Infof("%#v", resp) 126 | if numSuccesses == c.HTTPHealthSucessCount { 127 | break 128 | } else { 129 | numSuccesses++ 130 | } 131 | } else { 132 | log.Debugf("Health check failed: %s. Waiting for %d seconds before next attempt", err, c.HTTPHealthInterval) 133 | time.Sleep(time.Duration(c.HTTPHealthInterval) * time.Second) 134 | } 135 | } 136 | log.Info("Health check succeeded. Starting message processing") 137 | } 138 | 139 | awsSess := session.Must(session.NewSessionWithOptions(session.Options{ 140 | SharedConfigState: session.SharedConfigEnable, 141 | })) 142 | 143 | sqsHttpClient := &http.Client{ 144 | Timeout: time.Duration(c.SQSHTTPTimeout) * time.Second, 145 | Transport: &http.Transport{ 146 | MaxIdleConns: c.HTTPMaxConns, 147 | MaxIdleConnsPerHost: c.HTTPMaxConns, 148 | }, 149 | } 150 | sqsConfig := aws.NewConfig(). 151 | WithRegion(c.QueueRegion). 152 | WithHTTPClient(sqsHttpClient) 153 | 154 | if len(c.AWSEndpoint) > 0 { 155 | sqsConfig.WithEndpoint(c.AWSEndpoint) 156 | } 157 | 158 | sqsSvc := sqs.New(awsSess, sqsConfig) 159 | 160 | wConf := supervisor.WorkerConfig{ 161 | QueueURL: c.QueueURL, 162 | QueueMaxMessages: c.QueueMaxMessages, 163 | QueueWaitTime: c.QueueWaitTime, 164 | 165 | HTTPURL: c.HTTPURL, 166 | HTTPContentType: c.HTTPContentType, 167 | 168 | HTTPAUTHORIZATIONHeader: c.HTTPAUTHORIZATIONHeader, 169 | HTTPAUTHORIZATIONHeaderName: c.HTTPAUTHORIZATIONHeaderName, 170 | 171 | HTTPHMACHeader: c.HTTPHMACHeader, 172 | HMACSecretKey: c.HMACSecretKey, 173 | 174 | UserAgent: c.UserAgent, 175 | } 176 | 177 | httpClient := &http.Client{ 178 | Transport: &http.Transport{ 179 | MaxIdleConns: c.HTTPMaxConns, 180 | MaxIdleConnsPerHost: c.HTTPMaxConns, 181 | TLSClientConfig: &tls.Config{ 182 | MaxVersion: tls.VersionTLS11, 183 | InsecureSkipVerify: !c.SSLVerify, 184 | }, 185 | }, 186 | Timeout: time.Duration(c.HTTPTimeout) * time.Second, 187 | } 188 | 189 | if "" == c.CronEndPoint { 190 | parts, err := url.Parse(c.HTTPURL) 191 | if nil == err { 192 | parts.RawQuery = "" 193 | parts.Path = "" 194 | parts.Fragment = "" 195 | c.CronEndPoint = parts.String() 196 | } 197 | } 198 | if "" != c.CronFile && "" == c.CronEndPoint { 199 | log.Fatal("You need to specify SQSD_CRON_URL") 200 | } 201 | cronDaemon := cron_worker.New(&cron_worker.Config{ 202 | File: c.CronFile, 203 | EndPoint: c.CronEndPoint, 204 | Timeout: time.Duration(c.CronTimeout) * time.Second, 205 | UserAgent: c.UserAgent, 206 | HTTPContentType: c.HTTPContentType, 207 | HTTPAUTHORIZATIONHeader: c.HTTPAUTHORIZATIONHeader, 208 | HTTPAUTHORIZATIONHeaderName: c.HTTPAUTHORIZATIONHeaderName, 209 | }) 210 | if nil != cronDaemon { 211 | go cronDaemon.Run() 212 | } 213 | 214 | s := supervisor.NewSupervisor(logger, sqsSvc, httpClient, wConf) 215 | s.Start(c.HTTPMaxConns) 216 | s.Wait() 217 | cronDaemon.Stop() 218 | } 219 | 220 | func getEnvInt(key string, def int) int { 221 | val, err := strconv.Atoi(os.Getenv(key)) 222 | if err != nil { 223 | return def 224 | } 225 | 226 | return val 227 | } 228 | 229 | var ErrEnvVarEmpty = errors.New("getenv: environment variable empty") 230 | 231 | func getenvStr(key string) (string, error) { 232 | v := os.Getenv(key) 233 | if len(v) == 0 { 234 | return v, ErrEnvVarEmpty 235 | } 236 | return v, nil 237 | } 238 | 239 | func getenvBool(key string, def bool) bool { 240 | s, err := getenvStr(key) 241 | if err != nil { 242 | return def 243 | } 244 | v, err := strconv.ParseBool(s) 245 | if err != nil { 246 | return def 247 | } 248 | return v 249 | } 250 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/aws/aws-sdk-go v1.36.18 h1:PvfZkE0cjM1k1EMQDSb2BrX8LETPx0IFFZ/YKkurmFg= 2 | github.com/aws/aws-sdk-go v1.36.18/go.mod h1:hcU610XS61/+aQV88ixoOzUoG7v3b31pl2zKMmprdro= 3 | github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= 4 | github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 5 | github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= 6 | github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWoS4= 7 | github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ= 8 | github.com/fsnotify/fsnotify v1.5.4 h1:jRbGcIw6P2Meqdwuo0H1p6JVLbL5DHKAKlYndzMwVZI= 9 | github.com/fsnotify/fsnotify v1.5.4/go.mod h1:OVB6XrOHzAwXMpEM7uPOzcehqUV2UqJxmVXmkdnm1bU= 10 | github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= 11 | github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8= 12 | github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA= 13 | github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs= 14 | github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w= 15 | github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0= 16 | github.com/golang/protobuf v1.4.2/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= 17 | github.com/golang/protobuf v1.4.3 h1:JjCZWpVbqXDqFVmTfYWEVTMIYrL/NPdPSCHPJ0T/raM= 18 | github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= 19 | github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= 20 | github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= 21 | github.com/google/go-cmp v0.4.0 h1:xsAVV57WRhGj6kEIi8ReJzQlHHqcBYCElAvkovg3B/4= 22 | github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= 23 | github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= 24 | github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg= 25 | github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo= 26 | github.com/jmespath/go-jmespath/internal/testify v1.5.1 h1:shLQSRRSCCPj3f2gpwzGwWFoC7ycTf1rcQZHOlsJ6N8= 27 | github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U= 28 | github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A= 29 | github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE= 30 | github.com/nxadm/tail v1.4.8/go.mod h1:+ncqLTQzXmGhMZNUePPaPqPvBxHAIsmXswZKocGu+AU= 31 | github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= 32 | github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk= 33 | github.com/onsi/ginkgo v1.15.1 h1:DsXNrKujDlkMS9Rsxmd+Fg7S6Kc5lhE+qX8tY6laOxc= 34 | github.com/onsi/ginkgo v1.15.1/go.mod h1:Dd6YFfwBW84ETqqtL0CPyPXillHgY6XhQH3uuCCTr/o= 35 | github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY= 36 | github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo= 37 | github.com/onsi/gomega v1.11.0 h1:+CqWgvj0OZycCaqclBD1pxKHAU+tOkHmQIWvDHq2aug= 38 | github.com/onsi/gomega v1.11.0/go.mod h1:azGKhqFUon9Vuj0YmTfLSmx0FUwqXYSTl5re8lQLTUg= 39 | github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= 40 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 41 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 42 | github.com/robfig/cron/v3 v3.0.0 h1:kQ6Cb7aHOHTSzNVNEhmp8EcWKLb4CbiMW9h9VyIhO4E= 43 | github.com/robfig/cron/v3 v3.0.0/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro= 44 | github.com/sirupsen/logrus v1.0.4 h1:gzbtLsZC3Ic5PptoRG+kQj4L60qjK7H7XszrU163JNQ= 45 | github.com/sirupsen/logrus v1.0.4/go.mod h1:pMByvHTf9Beacp5x1UXfOR9xyW/9antXMhjMPG0dEzc= 46 | github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= 47 | github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w= 48 | github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= 49 | github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= 50 | golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= 51 | golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= 52 | golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9 h1:psW17arqaxU48Z5kZ0CQnkZWQJsqcURM6tKiBApRjXI= 53 | golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= 54 | golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= 55 | golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= 56 | golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= 57 | golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= 58 | golang.org/x/net v0.0.0-20200520004742-59133d7f0dd7/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A= 59 | golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= 60 | golang.org/x/net v0.0.0-20201110031124-69a78807bb2b/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= 61 | golang.org/x/net v0.0.0-20201202161906-c7110b5ffcbb h1:eBmm0M9fYhWpKZLjQUUKka/LtIxf46G4fxeEz5KJr9U= 62 | golang.org/x/net v0.0.0-20201202161906-c7110b5ffcbb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= 63 | golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 64 | golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 65 | golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 66 | golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= 67 | golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= 68 | golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 69 | golang.org/x/sys v0.0.0-20190904154756-749cb33beabd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 70 | golang.org/x/sys v0.0.0-20191005200804-aed5e4c7ecf9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 71 | golang.org/x/sys v0.0.0-20191120155948-bd437916bb0e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 72 | golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 73 | golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 74 | golang.org/x/sys v0.0.0-20210112080510-489259a85091 h1:DMyOG0U+gKfu8JZzg2UQe9MeaC1X+xQWlAKcRnjxjCw= 75 | golang.org/x/sys v0.0.0-20210112080510-489259a85091/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 76 | golang.org/x/sys v0.0.0-20220412211240-33da011f77ad h1:ntjMns5wyP/fN65tdBD4g8J5w8n015+iIIs9rtjXkY0= 77 | golang.org/x/sys v0.0.0-20220412211240-33da011f77ad/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 78 | golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= 79 | golang.org/x/text v0.3.3 h1:cokOdA+Jmi5PJGXLlLllQSgYigAEfHXJAERHVMaCc2k= 80 | golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= 81 | golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= 82 | golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= 83 | golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= 84 | golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 85 | golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 86 | golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 87 | golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE= 88 | golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 89 | google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8= 90 | google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0= 91 | google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM= 92 | google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE= 93 | google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo= 94 | google.golang.org/protobuf v1.23.0 h1:4MY060fB1DLGMB/7MBTLnwQUY6+F09GEiz6SsrNqyzM= 95 | google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= 96 | gopkg.in/airbrake/gobrake.v2 v2.0.9 h1:7z2uVWwn7oVeeugY1DtlPAy5H+KYgB1KeKTnqjNatLo= 97 | gopkg.in/airbrake/gobrake.v2 v2.0.9/go.mod h1:/h5ZAUhDkGaJfjzjKLSjv6zCL6O0LLBxU4K+aSYdM/U= 98 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= 99 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 100 | gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= 101 | gopkg.in/gemnasium/logrus-airbrake-hook.v2 v2.1.2 h1:OAj3g0cR6Dx/R07QgQe8wkA9RNjB2u4i700xBkIT4e0= 102 | gopkg.in/gemnasium/logrus-airbrake-hook.v2 v2.1.2/go.mod h1:Xk6kEKp8OKb+X14hQBKWaSkCsqBpgog8nAV2xsGOxlo= 103 | gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ= 104 | gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= 105 | gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= 106 | gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= 107 | gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= 108 | gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= 109 | gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= 110 | -------------------------------------------------------------------------------- /supervisor/supervisor_test.go: -------------------------------------------------------------------------------- 1 | package supervisor 2 | 3 | import ( 4 | "crypto/hmac" 5 | "crypto/sha256" 6 | "encoding/hex" 7 | "fmt" 8 | "io/ioutil" 9 | "net/http" 10 | "net/http/httptest" 11 | "testing" 12 | "time" 13 | 14 | "github.com/aws/aws-sdk-go/aws" 15 | "github.com/aws/aws-sdk-go/service/sqs" 16 | "github.com/aws/aws-sdk-go/service/sqs/sqsiface" 17 | log "github.com/sirupsen/logrus" 18 | "github.com/stretchr/testify/assert" 19 | ) 20 | 21 | type mockSQS struct { 22 | sqsiface.SQSAPI 23 | 24 | receiveMessageFunc func(*sqs.ReceiveMessageInput) (*sqs.ReceiveMessageOutput, error) 25 | deleteMessageBatchFunc func(*sqs.DeleteMessageBatchInput) (*sqs.DeleteMessageBatchOutput, error) 26 | changeMessageVisibilityBatchFunc func(*sqs.ChangeMessageVisibilityBatchInput) (*sqs.ChangeMessageVisibilityBatchOutput, error) 27 | } 28 | 29 | func (m *mockSQS) ReceiveMessage(input *sqs.ReceiveMessageInput) (*sqs.ReceiveMessageOutput, error) { 30 | if m.receiveMessageFunc != nil { 31 | return m.receiveMessageFunc(input) 32 | } 33 | 34 | return nil, nil 35 | } 36 | 37 | func (m *mockSQS) DeleteMessageBatch(input *sqs.DeleteMessageBatchInput) (*sqs.DeleteMessageBatchOutput, error) { 38 | if m.deleteMessageBatchFunc != nil { 39 | return m.deleteMessageBatchFunc(input) 40 | } 41 | 42 | return nil, nil 43 | } 44 | 45 | func (m *mockSQS) ChangeMessageVisibilityBatch(input *sqs.ChangeMessageVisibilityBatchInput) (*sqs.ChangeMessageVisibilityBatchOutput, error) { 46 | if m.changeMessageVisibilityBatchFunc != nil { 47 | return m.changeMessageVisibilityBatchFunc(input) 48 | } 49 | 50 | return nil, nil 51 | } 52 | 53 | func TestSupervisorSuccess(t *testing.T) { 54 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 55 | assert.Equal(t, "application/json", r.Header.Get("Content-Type")) 56 | 57 | w.WriteHeader(http.StatusOK) 58 | })) 59 | defer ts.Close() 60 | 61 | log.SetOutput(ioutil.Discard) 62 | logger := log.WithFields(log.Fields{}) 63 | mockSQS := &mockSQS{} 64 | config := WorkerConfig{ 65 | HTTPURL: ts.URL, 66 | HTTPContentType: "application/json", 67 | } 68 | 69 | mockSQS.receiveMessageFunc = func(*sqs.ReceiveMessageInput) (*sqs.ReceiveMessageOutput, error) { 70 | return &sqs.ReceiveMessageOutput{ 71 | Messages: []*sqs.Message{{ 72 | Body: aws.String("message 1"), 73 | MessageId: aws.String("m1"), 74 | ReceiptHandle: aws.String("r1"), 75 | }, { 76 | Body: aws.String("message 2"), 77 | MessageId: aws.String("m2"), 78 | ReceiptHandle: aws.String("r2"), 79 | }, { 80 | Body: aws.String("message 3"), 81 | MessageId: aws.String("m3"), 82 | ReceiptHandle: aws.String("r3"), 83 | }}, 84 | }, nil 85 | } 86 | 87 | supervisor := NewSupervisor(logger, mockSQS, &http.Client{}, config) 88 | 89 | mockSQS.deleteMessageBatchFunc = func(input *sqs.DeleteMessageBatchInput) (*sqs.DeleteMessageBatchOutput, error) { 90 | defer supervisor.Shutdown() 91 | 92 | assert.Len(t, input.Entries, 3) 93 | 94 | return nil, nil 95 | } 96 | 97 | mockSQS.changeMessageVisibilityBatchFunc = func(input *sqs.ChangeMessageVisibilityBatchInput) (*sqs.ChangeMessageVisibilityBatchOutput, error) { 98 | assert.Fail(t, "ChangeMessageVisibilityBatchFunc was called") 99 | return nil, nil 100 | } 101 | 102 | supervisor.Start(1) 103 | supervisor.Wait() 104 | } 105 | 106 | func TestSupervisorHTTPError(t *testing.T) { 107 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 108 | w.WriteHeader(http.StatusInternalServerError) 109 | })) 110 | defer ts.Close() 111 | 112 | log.SetOutput(ioutil.Discard) 113 | logger := log.WithFields(log.Fields{}) 114 | mockSQS := &mockSQS{} 115 | config := WorkerConfig{ 116 | HTTPURL: ts.URL, 117 | } 118 | 119 | supervisor := NewSupervisor(logger, mockSQS, &http.Client{}, config) 120 | 121 | receiveCount := 0 122 | mockSQS.receiveMessageFunc = func(*sqs.ReceiveMessageInput) (*sqs.ReceiveMessageOutput, error) { 123 | receiveCount++ 124 | 125 | if receiveCount == 2 { 126 | supervisor.Shutdown() 127 | 128 | return &sqs.ReceiveMessageOutput{ 129 | Messages: []*sqs.Message{}, 130 | }, nil 131 | } 132 | 133 | return &sqs.ReceiveMessageOutput{ 134 | Messages: []*sqs.Message{{ 135 | Body: aws.String("message 1"), 136 | MessageId: aws.String("m1"), 137 | ReceiptHandle: aws.String("r1"), 138 | }, { 139 | Body: aws.String("message 2"), 140 | MessageId: aws.String("m2"), 141 | ReceiptHandle: aws.String("r2"), 142 | }, { 143 | Body: aws.String("message 3"), 144 | MessageId: aws.String("m3"), 145 | ReceiptHandle: aws.String("r3"), 146 | }}, 147 | }, nil 148 | } 149 | 150 | mockSQS.deleteMessageBatchFunc = func(input *sqs.DeleteMessageBatchInput) (*sqs.DeleteMessageBatchOutput, error) { 151 | assert.Fail(t, "DeleteMessageBatchInput was called") 152 | 153 | return nil, nil 154 | } 155 | 156 | mockSQS.changeMessageVisibilityBatchFunc = func(input *sqs.ChangeMessageVisibilityBatchInput) (*sqs.ChangeMessageVisibilityBatchOutput, error) { 157 | assert.Fail(t, "ChangeMessageVisibilityBatchFunc was called") 158 | return nil, nil 159 | } 160 | 161 | supervisor.Start(1) 162 | supervisor.Wait() 163 | } 164 | 165 | func TestSupervisorHMAC(t *testing.T) { 166 | hmacHeader := "hmac" 167 | hmacSecretKey := []byte("foobar") 168 | hmacSuccess := false 169 | 170 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 171 | mac := hmac.New(sha256.New, hmacSecretKey) 172 | 173 | body, _ := ioutil.ReadAll(r.Body) 174 | r.Body.Close() 175 | 176 | mac.Write([]byte(fmt.Sprintf("%s %s\n%s", r.Method, fmt.Sprintf("http://%s", r.Host), string(body)))) 177 | expectedMAC := hex.EncodeToString(mac.Sum(nil)) 178 | 179 | hmacSuccess = hmac.Equal([]byte(r.Header.Get(hmacHeader)), []byte(expectedMAC)) 180 | })) 181 | defer ts.Close() 182 | 183 | log.SetOutput(ioutil.Discard) 184 | logger := log.WithFields(log.Fields{}) 185 | mockSQS := &mockSQS{} 186 | config := WorkerConfig{ 187 | HTTPURL: ts.URL, 188 | 189 | HTTPHMACHeader: hmacHeader, 190 | HMACSecretKey: hmacSecretKey, 191 | } 192 | 193 | supervisor := NewSupervisor(logger, mockSQS, &http.Client{}, config) 194 | 195 | mockSQS.receiveMessageFunc = func(*sqs.ReceiveMessageInput) (*sqs.ReceiveMessageOutput, error) { 196 | defer supervisor.Shutdown() 197 | 198 | return &sqs.ReceiveMessageOutput{ 199 | Messages: []*sqs.Message{{ 200 | Body: aws.String("message 1"), 201 | MessageId: aws.String("m1"), 202 | ReceiptHandle: aws.String("r1"), 203 | }}, 204 | }, nil 205 | } 206 | 207 | mockSQS.changeMessageVisibilityBatchFunc = func(input *sqs.ChangeMessageVisibilityBatchInput) (*sqs.ChangeMessageVisibilityBatchOutput, error) { 208 | assert.Fail(t, "ChangeMessageVisibilityBatchFunc was called") 209 | return nil, nil 210 | } 211 | 212 | supervisor.Start(1) 213 | supervisor.Wait() 214 | 215 | assert.True(t, hmacSuccess) 216 | } 217 | 218 | func TestSupervisorTooManyRequests(t *testing.T) { 219 | delayTime := time.Duration(1 * time.Hour) 220 | requestCount := 0 221 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 222 | requestCount++ 223 | 224 | assert.Equal(t, "application/json", r.Header.Get("Content-Type")) 225 | w.Header().Set("Retry-After", fmt.Sprintf("%v", delayTime.Seconds())) 226 | w.WriteHeader(http.StatusTooManyRequests) 227 | })) 228 | defer ts.Close() 229 | 230 | log.SetOutput(ioutil.Discard) 231 | logger := log.WithFields(log.Fields{}) 232 | mockSQS := &mockSQS{} 233 | config := WorkerConfig{ 234 | HTTPURL: ts.URL, 235 | HTTPContentType: "application/json", 236 | } 237 | 238 | mockSQS.receiveMessageFunc = func(*sqs.ReceiveMessageInput) (*sqs.ReceiveMessageOutput, error) { 239 | return &sqs.ReceiveMessageOutput{Messages: []*sqs.Message{{ 240 | Body: aws.String("message 1"), 241 | MessageId: aws.String("m1"), 242 | ReceiptHandle: aws.String("r1"), 243 | }, { 244 | Body: aws.String("message 2"), 245 | MessageId: aws.String("m2"), 246 | ReceiptHandle: aws.String("r2"), 247 | }, { 248 | Body: aws.String("message 3"), 249 | MessageId: aws.String("m3"), 250 | ReceiptHandle: aws.String("r3"), 251 | }}, 252 | }, nil 253 | } 254 | 255 | mockSQS.deleteMessageBatchFunc = func(input *sqs.DeleteMessageBatchInput) (*sqs.DeleteMessageBatchOutput, error) { 256 | assert.Fail(t, "DeleteMessageBatchFunc was called") 257 | return nil, nil 258 | } 259 | 260 | supervisor := NewSupervisor(logger, mockSQS, &http.Client{}, config) 261 | 262 | mockSQS.changeMessageVisibilityBatchFunc = func(input *sqs.ChangeMessageVisibilityBatchInput) (*sqs.ChangeMessageVisibilityBatchOutput, error) { 263 | defer supervisor.Shutdown() 264 | 265 | assert.Len(t, input.Entries, 3) 266 | for _, entry := range input.Entries { 267 | VisibilityTimeout := *entry.VisibilityTimeout 268 | timeoutDiff := int64(delayTime.Seconds()) - VisibilityTimeout 269 | assert.True(t, timeoutDiff >= 0) 270 | assert.True(t, timeoutDiff < 5) 271 | } 272 | 273 | return nil, nil 274 | } 275 | 276 | supervisor.Start(1) 277 | supervisor.Wait() 278 | } 279 | 280 | func TestSupervisorTooManyRequestsBadRetryAfter(t *testing.T) { 281 | ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 282 | assert.Equal(t, "application/json", r.Header.Get("Content-Type")) 283 | 284 | w.Header().Set("Retry-After", "invalid") 285 | w.WriteHeader(http.StatusTooManyRequests) 286 | })) 287 | defer ts.Close() 288 | 289 | log.SetOutput(ioutil.Discard) 290 | logger := log.WithFields(log.Fields{}) 291 | mockSQS := &mockSQS{} 292 | config := WorkerConfig{ 293 | HTTPURL: ts.URL, 294 | HTTPContentType: "application/json", 295 | } 296 | 297 | supervisor := NewSupervisor(logger, mockSQS, &http.Client{}, config) 298 | 299 | receiveCount := 0 300 | mockSQS.receiveMessageFunc = func(*sqs.ReceiveMessageInput) (*sqs.ReceiveMessageOutput, error) { 301 | receiveCount++ 302 | 303 | if receiveCount == 2 { 304 | supervisor.Shutdown() 305 | 306 | return &sqs.ReceiveMessageOutput{ 307 | Messages: []*sqs.Message{}, 308 | }, nil 309 | } 310 | 311 | return &sqs.ReceiveMessageOutput{ 312 | Messages: []*sqs.Message{{ 313 | Body: aws.String("message 1"), 314 | MessageId: aws.String("m1"), 315 | ReceiptHandle: aws.String("r1"), 316 | }, { 317 | Body: aws.String("message 2"), 318 | MessageId: aws.String("m2"), 319 | ReceiptHandle: aws.String("r2"), 320 | }, { 321 | Body: aws.String("message 3"), 322 | MessageId: aws.String("m3"), 323 | ReceiptHandle: aws.String("r3"), 324 | }}, 325 | }, nil 326 | } 327 | 328 | mockSQS.deleteMessageBatchFunc = func(input *sqs.DeleteMessageBatchInput) (*sqs.DeleteMessageBatchOutput, error) { 329 | assert.Fail(t, "DeleteMessageBatchInput was called") 330 | return nil, nil 331 | } 332 | 333 | mockSQS.changeMessageVisibilityBatchFunc = func(input *sqs.ChangeMessageVisibilityBatchInput) (*sqs.ChangeMessageVisibilityBatchOutput, error) { 334 | assert.Fail(t, "ChangeMessageVisibilityBatchFunc was called") 335 | return nil, nil 336 | } 337 | 338 | supervisor.Start(1) 339 | supervisor.Wait() 340 | } 341 | --------------------------------------------------------------------------------