github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/clients/pkg/promtail/client/batch.go (about)

     1  package client
     2  
     3  import (
     4  	"fmt"
     5  	"sort"
     6  	"strings"
     7  	"time"
     8  
     9  	"github.com/gogo/protobuf/proto"
    10  	"github.com/golang/snappy"
    11  	"github.com/prometheus/common/model"
    12  
    13  	"github.com/grafana/loki/clients/pkg/promtail/api"
    14  
    15  	"github.com/grafana/loki/pkg/logproto"
    16  )
    17  
    18  // batch holds pending log streams waiting to be sent to Loki, and it's used
    19  // to reduce the number of push requests to Loki aggregating multiple log streams
    20  // and entries in a single batch request. In case of multi-tenant Promtail, log
    21  // streams for each tenant are stored in a dedicated batch.
    22  type batch struct {
    23  	streams   map[string]*logproto.Stream
    24  	bytes     int
    25  	createdAt time.Time
    26  }
    27  
    28  func newBatch(entries ...api.Entry) *batch {
    29  	b := &batch{
    30  		streams:   map[string]*logproto.Stream{},
    31  		bytes:     0,
    32  		createdAt: time.Now(),
    33  	}
    34  
    35  	// Add entries to the batch
    36  	for _, entry := range entries {
    37  		b.add(entry)
    38  	}
    39  
    40  	return b
    41  }
    42  
    43  // add an entry to the batch
    44  func (b *batch) add(entry api.Entry) {
    45  	b.bytes += len(entry.Line)
    46  
    47  	// Append the entry to an already existing stream (if any)
    48  	labels := labelsMapToString(entry.Labels, ReservedLabelTenantID)
    49  	if stream, ok := b.streams[labels]; ok {
    50  		stream.Entries = append(stream.Entries, entry.Entry)
    51  		return
    52  	}
    53  
    54  	// Add the entry as a new stream
    55  	b.streams[labels] = &logproto.Stream{
    56  		Labels:  labels,
    57  		Entries: []logproto.Entry{entry.Entry},
    58  	}
    59  }
    60  
    61  func labelsMapToString(ls model.LabelSet, without ...model.LabelName) string {
    62  	lstrs := make([]string, 0, len(ls))
    63  Outer:
    64  	for l, v := range ls {
    65  		for _, w := range without {
    66  			if l == w {
    67  				continue Outer
    68  			}
    69  		}
    70  		lstrs = append(lstrs, fmt.Sprintf("%s=%q", l, v))
    71  	}
    72  
    73  	sort.Strings(lstrs)
    74  	return fmt.Sprintf("{%s}", strings.Join(lstrs, ", "))
    75  }
    76  
    77  // sizeBytes returns the current batch size in bytes
    78  func (b *batch) sizeBytes() int {
    79  	return b.bytes
    80  }
    81  
    82  // sizeBytesAfter returns the size of the batch after the input entry
    83  // will be added to the batch itself
    84  func (b *batch) sizeBytesAfter(entry api.Entry) int {
    85  	return b.bytes + len(entry.Line)
    86  }
    87  
    88  // age of the batch since its creation
    89  func (b *batch) age() time.Duration {
    90  	return time.Since(b.createdAt)
    91  }
    92  
    93  // encode the batch as snappy-compressed push request, and returns
    94  // the encoded bytes and the number of encoded entries
    95  func (b *batch) encode() ([]byte, int, error) {
    96  	req, entriesCount := b.createPushRequest()
    97  	buf, err := proto.Marshal(req)
    98  	if err != nil {
    99  		return nil, 0, err
   100  	}
   101  	buf = snappy.Encode(nil, buf)
   102  	return buf, entriesCount, nil
   103  }
   104  
   105  // creates push request and returns it, together with number of entries
   106  func (b *batch) createPushRequest() (*logproto.PushRequest, int) {
   107  	req := logproto.PushRequest{
   108  		Streams: make([]logproto.Stream, 0, len(b.streams)),
   109  	}
   110  
   111  	entriesCount := 0
   112  	for _, stream := range b.streams {
   113  		req.Streams = append(req.Streams, *stream)
   114  		entriesCount += len(stream.Entries)
   115  	}
   116  	return &req, entriesCount
   117  }