fix(loki.write): Potential deadlock when queue is full and component shuts down

docs/sources/reference/components/loki/loki.write.md

@@ -81,7 +81,7 @@ The following arguments are supported:
 | `http_headers`           | `map(list(secret))` | Custom HTTP headers to be sent along with each request. The map key is the header name.          |           | no       |
 | `headers`                | `map(string)`       | Extra headers to deliver with the request.                                                       |           | no       |
 | `max_backoff_period`     | `duration`          | Maximum backoff time between retries.                                                            | `"5m"`    | no       |
-| `max_backoff_retries`    | `int`               | Maximum number of retries.                                                                       | `10`      | no       |
+| `max_backoff_retries`    | `int`               | Maximum number of retries. Set to `0` to retry indefinitely.                                     | `10`      | no       |
 | `min_backoff_period`     | `duration`          | Initial backoff time between retries.                                                            | `"500ms"` | no       |
 | `name`                   | `string`            | Optional name to identify this endpoint with.                                                    |           | no       |
 | `no_proxy`               | `string`            | Comma-separated list of IP addresses, CIDR notations, and domain names to exclude from proxying. |           | no       |
@@ -139,7 +139,7 @@ The following arguments are supported:
 | -------------------- | ---------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | -------- |
 | `block_on_overflow`  | `bool`     | If `true`, block until there is space in the queue; if `false`, drop entries when queue is full.                                                                              | `true`  | no       |
 | `capacity`           | `string`   | Controls the size of the underlying send queue buffer. This setting should be considered a worst-case scenario of memory consumption, in which all enqueued batches are full. | `10MiB` | no       |
-| `drain_timeout`      | `duration` | Configures the maximum time the client can take to drain the send queue upon shutdown. During that time, it enqueues pending batches and drains the send queue sending each.  | `"1m"`  | no       |
+| `drain_timeout`      | `duration` | Configures the maximum time the client can take to drain the send queue upon shutdown. During that time, it enqueues pending batches and drains the send queue sending each.  | `"15s"` | no       |
 | `min_shards`         | `number`   | Minimum number of concurrent shards sending samples to the endpoint.                                                                                                          | `1`     | no       |
 
 Each endpoint is divided into a number of concurrent _shards_ which are responsible for sending a fraction of batches. The number of shards is controlled with `min_shards` argument.
@@ -171,7 +171,7 @@ The following arguments are supported:
 
 | Name                 | Type       | Description                                                                                                    | Default   | Required |
 | -------------------- | ---------- | -------------------------------------------------------------------------------------------------------------- | --------- | -------- |
-| `drain_timeout`      | `duration` | Maximum time the WAL drain procedure can take, before being forcefully stopped.                                | `"30s"`   | no       |
+| `drain_timeout`      | `duration` | Maximum time the WAL drain procedure can take, before being forcefully stopped.                                | `"15s"`   | no       |
 | `enabled`            | `bool`     | Whether to enable the WAL.                                                                                     | `false`   | no       |
 | `max_read_frequency` | `duration` | Maximum backoff time in the backup read mechanism.                                                             | `"1s"`    | no       |
 | `max_segment_age`    | `duration` | Maximum time a WAL segment should be allowed to live. Segments older than this setting are eventually deleted. | `"1h"`    | no       |

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

@@ -1,11 +1,14 @@
 package client
 
-import "github.com/grafana/alloy/internal/component/common/loki"
+import (
+	"context"
 
-// Consumer is an interface for consuming Loki log entries. It provides a channel
-// to send entries to and a method to stop the consumer.
+	"github.com/grafana/alloy/internal/component/common/loki"
+)
+
+// Consumer is an interface for consuming Loki log entries.
 type Consumer interface {
-	Chan() chan<- loki.Entry
+	ConsumeEntry(ctx context.Context, entry loki.Entry) error
 	Stop()
 }
 

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

@@ -1,6 +1,8 @@
 package client
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"log/slog"
 	"sync"
@@ -16,9 +18,8 @@ func NewFanoutConsumer(logger *slog.Logger, reg prometheus.Registerer, cfgs ...C
 		return nil, fmt.Errorf("at least one endpoint config must be provided")
 	}
 
-	m := &FanoutConsumer{
+	c := &FanoutConsumer{
 		endpoints: make([]*endpoint, 0, len(cfgs)),
-		recv:      make(chan loki.Entry),
 	}
 
 	var (
@@ -39,40 +40,36 @@ func NewFanoutConsumer(logger *slog.Logger, reg prometheus.Registerer, cfgs ...C
 			return nil, fmt.Errorf("error starting endpoint: %w", err)
 		}
 
-		m.endpoints = append(m.endpoints, endpoint)
+		c.endpoints = append(c.endpoints, endpoint)
 	}
 
-	m.wg.Go(m.run)
-	return m, nil
+	return c, nil
 }
 
 var _ Consumer = (*FanoutConsumer)(nil)
 
 type FanoutConsumer struct {
 	endpoints []*endpoint
-	wg        sync.WaitGroup
-	once      sync.Once
-	recv      chan loki.Entry
 }
 
