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

     1  package reader
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/Jeffail/benthos/v3/lib/types"
    10  )
    11  
    12  //------------------------------------------------------------------------------
    13  
    14  // AsyncBundleUnacks is a wrapper for reader.Async implementations that,
    15  // whenever an unack is given as a response to an async message, keeps the
    16  // underlying ack function buffered. Once a non-unack response is received it is
    17  // provided to all buffered ack functions.
    18  //
    19  // Unacks are only returned by the batch processor, and once it is removed this
    20  // component can also be removed.
    21  //
    22  // TODO: V4 Remove this.
    23  type AsyncBundleUnacks struct {
    24  	pendingAcks []AsyncAckFn
    25  	acksMut     sync.RWMutex
    26  
    27  	r Async
    28  }
    29  
    30  // NewAsyncBundleUnacks returns a new AsyncBundleUnacks wrapper around a
    31  // reader.Async.
    32  func NewAsyncBundleUnacks(r Async) *AsyncBundleUnacks {
    33  	return &AsyncBundleUnacks{
    34  		r: r,
    35  	}
    36  }
    37  
    38  //------------------------------------------------------------------------------
    39  
    40  // ConnectWithContext attempts to establish a connection to the source, if
    41  // unsuccessful returns an error. If the attempt is successful (or not
    42  // necessary) returns nil.
    43  func (p *AsyncBundleUnacks) ConnectWithContext(ctx context.Context) error {
    44  	return p.r.ConnectWithContext(ctx)
    45  }
    46  
    47  func (p *AsyncBundleUnacks) wrapAckFn(ackFn AsyncAckFn) AsyncAckFn {
    48  	p.acksMut.Lock()
    49  	accumulatedAcks := p.pendingAcks
    50  	p.pendingAcks = nil
    51  	p.acksMut.Unlock()
    52  	return func(ctx context.Context, res types.Response) error {
    53  		if res.Error() == nil && res.SkipAck() {
    54  			p.acksMut.Lock()
    55  			p.pendingAcks = append(p.pendingAcks, accumulatedAcks...)
    56  			p.pendingAcks = append(p.pendingAcks, ackFn)
    57  			p.acksMut.Unlock()
    58  			return nil
    59  		}
    60  
    61  		nPendingAcks := len(accumulatedAcks)
    62  		if nPendingAcks == 0 {
    63  			return ackFn(ctx, res)
    64  		}
    65  
    66  		errs := []error{}
    67  		for _, aFn := range accumulatedAcks {
    68  			if err := aFn(ctx, res); err != nil {
    69  				errs = append(errs, err)
    70  			}
    71  		}
    72  		if err := ackFn(ctx, res); err != nil {
    73  			errs = append(errs, err)
    74  		}
    75  		if len(errs) == 1 {
    76  			return errs[0]
    77  		}
    78  		if len(errs) > 0 {
    79  			return fmt.Errorf("failed to send grouped acknowledgements: %s", errs)
    80  		}
    81  		return nil
    82  	}
    83  }
    84  
    85  // ReadWithContext attempts to read a new message from the source.
    86  func (p *AsyncBundleUnacks) ReadWithContext(ctx context.Context) (types.Message, AsyncAckFn, error) {
    87  	msg, aFn, err := p.r.ReadWithContext(ctx)
    88  	if err != nil {
    89  		return nil, nil, err
    90  	}
    91  	return msg, p.wrapAckFn(aFn), nil
    92  }
    93  
    94  // CloseAsync triggers the asynchronous closing of the reader.
    95  func (p *AsyncBundleUnacks) CloseAsync() {
    96  	p.r.CloseAsync()
    97  }
    98  
    99  // WaitForClose blocks until either the reader is finished closing or a timeout
   100  // occurs.
   101  func (p *AsyncBundleUnacks) WaitForClose(tout time.Duration) error {
   102  	return p.r.WaitForClose(tout)
   103  }
   104  
   105  //------------------------------------------------------------------------------