github.com/Jeffail/benthos/v3@v3.65.0/internal/batch/count.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 batchedCountKeyType int
    11  
    12  const batchedCountKey batchedCountKeyType = iota
    13  
    14  // CollapsedCount attempts to extract the actual number of messages that were
    15  // collapsed into the resulting message part. This value could be greater than 1
    16  // when users configure processors that archive batched message parts.
    17  func CollapsedCount(p types.Part) int {
    18  	if v, ok := message.GetContext(p).Value(batchedCountKey).(int); ok {
    19  		return v
    20  	}
    21  	return 1
    22  }
    23  
    24  // MessageCollapsedCount attempts to extract the actual number of messages that
    25  // were combined into the resulting batched message parts. This value could
    26  // differ from message.Len() when users configure processors that archive
    27  // batched message parts.
    28  func MessageCollapsedCount(m types.Message) int {
    29  	total := 0
    30  	m.Iter(func(i int, p types.Part) error {
    31  		total += CollapsedCount(p)
    32  		return nil
    33  	})
    34  	return total
    35  }
    36  
    37  // WithCollapsedCount returns a message part with a context indicating that this
    38  // message is the result of collapsing a number of messages. This allows
    39  // downstream components to know how many total messages were combined.
    40  func WithCollapsedCount(p types.Part, count int) types.Part {
    41  	// Start with the previous length which could also be >1.
    42  	ctx := message.GetContext(p)
    43  	base := 1
    44  	if v, ok := ctx.Value(batchedCountKey).(int); ok {
    45  		base = v
    46  	}
    47  
    48  	// The new length is the previous length plus the total messages put into
    49  	// this batch (minus one to prevent double counting the original part).
    50  	ctx = context.WithValue(ctx, batchedCountKey, base-1+count)
    51  
    52  	return message.WithContext(ctx, p)
    53  }