chore: refactor loki client code (#4840)

* move metrics to it's own file

* remove unused function StopNow

* Remove unused Name() function

* move MaxStreams to config

* move wal writer into client manager

* refactor batch to now have two different methods.

Author: kalleep

internal/component/common/loki/client/batch.go

@@ -54,44 +54,21 @@ func newBatch(maxStreams int, entries ...loki.Entry) *batch {
 	// Add entries to the batch
 	for _, entry := range entries {
 		//never error here
-		_ = b.add(entry)
+		_ = b.add(entry, 0)
 	}
 
 	return b
 }
 
-// add an entry to the batch
-func (b *batch) add(entry loki.Entry) error {
+// add an entry to the batch. segmentNum is used to associate batch with a segment from WAL.
+// If entry is added from non backed WAL client it can be anything and is unused.
+func (b *batch) add(entry loki.Entry, segmentNum int) error {
 	b.totalBytes += entrySize(entry.Entry)
 
 	// Append the entry to an already existing stream (if any)
 	labels := labelsMapToString(entry.Labels)
 	if stream, ok := b.streams[labels]; ok {
 		stream.Entries = append(stream.Entries, entry.Entry)
-		return nil
-	}
-
-	streams := len(b.streams)
-	if b.maxStreams > 0 && streams >= b.maxStreams {
-		return fmt.Errorf("%w, streams: %d exceeds limit: %d, stream: '%s'", errMaxStreamsLimitExceeded, streams, b.maxStreams, labels)
-	}
-	// Add the entry as a new stream
-	b.streams[labels] = &push.Stream{
-		Labels:  labels,
-		Entries: []push.Entry{entry.Entry},
-	}
-	return nil
-}
-
-// addFromWAL adds an entry to the batch, tracking that the data being added comes from segment segmentNum read from the
-// WAL.
-func (b *batch) addFromWAL(lbs model.LabelSet, entry push.Entry, segmentNum int) error {
-	b.totalBytes += len(entry.Line)
-
-	// Append the entry to an already existing stream (if any)
-	labels := labelsMapToString(lbs)
-	if stream, ok := b.streams[labels]; ok {
-		stream.Entries = append(stream.Entries, entry)
 		b.countForSegment(segmentNum)
 		return nil
 	}
@@ -100,14 +77,12 @@ func (b *batch) addFromWAL(lbs model.LabelSet, entry push.Entry, segmentNum int)
 	if b.maxStreams > 0 && streams >= b.maxStreams {
 		return fmt.Errorf("%w, streams: %d exceeds limit: %d, stream: '%s'", errMaxStreamsLimitExceeded, streams, b.maxStreams, labels)
 	}
-
 	// Add the entry as a new stream
 	b.streams[labels] = &push.Stream{
 		Labels:  labels,
-		Entries: []push.Entry{entry},
+		Entries: []push.Entry{entry.Entry},
 	}
 	b.countForSegment(segmentNum)
-
 	return nil
 }
 

internal/component/common/loki/client/batch_test.go

@@ -27,7 +27,7 @@ func TestBatch_MaxStreams(t *testing.T) {
 
 	errCount := 0
 	for _, entry := range inputEntries {
-		err := b.add(entry)
+		err := b.add(entry, 0)
 		if err != nil {
 			errCount++
 			assert.ErrorIs(t, err, errMaxStreamsLimitExceeded)
@@ -78,7 +78,7 @@ func TestBatch_add(t *testing.T) {
 			b := newBatch(0)
 
 			for _, entry := range testData.inputEntries {
-				err := b.add(entry)
+				err := b.add(entry, 0)
 				assert.NoError(t, err)
 			}
 
@@ -146,9 +146,9 @@ func TestHashCollisions(t *testing.T) {
 	const entriesPerLabel = 10
 
 	for i := 0; i < entriesPerLabel; i++ {
-		_ = b.add(loki.Entry{Labels: ls1, Entry: push.Entry{Timestamp: time.Now(), Line: fmt.Sprintf("line %d", i)}})
+		_ = b.add(loki.Entry{Labels: ls1, Entry: push.Entry{Timestamp: time.Now(), Line: fmt.Sprintf("line %d", i)}}, 0)
 
-		_ = b.add(loki.Entry{Labels: ls2, Entry: push.Entry{Timestamp: time.Now(), Line: fmt.Sprintf("line %d", i)}})
+		_ = b.add(loki.Entry{Labels: ls2, Entry: push.Entry{Timestamp: time.Now(), Line: fmt.Sprintf("line %d", i)}}, 0)
 	}
 
 	// make sure that colliding labels are stored properly as independent streams

internal/component/common/loki/client/client.go

@@ -14,7 +14,6 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/grafana/dskit/backoff"
-	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/common/config"
 	"github.com/prometheus/common/model"
 
@@ -23,7 +22,6 @@ import (
 	"github.com/grafana/alloy/internal/component/common/loki"
 	"github.com/grafana/alloy/internal/runtime/logging/level"
 	"github.com/grafana/alloy/internal/useragent"
-	"github.com/grafana/alloy/internal/util"
 )
 
 const (
@@ -33,111 +31,18 @@ const (
 	// Label reserved to override the tenant ID while processing
 	// pipeline stages
 	ReservedLabelTenantID = "__tenant_id__"
-
-	LatencyLabel = "filename"
-	HostLabel    = "host"
-	ClientLabel  = "client"
-	TenantLabel  = "tenant"
-	ReasonLabel  = "reason"
-
-	ReasonGeneric       = "ingester_error"
-	ReasonRateLimited   = "rate_limited"
-	ReasonStreamLimited = "stream_limited"
-	ReasonLineTooLong   = "line_too_long"
 )
 
-var Reasons = []string{ReasonGeneric, ReasonRateLimited, ReasonStreamLimited, ReasonLineTooLong}
-
 var userAgent = useragent.Get()
 
-type Metrics struct {
-	encodedBytes                 *prometheus.CounterVec
-	sentBytes                    *prometheus.CounterVec
-	droppedBytes                 *prometheus.CounterVec
-	sentEntries                  *prometheus.CounterVec
-	droppedEntries               *prometheus.CounterVec
-	mutatedEntries               *prometheus.CounterVec
-	mutatedBytes                 *prometheus.CounterVec
-	requestDuration              *prometheus.HistogramVec
-	batchRetries                 *prometheus.CounterVec
-	countersWithHostTenant       []*prometheus.CounterVec
-	countersWithHostTenantReason []*prometheus.CounterVec
-}
-
-func NewMetrics(reg prometheus.Registerer) *Metrics {
-	var m Metrics
-
-	m.encodedBytes = prometheus.NewCounterVec(prometheus.CounterOpts{
-		Name: "loki_write_encoded_bytes_total",
-		Help: "Number of bytes encoded and ready to send.",
-	}, []string{HostLabel, TenantLabel})
-	m.sentBytes = prometheus.NewCounterVec(prometheus.CounterOpts{
-		Name: "loki_write_sent_bytes_total",
-		Help: "Number of bytes sent.",
-	}, []string{HostLabel, TenantLabel})
-	m.droppedBytes = prometheus.NewCounterVec(prometheus.CounterOpts{
-		Name: "loki_write_dropped_bytes_total",
-		Help: "Number of bytes dropped because failed to be sent to the ingester after all retries.",
-	}, []string{HostLabel, TenantLabel, ReasonLabel})
-	m.sentEntries = prometheus.NewCounterVec(prometheus.CounterOpts{
-		Name: "loki_write_sent_entries_total",
-		Help: "Number of log entries sent to the ingester.",
-	}, []string{HostLabel, TenantLabel})
-	m.droppedEntries = prometheus.NewCounterVec(prometheus.CounterOpts{
-		Name: "loki_write_dropped_entries_total",
-		Help: "Number of log entries dropped because failed to be sent to the ingester after all retries.",
-	}, []string{HostLabel, TenantLabel, ReasonLabel})
-	m.mutatedEntries = prometheus.NewCounterVec(prometheus.CounterOpts{
-		Name: "loki_write_mutated_entries_total",
-		Help: "The total number of log entries that have been mutated.",
-	}, []string{HostLabel, TenantLabel, ReasonLabel})
-	m.mutatedBytes = prometheus.NewCounterVec(prometheus.CounterOpts{
-		Name: "loki_write_mutated_bytes_total",
-		Help: "The total number of bytes that have been mutated.",
-	}, []string{HostLabel, TenantLabel, ReasonLabel})
-	m.requestDuration = prometheus.NewHistogramVec(prometheus.HistogramOpts{
-		Name: "loki_write_request_duration_seconds",
-		Help: "Duration of send requests.",
-	}, []string{"status_code", HostLabel, TenantLabel})
-	m.batchRetries = prometheus.NewCounterVec(prometheus.CounterOpts{
-		Name: "loki_write_batch_retries_total",
-		Help: "Number of times batches has had to be retried.",
-	}, []string{HostLabel, TenantLabel})
-
-	m.countersWithHostTenant = []*prometheus.CounterVec{
-		m.batchRetries, m.encodedBytes, m.sentBytes, m.sentEntries,
-	}
-
-	m.countersWithHostTenantReason = []*prometheus.CounterVec{
-		m.droppedBytes, m.droppedEntries, m.mutatedEntries, m.mutatedBytes,
-	}
-
-	if reg != nil {
-		m.encodedBytes = util.MustRegisterOrGet(reg, m.encodedBytes).(*prometheus.CounterVec)
-		m.sentBytes = util.MustRegisterOrGet(reg, m.sentBytes).(*prometheus.CounterVec)
-		m.droppedBytes = util.MustRegisterOrGet(reg, m.droppedBytes).(*prometheus.CounterVec)
-		m.sentEntries = util.MustRegisterOrGet(reg, m.sentEntries).(*prometheus.CounterVec)
-		m.droppedEntries = util.MustRegisterOrGet(reg, m.droppedEntries).(*prometheus.CounterVec)
-		m.mutatedEntries = util.MustRegisterOrGet(reg, m.mutatedEntries).(*prometheus.CounterVec)
-		m.mutatedBytes = util.MustRegisterOrGet(reg, m.mutatedBytes).(*prometheus.CounterVec)
-		m.requestDuration = util.MustRegisterOrGet(reg, m.requestDuration).(*prometheus.HistogramVec)
-		m.batchRetries = util.MustRegisterOrGet(reg, m.batchRetries).(*prometheus.CounterVec)
-	}
-
-	return &m
-}
-
 // Client pushes entries to Loki and can be stopped
 type Client interface {
-	loki.EntryHandler
-	// Stop goroutine sending batch of entries without retries.
-	StopNow()
-	Name() string
+	Chan() chan<- loki.Entry
+	Stop()
 }
 
 // Client for pushing logs in snappy-compressed protos over HTTP.
 type client struct {
-	name    string
 	metrics *Metrics
 	logger  log.Logger
 	cfg     Config
@@ -148,17 +53,16 @@ type client struct {
 	wg   sync.WaitGroup
 
 	// ctx is used in any upstream calls from the `client`.
-	ctx        context.Context
-	cancel     context.CancelFunc
-	maxStreams int
+	ctx    context.Context
+	cancel context.CancelFunc
 }
 
 // New makes a new Client.
-func New(metrics *Metrics, cfg Config, maxStreams int, logger log.Logger) (Client, error) {
-	return newClient(metrics, cfg, maxStreams, logger)
+func New(metrics *Metrics, cfg Config, logger log.Logger) (Client, error) {
+	return newClient(metrics, cfg, logger)
 }
 
-func newClient(metrics *Metrics, cfg Config, maxStreams int, logger log.Logger) (*client, error) {
+func newClient(metrics *Metrics, cfg Config, logger log.Logger) (*client, error) {
 	if cfg.URL.URL == nil {
 		return nil, errors.New("client needs target URL")
 	}
@@ -169,17 +73,12 @@ func newClient(metrics *Metrics, cfg Config, maxStreams int, logger log.Logger)
 	ctx, cancel := context.WithCancel(context.Background())
 
 	c := &client{
-		logger:     log.With(logger, "component", "client", "host", cfg.URL.Host),
-		cfg:        cfg,
-		entries:    make(chan loki.Entry),
-		metrics:    metrics,
-		name:       GetClientName(cfg),
-		ctx:        ctx,
-		cancel:     cancel,
-		maxStreams: maxStreams,
-	}
-	if cfg.Name != "" {
-		c.name = cfg.Name
+		logger:  log.With(logger, "component", "client", "host", cfg.URL.Host),
+		cfg:     cfg,
+		entries: make(chan loki.Entry),
+		metrics: metrics,
+		ctx:     ctx,
+		cancel:  cancel,
 	}
 
 	err := cfg.Client.Validate()
@@ -194,8 +93,7 @@ func newClient(metrics *Metrics, cfg Config, maxStreams int, logger log.Logger)
 
 	c.client.Timeout = cfg.Timeout
 
-	c.wg.Add(1)
-	go c.run()
+	c.wg.Go(func() { c.run() })
 	return c, nil
 }
 
@@ -233,8 +131,6 @@ func (c *client) run() {
 		for tenantID, batch := range batches {
 			c.sendBatch(tenantID, batch)
 		}
-
-		c.wg.Done()
 	}()
 
 	for {
@@ -249,7 +145,7 @@ func (c *client) run() {
 
 			// If the batch doesn't exist yet, we create a new one with the entry
 			if !ok {
-				batches[tenantID] = newBatch(c.maxStreams, e)
+				batches[tenantID] = newBatch(c.cfg.MaxStreams, e)
 				c.initBatchMetrics(tenantID)
 				break
 			}
@@ -259,12 +155,12 @@ func (c *client) run() {
 			if batch.sizeBytesAfter(e.Entry) > c.cfg.BatchSize {
 				c.sendBatch(tenantID, batch)
 
-				batches[tenantID] = newBatch(c.maxStreams, e)
+				batches[tenantID] = newBatch(c.cfg.MaxStreams, e)
 				break
 			}
 
 			// The max size of the batch isn't reached, so we can add the entry
-			err := batch.add(e)
+			err := batch.add(e, 0)
 			if err != nil {
 				level.Error(c.logger).Log("msg", "batch add err", "tenant", tenantID, "error", err)
 				reason := ReasonGeneric
@@ -422,18 +318,7 @@ func (c *client) Stop() {
 	c.wg.Wait()
 }
 
-// StopNow stops the client without retries
-func (c *client) StopNow() {
-	// cancel will stop retrying http requests.
-	c.cancel()
-	c.Stop()
-}
-
 func (c *client) processEntry(e loki.Entry) (loki.Entry, string) {
 	tenantID := c.getTenantID(e.Labels)
 	return e, tenantID
 }
-
-func (c *client) Name() string {
-	return c.name
-}