-func (c *FanoutConsumer) run() {
-	for e := range c.recv {
-		for _, c := range c.endpoints {
-			// NOTE: For now it's fine to ignore error because we can't act on it.
-			_ = c.enqueue(e, 0)
+func (c *FanoutConsumer) ConsumeEntry(ctx context.Context, entry loki.Entry) error {
+	for _, e := range c.endpoints {
+		err := e.enqueue(ctx, entry, 0)
+
+		if err != nil {
+			// If we get errQueueIsFull we skipped the entry for this endpoints
+			// and should try next.
+			if errors.Is(err, errQueueIsFull) {
+				continue
+			}
+			// For any other error we assume it's not useful to send to the next endpoint
+			return err
 		}
 	}
-}
-
-func (c *FanoutConsumer) Chan() chan<- loki.Entry {
-	return c.recv
+	return nil
 }
 
 func (c *FanoutConsumer) Stop() {
-	// First stop the receiving channel.
-	c.once.Do(func() { close(c.recv) })
-	c.wg.Wait()
-
 	var stopWG sync.WaitGroup
 	// Stop all endpoints.
 	for _, c := range c.endpoints {

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

@@ -43,13 +43,19 @@ func TestFanoutConsumer(t *testing.T) {
 	}
 	var totalLines = 100
 	for i := range totalLines {
-		consumer.Chan() <- loki.Entry{
-			Labels: testLabels,
-			Entry: push.Entry{
-				Timestamp: time.Now(),
-				Line:      fmt.Sprintf("line%d", i),
-			},
-		}
+		require.NoError(
+			t,
+			consumer.ConsumeEntry(
+				t.Context(),
+				loki.Entry{
+					Labels: testLabels,
+					Entry: push.Entry{
+						Timestamp: time.Now(),
+						Line:      fmt.Sprintf("line%d", i),
+					},
+				},
+			),
+		)
 	}
 
 	require.Eventually(t, func() bool {
@@ -104,13 +110,19 @@ func TestFanoutConsumer_MultipleConfigs(t *testing.T) {
 	}
 	var totalLines = 100
 	for i := range totalLines {
-		consumer.Chan() <- loki.Entry{
-			Labels: testLabels,
-			Entry: push.Entry{
-				Timestamp: time.Now(),
-				Line:      fmt.Sprintf("line%d", i),
-			},
-		}
+		require.NoError(
+			t,
+			consumer.ConsumeEntry(
+				t.Context(),
+				loki.Entry{
+					Labels: testLabels,
+					Entry: push.Entry{
+						Timestamp: time.Now(),
+						Line:      fmt.Sprintf("line%d", i),
+					},
+				},
+			),
+		)
 	}
 
 	// times 2 due to endpoints being run
@@ -205,3 +217,51 @@ func newServerAndEndpointConfig(t *testing.T) (Config, chan util.RemoteWriteRequ
 		close(receivedReqsChan)
 	}
 }
+
+func TestFanoutConsumer_StopWithFullSendQueue(t *testing.T) {
+	const drainTimeout = time.Second
+
+	server, blocked, release := newBlockedServer()
+	defer server.Close()
+	defer release()
+
+	serverURL, err := url.Parse(server.URL)
+	require.NoError(t, err)
+
+	endpointConfig := Config{
+		Name: "test-client",
+		URL:  flagext.URLValue{URL: serverURL},
+		// Long enough that the in-flight request stays parked for the whole test.
+		Timeout:   time.Minute,
+		BatchSize: 1,
+		BackoffConfig: backoff.Config{
+			MinBackoff: time.Millisecond,
+			MaxBackoff: 10 * time.Millisecond,
+			MaxRetries: 0,
+		},
+		QueueConfig: QueueConfig{
+			Capacity:        1,
+			MinShards:       1,
+			DrainTimeout:    drainTimeout,
+			BlockOnOverflow: true,
+		},
+	}
+
+	consumer, err := NewFanoutConsumer(logging.NewSlogNop(), prometheus.NewRegistry(), endpointConfig)
+	require.NoError(t, err)
+
+	feedUntilBlocked(t, blocked, consumer)
+
+	done := make(chan struct{})
+	go func() {
+		consumer.Stop()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(5 * drainTimeout):
+		release()
+		t.Fatal("Stop did not finish in time")
+	}
+}

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

@@ -1,10 +1,12 @@
 package client
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"log/slog"
 	"sync"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/common/model"
@@ -79,7 +81,7 @@ func NewWALConsumer(logger *slog.Logger, reg prometheus.Registerer, walCfg wal.C
 		})
 	}
 
-	writer.Start(walCfg.MaxSegmentAge)
+	writer.Start()
 
 	return m, nil
 }
@@ -89,16 +91,16 @@ type endpointWatcherPair struct {
 	endpoint *walEndpointAdapter
 }
 
-// Stop will proceed to stop, in order, watcher and the endpoint.
-func (p endpointWatcherPair) Stop(drain bool) {
+// stop will proceed to stop, in order, watcher and the endpoint.
+func (p endpointWatcherPair) stop(drain bool) {
 	// If drain enabled, drain the WAL.
 	if drain {
 		p.watcher.Drain()
 	}
 	p.watcher.Stop()
 
 	// Subsequently stop the endpoint.
-	p.endpoint.Stop()
+	p.endpoint.stop()
 }
 
 var _ DrainableConsumer = (*WALConsumer)(nil)
@@ -108,8 +110,9 @@ type WALConsumer struct {
 	pairs  []endpointWatcherPair
 }
 
-func (m *WALConsumer) Chan() chan<- loki.Entry {
-	return m.writer.Chan()
+// ConsumeEntry implements DrainableConsumer.
+func (m *WALConsumer) ConsumeEntry(ctx context.Context, entry loki.Entry) error {
+	return m.writer.WriteEntry(entry)
 }
 
 func (m *WALConsumer) Stop() {
@@ -132,7 +135,7 @@ func (m *WALConsumer) stop(drain bool) {
 	// endpoint config, each (watcher, queue) pair is stopped concurrently.
 	for _, pair := range m.pairs {
 		stopWG.Go(func() {
-			pair.Stop(drain)
+			pair.stop(drain)
 		})
 	}
 
@@ -193,52 +196,55 @@ func (c *walEndpointAdapter) StoreSeries(series []record.RefSeries, segment int)
 	}
 }
 
-func (c *walEndpointAdapter) AppendEntries(entries wal.RefEntries, segment int) error {
+func (c *walEndpointAdapter) AppendEntries(ctx context.Context, entries wal.RefEntries, segment int) error {
 	c.seriesLock.RLock()
 	l, ok := c.series[entries.Ref]
 	c.seriesLock.RUnlock()
 
 	var (
 		queuedEntries    int
-		maxSeenTimestamp int64 = -1
+		maxSeenTimestamp time.Time
 	)
 
-	if ok {
-		for i := range entries.Entries {
-			e := entries.EntryAt(l, i)
-			err := c.endpoint.enqueue(e, segment)
-			// We can receive errQueueIsFull if we have configured endpoint with BlockOnOverflow.
-			// Here we just skip the entry and try with the next one.
-			if errors.Is(err, errQueueIsFull) {
-				continue
-			}
-			// NOTE: The only other error that can be returned is context.Canceled and that happens
-			// if endpoint was stopped.
-			if err != nil {
-				return nil
-			}
-
-			queuedEntries += 1
-			if e.Timestamp.Unix() > maxSeenTimestamp {
-				maxSeenTimestamp = e.Timestamp.Unix()
-			}
-		}
-		// update marker with all successfully queued entries.
-		c.tracker.UpdateReceivedData(segment, queuedEntries)
-	} else {
+	if !ok {
 		// TODO(thepalbi): Add metric here
-		c.logger.Debug("series for entry not found")
+		c.logger.Debug("series for entries not found")
+		return nil
+	}
+
+	for i := range entries.Entries {
+		entry := entries.EntryAt(l, i)
+		err := c.endpoint.enqueue(ctx, entry, segment)
+
+		// If we get errQueueIsFull we skipped the entry and should
+		// not count it as queued and should move on to the next one.
+		if errors.Is(err, errQueueIsFull) {
+			continue
+		}
+
+		if err != nil {
+			return err
+		}
+
+		queuedEntries += 1
+
+		if entry.Timestamp.After(maxSeenTimestamp) {
+			maxSeenTimestamp = entry.Timestamp
+		}
+	}
+	// update tracker with all successfully queued entries.
+	c.tracker.UpdateReceivedData(segment, queuedEntries)
+
+	if queuedEntries > 0 {
+		c.metrics.lastReadTimestamp.WithLabelValues().Set(float64(maxSeenTimestamp.Unix()))
 	}
 
-	// It's safe to assume that upon an AppendEntries call, there will always be at least
-	// one entry.
-	c.metrics.lastReadTimestamp.WithLabelValues().Set(float64(maxSeenTimestamp))
 	return nil
 }
 
-// Stop the endpoint, enqueueing pending batches and draining the send queue accordingly. Both closing operations are
-// limited by a deadline, controlled by a configured drain timeout, which is global to the Stop call.
-func (c *walEndpointAdapter) Stop() {
+// stop the endpoint, enqueueing pending batches and draining the send queue accordingly. Both closing operations are
+// limited by a deadline, controlled by a configured drain timeout, which is global to the stop call.
+func (c *walEndpointAdapter) stop() {
 	c.endpoint.stop()
 	c.tracker.Stop()
 }