github.com/Jeffail/benthos/v3@v3.65.0/internal/batch/tracker.go (about)

     1  package batch
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/Jeffail/benthos/v3/lib/message"
     7  	"github.com/Jeffail/benthos/v3/lib/types"
     8  )
     9  
    10  type tag struct {
    11  }
    12  
    13  type tagType *tag
    14  
    15  type tagKeyType int
    16  
    17  const tagKey tagKeyType = iota
    18  
    19  type tagChecker interface {
    20  	HasTag(t tagType) bool
    21  }
    22  
    23  type tagValue struct {
    24  	tag      tagType
    25  	previous tagChecker
    26  }
    27  
    28  func (t tagValue) HasTag(tag tagType) bool {
    29  	if t.tag == tag {
    30  		return true
    31  	}
    32  	if t.previous != nil {
    33  		return t.previous.HasTag(tag)
    34  	}
    35  	return false
    36  }
    37  
    38  func hasTag(p types.Part, tag tagType) bool {
    39  	ctx := message.GetContext(p)
    40  
    41  	v, ok := ctx.Value(tagKey).(tagChecker)
    42  	if !ok {
    43  		return false
    44  	}
    45  
    46  	return v.HasTag(tag)
    47  }
    48  
    49  func withTag(tag tagType, p types.Part) types.Part {
    50  	ctx := message.GetContext(p)
    51  
    52  	var prev tagChecker
    53  	if v, ok := ctx.Value(tagKey).(tagChecker); ok {
    54  		prev = v
    55  	}
    56  
    57  	ctx = context.WithValue(ctx, tagKey, tagValue{
    58  		tag:      tag,
    59  		previous: prev,
    60  	})
    61  
    62  	return message.WithContext(ctx, p)
    63  }