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

     1  package service
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  
     7  	"github.com/Jeffail/benthos/v3/internal/component/buffer"
     8  	"github.com/Jeffail/benthos/v3/internal/shutdown"
     9  	"github.com/Jeffail/benthos/v3/lib/message"
    10  	"github.com/Jeffail/benthos/v3/lib/types"
    11  )
    12  
    13  // BatchBuffer is an interface implemented by Buffers able to read and write
    14  // message batches. Buffers are a component type that are placed after inputs,
    15  // and decouples the acknowledgement system of the inputs from the rest of the
    16  // pipeline.
    17  //
    18  // Buffers are useful when implementing buffers intended to relieve back
    19  // pressure from upstream components, or when implementing message aggregators
    20  // where the concept of discrete messages running through a pipeline no longer
    21  // applies (such as with windowing algorithms).
    22  //
    23  // Buffers are advanced component types that weaken delivery guarantees of a
    24  // Benthos pipeline. Therefore, if you aren't absolutely sure that a component
    25  // you wish to build should be a buffer type then it likely shouldn't be.
    26  type BatchBuffer interface {
    27  	// Write a batch of messages to the buffer, the batch is accompanied with an
    28  	// acknowledge function. A non-nil error should be returned if it is not
    29  	// possible to store the given message batch in the buffer.
    30  	//
    31  	// If a nil error is returned the buffer assumes responsibility for calling
    32  	// the acknowledge function at least once during the lifetime of the
    33  	// message.
    34  	//
    35  	// This could be at the point where the message is written to the buffer,
    36  	// which weakens delivery guarantees but can be useful for decoupling the
    37  	// input from downstream components. Alternatively, this could be when the
    38  	// associated batch has been read from the buffer and acknowledged
    39  	// downstream, which preserves delivery guarantees.
    40  	WriteBatch(context.Context, MessageBatch, AckFunc) error
    41  
    42  	// Read a batch of messages from the buffer. This call should block until
    43  	// either a batch is ready to consume, the provided context is cancelled or
    44  	// EndOfInput has been called which indicates that the buffer is no longer
    45  	// being populated with new messages.
    46  	//
    47  	// The returned acknowledge function will be called when a consumed message
    48  	// batch has been processed and sent downstream. It is up to the buffer
    49  	// implementation whether the ack function is used, it might be used in
    50  	// order to "commit" the removal of a message from the buffer in cases where
    51  	// the buffer is a persisted storage solution, or in cases where the output
    52  	// of the buffer is temporal (a windowing algorithm, etc) it might be
    53  	// considered correct to simply drop message batches that are not acked.
    54  	//
    55  	// When the buffer is closed (EndOfInput has been called and no more
    56  	// messages are available) this method should return an ErrEndOfBuffer in
    57  	// order to indicate the end of the buffered stream.
    58  	//
    59  	// It is valid to return a batch of only one message.
    60  	ReadBatch(context.Context) (MessageBatch, AckFunc, error)
    61  
    62  	// EndOfInput indicates to the buffer that the input has ended and that once
    63  	// the buffer is depleted it should return ErrEndOfBuffer from ReadBatch in
    64  	// order to gracefully shut down the pipeline.
    65  	//
    66  	// EndOfInput should be idempotent as it may be called more than once.
    67  	EndOfInput()
    68  
    69  	Closer
    70  }
    71  
    72  //------------------------------------------------------------------------------
    73  
    74  // Implements buffer.ReaderWriter
    75  type airGapBatchBuffer struct {
    76  	b   BatchBuffer
    77  	sig *shutdown.Signaller
    78  }
    79  
    80  func newAirGapBatchBuffer(b BatchBuffer) buffer.ReaderWriter {
    81  	return &airGapBatchBuffer{b, shutdown.NewSignaller()}
    82  }
    83  
    84  func (a *airGapBatchBuffer) Write(ctx context.Context, msg types.Message, aFn buffer.AckFunc) error {
    85  	parts := make([]*Message, msg.Len())
    86  	_ = msg.Iter(func(i int, part types.Part) error {
    87  		// Copy because we ack the message after returning, therefore we lose
    88  		// ownership of the underlying.
    89  		parts[i] = newMessageFromPart(part).Copy()
    90  		return nil
    91  	})
    92  	return a.b.WriteBatch(ctx, parts, AckFunc(aFn))
    93  }
    94  
    95  func (a *airGapBatchBuffer) Read(ctx context.Context) (types.Message, buffer.AckFunc, error) {
    96  	batch, ackFn, err := a.b.ReadBatch(ctx)
    97  	if err != nil {
    98  		if errors.Is(err, ErrEndOfBuffer) {
    99  			err = types.ErrTypeClosed
   100  		}
   101  		return nil, nil, err
   102  	}
   103  	tMsg := message.New(nil)
   104  	for _, msg := range batch {
   105  		tMsg.Append(msg.part)
   106  	}
   107  	return tMsg, func(c context.Context, aerr error) error {
   108  		return ackFn(c, aerr)
   109  	}, nil
   110  }
   111  
   112  func (a *airGapBatchBuffer) EndOfInput() {
   113  	a.b.EndOfInput()
   114  }
   115  
   116  func (a *airGapBatchBuffer) Close(ctx context.Context) error {
   117  	return a.b.Close(ctx)
   118  }