github.com/Jeffail/benthos/v3@v3.65.0/public/service/input_auto_retry.go (about)

     1  package service
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/cenkalti/backoff/v4"
     9  )
    10  
    11  // AutoRetryNacks wraps an input implementation with a component that
    12  // automatically reattempts messages that fail downstream. This is useful for
    13  // inputs that do not support nacks, and therefore don't have an answer for
    14  // when an ack func is called with an error.
    15  //
    16  // When messages fail to be delivered they will be reattempted with back off
    17  // until success or the stream is stopped.
    18  func AutoRetryNacks(i Input) Input {
    19  	return &autoRetryInput{
    20  		child:           i,
    21  		resendInterrupt: func() {},
    22  	}
    23  }
    24  
    25  //------------------------------------------------------------------------------
    26  
    27  type messageRetry struct {
    28  	boff     backoff.BackOff
    29  	attempts int
    30  	msg      *Message
    31  	ackFn    AckFunc
    32  }
    33  
    34  func newMessageRetry(msg *Message, ackFn AckFunc) messageRetry {
    35  	boff := backoff.NewExponentialBackOff()
    36  	boff.InitialInterval = time.Millisecond
    37  	boff.MaxInterval = time.Second
    38  	boff.Multiplier = 1.1
    39  	boff.MaxElapsedTime = 0
    40  	return messageRetry{boff, 0, msg, ackFn}
    41  }
    42  
    43  type autoRetryInput struct {
    44  	resendMessages  []messageRetry
    45  	resendInterrupt func()
    46  	msgsMut         sync.Mutex
    47  
    48  	child Input
    49  }
    50  
    51  func (i *autoRetryInput) Connect(ctx context.Context) error {
    52  	return i.child.Connect(ctx)
    53  }
    54  
    55  func (i *autoRetryInput) wrapAckFunc(m messageRetry) (*Message, AckFunc) {
    56  	return m.msg, func(ctx context.Context, err error) error {
    57  		if err != nil {
    58  			i.msgsMut.Lock()
    59  			i.resendMessages = append(i.resendMessages, m)
    60  			i.resendInterrupt()
    61  			i.msgsMut.Unlock()
    62  			return nil
    63  		}
    64  		return m.ackFn(ctx, nil)
    65  	}
    66  }
    67  
    68  func (i *autoRetryInput) Read(ctx context.Context) (*Message, AckFunc, error) {
    69  	var cancel func()
    70  	ctx, cancel = context.WithCancel(ctx)
    71  	defer cancel()
    72  
    73  	// If we have messages queued to be resent we prioritise them over reading
    74  	// new messages.
    75  	i.msgsMut.Lock()
    76  	if lMsgs := len(i.resendMessages); lMsgs > 0 {
    77  		resend := i.resendMessages[0]
    78  		if lMsgs > 1 {
    79  			i.resendMessages = i.resendMessages[1:]
    80  		} else {
    81  			i.resendMessages = nil
    82  		}
    83  		i.msgsMut.Unlock()
    84  
    85  		resend.attempts++
    86  		if resend.attempts > 2 {
    87  			// This sleep prevents a busy loop on permanently failed messages.
    88  			if tout := resend.boff.NextBackOff(); tout > 0 {
    89  				select {
    90  				case <-time.After(tout):
    91  				case <-ctx.Done():
    92  					return nil, nil, ctx.Err()
    93  				}
    94  			}
    95  		}
    96  		sendMsg, ackFn := i.wrapAckFunc(resend)
    97  		return sendMsg, ackFn, nil
    98  	}
    99  	i.resendInterrupt = cancel
   100  	i.msgsMut.Unlock()
   101  
   102  	msg, aFn, err := i.child.Read(ctx)
   103  	if err != nil {
   104  		return nil, nil, err
   105  	}
   106  	sendMsg, ackFn := i.wrapAckFunc(newMessageRetry(msg, aFn))
   107  	return sendMsg, ackFn, nil
   108  }
   109  
   110  func (i *autoRetryInput) Close(ctx context.Context) error {
   111  	return i.child.Close(ctx)
   112  }