github.com/Jeffail/benthos/v3@v3.65.0/lib/message/roundtrip/result_store.go (about)

     1  package roundtrip
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"sync"
     7  
     8  	"github.com/Jeffail/benthos/v3/lib/message"
     9  	"github.com/Jeffail/benthos/v3/lib/types"
    10  )
    11  
    12  //------------------------------------------------------------------------------
    13  
    14  // ErrNoStore is an error returned by components attempting to write a message
    15  // batch to a ResultStore but are unable to locate the store within the batch
    16  // context.
    17  var ErrNoStore = errors.New("result store not found within batch context")
    18  
    19  // ResultStoreKeyType is the recommended type of a context key for adding
    20  // ResultStores to a message context.
    21  type ResultStoreKeyType int
    22  
    23  // ResultStoreKey is the recommended key value for adding ResultStores to a
    24  // message context.
    25  const ResultStoreKey ResultStoreKeyType = iota
    26  
    27  // ResultStore is a type designed to be propagated along with a message as a way
    28  // for an output destination to store the final version of the message payload
    29  // as it saw it.
    30  //
    31  // It is intended that this structure is placed within a message via an attached
    32  // context, usually under the key 'result_store'.
    33  type ResultStore interface {
    34  	// Add a message to the store. The message will be deep copied and have its
    35  	// context wiped before storing, and is therefore safe to add even when
    36  	// ownership of the message is about to be yielded.
    37  	Add(msg types.Message)
    38  
    39  	// Get the stored slice of messages.
    40  	Get() []types.Message
    41  
    42  	// Clear any currently stored messages.
    43  	Clear()
    44  }
    45  
    46  //------------------------------------------------------------------------------
    47  
    48  type resultStoreImpl struct {
    49  	payloads []types.Message
    50  	sync.RWMutex
    51  }
    52  
    53  func (r *resultStoreImpl) Add(msg types.Message) {
    54  	r.Lock()
    55  	defer r.Unlock()
    56  	strippedParts := make([]types.Part, msg.Len())
    57  	msg.DeepCopy().Iter(func(i int, p types.Part) error {
    58  		strippedParts[i] = message.WithContext(context.Background(), p)
    59  		return nil
    60  	})
    61  	msg.SetAll(strippedParts)
    62  	r.payloads = append(r.payloads, msg)
    63  }
    64  
    65  func (r *resultStoreImpl) Get() []types.Message {
    66  	r.RLock()
    67  	defer r.RUnlock()
    68  	return r.payloads
    69  }
    70  
    71  func (r *resultStoreImpl) Clear() {
    72  	r.Lock()
    73  	r.payloads = nil
    74  	r.Unlock()
    75  }
    76  
    77  //------------------------------------------------------------------------------
    78  
    79  // NewResultStore returns an implementation of ResultStore.
    80  func NewResultStore() ResultStore {
    81  	return &resultStoreImpl{}
    82  }
    83  
    84  //------------------------------------------------------------------------------
    85  
    86  // AddResultStore sets a result store within the context of the provided message
    87  // that allows a roundtrip.Writer or any other component to propagate a
    88  // resulting message back to the origin.
    89  func AddResultStore(msg types.Message, store ResultStore) {
    90  	parts := make([]types.Part, msg.Len())
    91  	msg.Iter(func(i int, p types.Part) error {
    92  		ctx := message.GetContext(p)
    93  		parts[i] = message.WithContext(context.WithValue(ctx, ResultStoreKey, store), p)
    94  		return nil
    95  	})
    96  	msg.SetAll(parts)
    97  }
    98  
    99  // SetAsResponse takes a mutated message and stores it as a response message,
   100  // this action fails if the message does not contain a valid ResultStore within
   101  // its context.
   102  func SetAsResponse(msg types.Message) error {
   103  	ctx := message.GetContext(msg.Get(0))
   104  	store, ok := ctx.Value(ResultStoreKey).(ResultStore)
   105  	if !ok {
   106  		return ErrNoStore
   107  	}
   108  	store.Add(msg)
   109  	return nil
   110  }
   111  
   112  //------------------------------------------------------------------------------