github.com/Jeffail/benthos/v3@v3.65.0/lib/input/reader/preserver.go (about)

     1  package reader
     2  
     3  import (
     4  	"time"
     5  
     6  	"github.com/Jeffail/benthos/v3/lib/types"
     7  	"github.com/Jeffail/benthos/v3/lib/util/throttle"
     8  )
     9  
    10  //------------------------------------------------------------------------------
    11  
    12  // Preserver is a wrapper for reader.Type implementations that keeps a buffer of
    13  // sent messages until they are acknowledged. If an error occurs during message
    14  // propagation the contents of the buffer will be resent instead of reading new
    15  // messages until it is depleted. Preserver implements reader.Type.
    16  type Preserver struct {
    17  	unAckMessages  []types.Message
    18  	resendMessages []types.Message
    19  
    20  	throt *throttle.Type
    21  
    22  	r Type
    23  }
    24  
    25  // NewPreserver returns a new Preserver wrapper around a reader.Type.
    26  func NewPreserver(r Type) *Preserver {
    27  	return &Preserver{
    28  		r: r,
    29  		throt: throttle.New(
    30  			throttle.OptThrottlePeriod(time.Second),
    31  		),
    32  	}
    33  }
    34  
    35  //------------------------------------------------------------------------------
    36  
    37  // Connect attempts to establish a connection to the source, if unsuccessful
    38  // returns an error. If the attempt is successful (or not necessary) returns
    39  // nil.
    40  func (p *Preserver) Connect() error {
    41  	return p.r.Connect()
    42  }
    43  
    44  // Acknowledge instructs whether messages read since the last Acknowledge call
    45  // were successfully propagated. If the error is nil this will be forwarded to
    46  // the underlying wrapped reader. If a non-nil error is returned the buffer of
    47  // messages will be resent.
    48  func (p *Preserver) Acknowledge(err error) error {
    49  	if err == nil {
    50  		p.throt.Reset()
    51  		p.unAckMessages = nil
    52  		if len(p.resendMessages) == 0 {
    53  			// Only propagate ack if we are done resending buffered messages.
    54  			return p.r.Acknowledge(err)
    55  		}
    56  		return nil
    57  	}
    58  
    59  	// Do not propagate errors since we are handling them here by resending.
    60  	p.resendMessages = append(p.resendMessages, p.unAckMessages...)
    61  	p.unAckMessages = nil
    62  	p.throt.Retry()
    63  	return nil
    64  }
    65  
    66  // Read attempts to read a new message from the source.
    67  func (p *Preserver) Read() (types.Message, error) {
    68  	// If we have messages queued to be resent we prioritise them over reading
    69  	// new messages.
    70  	if lMsgs := len(p.resendMessages); lMsgs > 0 {
    71  		msg := p.resendMessages[0]
    72  		if lMsgs > 1 {
    73  			p.resendMessages = p.resendMessages[1:]
    74  		} else {
    75  			p.resendMessages = nil
    76  		}
    77  		p.unAckMessages = append(p.unAckMessages, msg)
    78  		return msg, nil
    79  	}
    80  	msg, err := p.r.Read()
    81  	if err == nil {
    82  		p.unAckMessages = append(p.unAckMessages, msg)
    83  	}
    84  	return msg, err
    85  }
    86  
    87  // CloseAsync triggers the asynchronous closing of the reader.
    88  func (p *Preserver) CloseAsync() {
    89  	p.r.CloseAsync()
    90  }
    91  
    92  // WaitForClose blocks until either the reader is finished closing or a timeout
    93  // occurs.
    94  func (p *Preserver) WaitForClose(tout time.Duration) error {
    95  	return p.r.WaitForClose(tout)
    96  }
    97  
    98  //------------------------------------------------------------------------------