github.com/Jeffail/benthos/v3@v3.65.0/lib/input/aws_sqs.go (about) 1 package input 2 3 import ( 4 "context" 5 "sync" 6 "time" 7 8 "github.com/Jeffail/benthos/v3/internal/docs" 9 "github.com/Jeffail/benthos/v3/internal/shutdown" 10 "github.com/Jeffail/benthos/v3/lib/input/reader" 11 "github.com/Jeffail/benthos/v3/lib/log" 12 "github.com/Jeffail/benthos/v3/lib/message" 13 "github.com/Jeffail/benthos/v3/lib/metrics" 14 "github.com/Jeffail/benthos/v3/lib/types" 15 sess "github.com/Jeffail/benthos/v3/lib/util/aws/session" 16 "github.com/aws/aws-sdk-go/aws" 17 "github.com/aws/aws-sdk-go/aws/awserr" 18 "github.com/aws/aws-sdk-go/aws/request" 19 "github.com/aws/aws-sdk-go/aws/session" 20 "github.com/aws/aws-sdk-go/service/sqs" 21 "github.com/cenkalti/backoff/v4" 22 ) 23 24 //------------------------------------------------------------------------------ 25 26 func init() { 27 Constructors[TypeAWSSQS] = TypeSpec{ 28 Status: docs.StatusStable, 29 constructor: fromSimpleConstructor(func(conf Config, mgr types.Manager, log log.Modular, stats metrics.Type) (Type, error) { 30 r, err := newAWSSQS(conf.AWSSQS, log, stats) 31 if err != nil { 32 return nil, err 33 } 34 return NewAsyncReader(TypeAWSSQS, false, r, log, stats) 35 }), 36 Summary: ` 37 Consume messages from an AWS SQS URL.`, 38 Description: ` 39 ### Credentials 40 41 By default Benthos will use a shared credentials file when connecting to AWS 42 services. It's also possible to set them explicitly at the component level, 43 allowing you to transfer data across accounts. You can find out more 44 [in this document](/docs/guides/cloud/aws). 45 46 ### Metadata 47 48 This input adds the following metadata fields to each message: 49 50 ` + "```text" + ` 51 - sqs_message_id 52 - sqs_receipt_handle 53 - sqs_approximate_receive_count 54 - All message attributes 55 ` + "```" + ` 56 57 You can access these metadata fields using 58 [function interpolation](/docs/configuration/interpolation#metadata).`, 59 FieldSpecs: append(docs.FieldSpecs{ 60 docs.FieldCommon("url", "The SQS URL to consume from."), 61 docs.FieldAdvanced("delete_message", "Whether to delete the consumed message once it is acked. Disabling allows you to handle the deletion using a different mechanism."), 62 docs.FieldAdvanced("reset_visibility", "Whether to set the visibility timeout of the consumed message to zero once it is nacked. Disabling honors the preset visibility timeout specified for the queue.").AtVersion("3.58.0"), 63 }, sess.FieldSpecs()...), 64 Categories: []Category{ 65 CategoryServices, 66 CategoryAWS, 67 }, 68 } 69 } 70 71 //------------------------------------------------------------------------------ 72 73 // AWSSQSConfig contains configuration values for the input type. 74 type AWSSQSConfig struct { 75 sess.Config `json:",inline" yaml:",inline"` 76 URL string `json:"url" yaml:"url"` 77 DeleteMessage bool `json:"delete_message" yaml:"delete_message"` 78 ResetVisibility bool `json:"reset_visibility" yaml:"reset_visibility"` 79 } 80 81 // NewAWSSQSConfig creates a new Config with default values. 82 func NewAWSSQSConfig() AWSSQSConfig { 83 return AWSSQSConfig{ 84 Config: sess.NewConfig(), 85 URL: "", 86 DeleteMessage: true, 87 ResetVisibility: true, 88 } 89 } 90 91 //------------------------------------------------------------------------------ 92 93 type awsSQS struct { 94 conf AWSSQSConfig 95 96 session *session.Session 97 sqs *sqs.SQS 98 99 messagesChan chan *sqs.Message 100 ackMessagesChan chan sqsMessageHandle 101 nackMessagesChan chan sqsMessageHandle 102 closeSignal *shutdown.Signaller 103 104 log log.Modular 105 stats metrics.Type 106 } 107 108 func newAWSSQS(conf AWSSQSConfig, log log.Modular, stats metrics.Type) (*awsSQS, error) { 109 return &awsSQS{ 110 conf: conf, 111 log: log, 112 stats: stats, 113 messagesChan: make(chan *sqs.Message), 114 ackMessagesChan: make(chan sqsMessageHandle), 115 nackMessagesChan: make(chan sqsMessageHandle), 116 closeSignal: shutdown.NewSignaller(), 117 }, nil 118 } 119 120 // ConnectWithContext attempts to establish a connection to the target SQS 121 // queue. 122 func (a *awsSQS) ConnectWithContext(ctx context.Context) error { 123 if a.session != nil { 124 return nil 125 } 126 127 sess, err := a.conf.GetSession() 128 if err != nil { 129 return err 130 } 131 132 a.sqs = sqs.New(sess) 133 a.session = sess 134 135 var wg sync.WaitGroup 136 wg.Add(2) 137 go a.readLoop(&wg) 138 go a.ackLoop(&wg) 139 go func() { 140 wg.Wait() 141 a.closeSignal.ShutdownComplete() 142 }() 143 144 a.log.Infof("Receiving Amazon SQS messages from URL: %v\n", a.conf.URL) 145 return nil 146 } 147 148 func (a *awsSQS) ackLoop(wg *sync.WaitGroup) { 149 defer wg.Done() 150 151 var pendingAcks []sqsMessageHandle 152 var pendingNacks []sqsMessageHandle 153 154 flushAcks := func() { 155 tmpAcks := pendingAcks 156 pendingAcks = nil 157 if len(tmpAcks) == 0 { 158 return 159 } 160 161 ctx, done := a.closeSignal.CloseNowCtx(context.Background()) 162 defer done() 163 if err := a.deleteMessages(ctx, tmpAcks...); err != nil { 164 a.log.Errorf("Failed to delete messages: %v", err) 165 } 166 } 167 168 flushNacks := func() { 169 tmpNacks := pendingNacks 170 pendingNacks = nil 171 if len(tmpNacks) == 0 { 172 return 173 } 174 175 ctx, done := a.closeSignal.CloseNowCtx(context.Background()) 176 defer done() 177 if err := a.resetMessages(ctx, tmpNacks...); err != nil { 178 a.log.Errorf("Failed to reset the visibility timeout of messages: %v", err) 179 } 180 } 181 182 flushTimer := time.NewTicker(time.Second) 183 defer flushTimer.Stop() 184 185 ackLoop: 186 for { 187 select { 188 case h := <-a.ackMessagesChan: 189 pendingAcks = append(pendingAcks, h) 190 if len(pendingAcks) >= 10 { 191 flushAcks() 192 } 193 case h := <-a.nackMessagesChan: 194 pendingNacks = append(pendingNacks, h) 195 if len(pendingNacks) >= 10 { 196 flushNacks() 197 } 198 case <-flushTimer.C: 199 flushAcks() 200 flushNacks() 201 case <-a.closeSignal.CloseAtLeisureChan(): 202 break ackLoop 203 } 204 } 205 206 flushAcks() 207 flushNacks() 208 } 209 210 func (a *awsSQS) readLoop(wg *sync.WaitGroup) { 211 defer wg.Done() 212 213 var pendingMsgs []*sqs.Message 214 defer func() { 215 if len(pendingMsgs) > 0 { 216 tmpNacks := make([]sqsMessageHandle, 0, len(pendingMsgs)) 217 for _, m := range pendingMsgs { 218 if m.MessageId == nil || m.ReceiptHandle == nil { 219 continue 220 } 221 tmpNacks = append(tmpNacks, sqsMessageHandle{ 222 id: *m.MessageId, 223 receiptHandle: *m.ReceiptHandle, 224 }) 225 } 226 ctx, done := a.closeSignal.CloseNowCtx(context.Background()) 227 defer done() 228 if err := a.resetMessages(ctx, tmpNacks...); err != nil { 229 a.log.Errorf("Failed to reset visibility timeout for pending messages: %v", err) 230 } 231 } 232 }() 233 234 backoff := backoff.NewExponentialBackOff() 235 backoff.InitialInterval = time.Millisecond 236 backoff.MaxInterval = time.Second 237 238 getMsgs := func() { 239 ctx, done := a.closeSignal.CloseAtLeisureCtx(context.Background()) 240 defer done() 241 res, err := a.sqs.ReceiveMessageWithContext(ctx, &sqs.ReceiveMessageInput{ 242 QueueUrl: aws.String(a.conf.URL), 243 MaxNumberOfMessages: aws.Int64(10), 244 AttributeNames: []*string{aws.String("All")}, 245 MessageAttributeNames: []*string{aws.String("All")}, 246 }) 247 if err != nil { 248 if aerr, ok := err.(awserr.Error); !ok || aerr.Code() != request.CanceledErrorCode { 249 a.log.Errorf("Failed to pull new SQS messages: %v", aerr) 250 } 251 return 252 } 253 if len(res.Messages) > 0 { 254 pendingMsgs = append(pendingMsgs, res.Messages...) 255 backoff.Reset() 256 } 257 } 258 259 for { 260 if len(pendingMsgs) == 0 { 261 getMsgs() 262 if len(pendingMsgs) == 0 { 263 select { 264 case <-time.After(backoff.NextBackOff()): 265 case <-a.closeSignal.CloseAtLeisureChan(): 266 return 267 } 268 continue 269 } 270 } 271 select { 272 case a.messagesChan <- pendingMsgs[0]: 273 pendingMsgs = pendingMsgs[1:] 274 case <-a.closeSignal.CloseAtLeisureChan(): 275 return 276 } 277 } 278 } 279 280 type sqsMessageHandle struct { 281 id, receiptHandle string 282 } 283 284 func (a *awsSQS) deleteMessages(ctx context.Context, msgs ...sqsMessageHandle) error { 285 for len(msgs) > 0 { 286 input := sqs.DeleteMessageBatchInput{ 287 QueueUrl: aws.String(a.conf.URL), 288 Entries: []*sqs.DeleteMessageBatchRequestEntry{}, 289 } 290 291 for _, msg := range msgs { 292 input.Entries = append(input.Entries, &sqs.DeleteMessageBatchRequestEntry{ 293 Id: aws.String(msg.id), 294 ReceiptHandle: aws.String(msg.receiptHandle), 295 }) 296 if len(input.Entries) == 10 { 297 break 298 } 299 } 300 301 msgs = msgs[len(input.Entries):] 302 response, err := a.sqs.DeleteMessageBatchWithContext(ctx, &input) 303 if err != nil { 304 return err 305 } 306 for _, fail := range response.Failed { 307 a.log.Errorf("Failed to delete consumed SQS message '%v', response code: %v\n", *fail.Id, *fail.Code) 308 } 309 } 310 return nil 311 } 312 313 func (a *awsSQS) resetMessages(ctx context.Context, msgs ...sqsMessageHandle) error { 314 for len(msgs) > 0 { 315 input := sqs.ChangeMessageVisibilityBatchInput{ 316 QueueUrl: aws.String(a.conf.URL), 317 Entries: []*sqs.ChangeMessageVisibilityBatchRequestEntry{}, 318 } 319 320 for _, msg := range msgs { 321 input.Entries = append(input.Entries, &sqs.ChangeMessageVisibilityBatchRequestEntry{ 322 Id: aws.String(msg.id), 323 ReceiptHandle: aws.String(msg.receiptHandle), 324 VisibilityTimeout: aws.Int64(0), 325 }) 326 if len(input.Entries) == 10 { 327 break 328 } 329 } 330 331 msgs = msgs[len(input.Entries):] 332 response, err := a.sqs.ChangeMessageVisibilityBatchWithContext(ctx, &input) 333 if err != nil { 334 return err 335 } 336 for _, fail := range response.Failed { 337 a.log.Errorf("Failed to delete consumed SQS message '%v', response code: %v\n", *fail.Id, *fail.Code) 338 } 339 } 340 return nil 341 } 342 343 func addSQSMetadata(p types.Part, sqsMsg *sqs.Message) { 344 meta := p.Metadata() 345 meta.Set("sqs_message_id", *sqsMsg.MessageId) 346 meta.Set("sqs_receipt_handle", *sqsMsg.ReceiptHandle) 347 if rCountStr := sqsMsg.Attributes["ApproximateReceiveCount"]; rCountStr != nil { 348 meta.Set("sqs_approximate_receive_count", *rCountStr) 349 } 350 for k, v := range sqsMsg.MessageAttributes { 351 if v.StringValue != nil { 352 meta.Set(k, *v.StringValue) 353 } 354 } 355 } 356 357 // ReadWithContext attempts to read a new message from the target SQS. 358 func (a *awsSQS) ReadWithContext(ctx context.Context) (types.Message, reader.AsyncAckFn, error) { 359 if a.session == nil { 360 return nil, nil, types.ErrNotConnected 361 } 362 363 var next *sqs.Message 364 var open bool 365 select { 366 case next, open = <-a.messagesChan: 367 if !open { 368 return nil, nil, types.ErrTypeClosed 369 } 370 case <-a.closeSignal.CloseAtLeisureChan(): 371 return nil, nil, types.ErrTypeClosed 372 case <-ctx.Done(): 373 return nil, nil, ctx.Err() 374 } 375 376 msg := message.New(nil) 377 if next.Body != nil { 378 part := message.NewPart([]byte(*next.Body)) 379 addSQSMetadata(part, next) 380 msg.Append(part) 381 } 382 if msg.Len() == 0 { 383 return nil, nil, types.ErrTimeout 384 } 385 386 mHandle := sqsMessageHandle{ 387 id: *next.MessageId, 388 } 389 if next.ReceiptHandle != nil { 390 mHandle.receiptHandle = *next.ReceiptHandle 391 } 392 return msg, func(rctx context.Context, res types.Response) error { 393 if mHandle.receiptHandle == "" { 394 return nil 395 } 396 397 if res.Error() == nil { 398 if !a.conf.DeleteMessage { 399 return nil 400 } 401 select { 402 case <-rctx.Done(): 403 return rctx.Err() 404 case <-a.closeSignal.CloseAtLeisureChan(): 405 return a.deleteMessages(rctx, mHandle) 406 case a.ackMessagesChan <- mHandle: 407 } 408 return nil 409 } 410 411 if !a.conf.ResetVisibility { 412 return nil 413 } 414 select { 415 case <-rctx.Done(): 416 return rctx.Err() 417 case <-a.closeSignal.CloseAtLeisureChan(): 418 return a.resetMessages(rctx, mHandle) 419 case a.nackMessagesChan <- mHandle: 420 } 421 return nil 422 }, nil 423 } 424 425 // CloseAsync begins cleaning up resources used by this reader asynchronously. 426 func (a *awsSQS) CloseAsync() { 427 a.closeSignal.CloseAtLeisure() 428 } 429 430 // WaitForClose will block until either the reader is closed or a specified 431 // timeout occurs. 432 func (a *awsSQS) WaitForClose(tout time.Duration) error { 433 go func() { 434 closeNowAt := tout - time.Second 435 if closeNowAt < time.Second { 436 closeNowAt = time.Second 437 } 438 <-time.After(closeNowAt) 439 a.closeSignal.CloseNow() 440 }() 441 select { 442 case <-time.After(tout): 443 return types.ErrTimeout 444 case <-a.closeSignal.HasClosedChan(): 445 return nil 446 } 447 }