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 //------------------------------------------------------------------------------