fix(loki.write): Don't hang on shutdown when the send queue is full

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. Writes blocked on a full queue also get this long to find room before being dropped. | `"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_fanout.go

@@ -71,6 +71,13 @@ func (c *FanoutConsumer) Chan() chan<- loki.Entry {
 func (c *FanoutConsumer) Stop() {
 	// First stop the receiving channel.
 	c.once.Do(func() { close(c.recv) })
+
+	// run may be blocked enqueueing to an endpoint whose queue is full, so release it before
+	// waiting on it.
+	for _, c := range c.endpoints {
+		c.stopAccepting()
+	}
+
 	c.wg.Wait()
 
 	var stopWG sync.WaitGroup

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

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"net/http/httptest"
 	"net/url"
 	"testing"
 	"time"
@@ -205,3 +206,87 @@ func newServerAndEndpointConfig(t *testing.T) (Config, chan util.RemoteWriteRequ
 		close(receivedReqsChan)
 	}
 }
+
+func TestFanoutConsumer_StopWhileBlockedOnOverflow(t *testing.T) {
+	// An endpoint that never succeeds, combined with infinite retries, means batches are
+	// never drained and the send queue stays full. With BlockOnOverflow the enqueue loop
+	// then blocks, and Stop must still be able to shut the consumer down.
+	//
+	// Queues are per shard and entries are routed by label fingerprint, so a single stream
+	// only ever fills one shard's queue. Cover more than one shard to make sure that still
+	// releases on shutdown.
+	for _, minShards := range []int{1, 3} {
+		t.Run(fmt.Sprintf("min_shards=%d", minShards), func(t *testing.T) {
+			testStopWhileBlockedOnOverflow(t, minShards)
+		})
+	}
+}
+
+func testStopWhileBlockedOnOverflow(t *testing.T, minShards int) {
+	const drainTimeout = time.Second
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+	}))
+	t.Cleanup(server.Close)
+
+	serverURL, _ := url.Parse(server.URL)
+	consumer, err := NewFanoutConsumer(logging.NewSlogNop(), prometheus.NewRegistry(), Config{
+		Name:      "blocked",
+		URL:       flagext.URLValue{URL: serverURL},
+		Timeout:   time.Second,
+		BatchSize: 1,
+		BackoffConfig: backoff.Config{
+			MinBackoff: time.Millisecond,
+			MaxBackoff: 10 * time.Millisecond,
+			MaxRetries: 0, // retry indefinitely, as max_backoff_retries = 0 does
+		},
+		QueueConfig: QueueConfig{
+			Capacity:        1,
+			MinShards:       minShards,
+			DrainTimeout:    drainTimeout,
+			BlockOnOverflow: true,
+		},
+	})
+	require.NoError(t, err)
+
+	// Push entries until the consumer stops accepting them, which means the enqueue loop
+	// is blocked on a full queue.
+	backpressure := make(chan struct{})
+	go func() {
+		defer close(backpressure)
+		for {
+			select {
+			case consumer.Chan() <- loki.Entry{
+				Labels: model.LabelSet{"test": "backpressure"},
+				Entry:  push.Entry{Timestamp: time.Now(), Line: "line"},
+			}:
+			case <-time.After(time.Second):
+				return
+			}
+		}
+	}()
+
+	select {
+	case <-backpressure:
+	case <-time.After(30 * time.Second):
+		t.Fatal("timed out waiting for the send queue to fill")
+	}
+
+	stopped := make(chan struct{})
+	start := time.Now()
+	go func() {
+		consumer.Stop()
+		close(stopped)
+	}()
+
+	select {
+	case <-stopped:
+	case <-time.After(30 * time.Second):
+		t.Fatal("Stop did not return while the enqueue loop was blocked on a full queue")
+	}
+
+	// The queue can never drain here, so shutdown is expected to spend the whole drain budget
+	// waiting for room before giving up on the entries it still holds.
+	require.GreaterOrEqual(t, time.Since(start), 9*drainTimeout/10, "Stop gave up before the drain timeout")
+}

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

@@ -91,6 +91,10 @@ type endpointWatcherPair struct {
 
 // Stop will proceed to stop, in order, watcher and the endpoint.
 func (p endpointWatcherPair) Stop(drain bool) {
+	// Start the endpoint's drain budget first. The watcher may be blocked enqueueing to a full
+	// queue, in which case draining reads nothing and both waits below stack up behind it.
+	p.endpoint.StopAccepting()
+
 	// If drain enabled, drain the WAL.
 	if drain {
 		p.watcher.Drain()
@@ -236,8 +240,13 @@ func (c *walEndpointAdapter) AppendEntries(entries wal.RefEntries, segment int)
 	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.
+// StopAccepting releases the watcher if it is blocked enqueueing to a full send queue.
+func (c *walEndpointAdapter) StopAccepting() {
+	c.endpoint.stopAccepting()
+}
+
+// Stop the endpoint, enqueueing pending batches and draining the send queue accordingly. Each of
+// those is limited by the configured drain timeout.
 func (c *walEndpointAdapter) Stop() {
 	c.endpoint.stop()
 	c.tracker.Stop()

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

@@ -3,7 +3,10 @@ package client
 import (
 	"context"
 	"fmt"
+	"net/http"
+	"net/http/httptest"
 	"net/url"
+	"sync"
 	"testing"
 	"time"
 
@@ -536,3 +539,120 @@ func runEndpointBenchCase(b *testing.B, bc testCase) {
 	endpoint.stop()
 	close(receivedReqsChan)
 }
+
+func TestWALConsumer_StopWhileBlockedOnOverflow(t *testing.T) {
+	// An endpoint that never succeeds, combined with infinite retries, means batches are
+	// never drained and the send queue stays full. With BlockOnOverflow the WAL watcher then
+	// blocks enqueueing, and shutting the consumer down must still work.
+	//
+	// Cover both shutdown paths: Update uses Stop, while component shutdown uses StopAndDrain.
+	t.Run("Stop", func(t *testing.T) {
+		testWALStopWhileBlockedOnOverflow(t, func(c *WALConsumer) { c.Stop() })
+	})
+	t.Run("StopAndDrain", func(t *testing.T) {
+		testWALStopWhileBlockedOnOverflow(t, func(c *WALConsumer) { c.StopAndDrain() })
+	})
+}
+
+func testWALStopWhileBlockedOnOverflow(t *testing.T, stop func(*WALConsumer)) {
+	const drainTimeout = time.Second
+
+	// The WAL drain has its own timeout, and a blocked watcher reads nothing during it, so the
+	// two waits must not stack behind each other.
+	watchConfig := wal.DefaultWatchConfig
+	watchConfig.DrainTimeout = drainTimeout
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+	}))
+	t.Cleanup(server.Close)
+
+	serverURL, _ := url.Parse(server.URL)
+	reg := prometheus.NewRegistry()
+	consumer, err := NewWALConsumer(logging.NewSlogNop(), reg, wal.Config{
+		Dir:           t.TempDir(),
+		Enabled:       true,
+		MaxSegmentAge: time.Minute,
+		WatchConfig:   watchConfig,
+	}, Config{
+		Name:      "blocked",
+		URL:       flagext.URLValue{URL: serverURL},
+		Timeout:   time.Second,
+		BatchSize: 1,
+		BackoffConfig: backoff.Config{
+			MinBackoff: time.Millisecond,
+			MaxBackoff: 10 * time.Millisecond,
+			MaxRetries: 0, // retry indefinitely, as max_backoff_retries = 0 does
+		},
+		QueueConfig: QueueConfig{
+			Capacity:        1,
+			MinShards:       1,
+			DrainTimeout:    drainTimeout,
+			BlockOnOverflow: true,
+		},
+	})
+	require.NoError(t, err)
+
+	// Keep writing so the watcher has more to read than the queue can hold.
+	writing, stopWriting := context.WithCancel(context.Background())
+	var writers sync.WaitGroup
+	writers.Go(func() {
+		for {
+			select {
+			case consumer.Chan() <- loki.Entry{
+				Labels: model.LabelSet{"test": "backpressure"},
+				Entry:  push.Entry{Timestamp: time.Now(), Line: "line"},
+			}:
+			case <-writing.Done():
+				return
+			}
+		}
+	})
+
+	// Retries mean a batch is stuck in flight and the queue behind it is filling.
+	require.Eventually(t, func() bool {
+		return counterValue(t, reg, "loki_write_batch_retries_total") > 0
+	}, 30*time.Second, 50*time.Millisecond, "timed out waiting for the endpoint to start retrying")
+
+	// Stop writing before Stop closes the WAL writer's channel underneath us. The watcher
+	// stays blocked either way, since the segments already hold more than the queue can take.
+	stopWriting()
+	writers.Wait()
+
+	stopped := make(chan struct{})
+	start := time.Now()
+	go func() {
+		stop(consumer)
+		close(stopped)
+	}()
+
+	select {
+	case <-stopped:
+	case <-time.After(30 * time.Second):
+		t.Fatal("shutdown did not return while the WAL watcher was blocked on a full queue")
+	}
+
+	// Releasing the blocked enqueue and draining the send queue each get their own budget, but
+	// the WAL drain runs concurrently with the first rather than ahead of it.
+	elapsed := time.Since(start)
+	t.Logf("shutdown took %s", elapsed)
+	require.Less(t, elapsed, 5*drainTimeout/2, "shutdown waits are stacking")
+}
+
+func counterValue(t *testing.T, g prometheus.Gatherer, name string) float64 {
+	t.Helper()
+
+	families, err := g.Gather()
+	require.NoError(t, err)
+
+	var total float64
+	for _, family := range families {
+		if family.GetName() != name {
+			continue
+		}
+		for _, m := range family.GetMetric() {
+			total += m.GetCounter().GetValue()
+		}
+	}
+	return total
+}

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

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/grafana/dskit/backoff"
+	"go.uber.org/atomic"
 
 	"github.com/grafana/alloy/internal/component/common/loki"
 	"github.com/grafana/alloy/internal/component/common/loki/client/internal/marker"
@@ -25,6 +26,9 @@ type endpoint struct {
 
 	shards  *shards
 	backoff *backoff.Backoff
+
+	// drainDeadline is when the shutdown drain budget expires, or nil while not draining.
+	drainDeadline atomic.Pointer[time.Time]
 }
 
 func newEndpoint(metrics *metrics, cfg Config, logger *slog.Logger, markerHandler marker.Tracker) (*endpoint, error) {
@@ -57,14 +61,14 @@ func newEndpoint(metrics *metrics, cfg Config, logger *slog.Logger, markerHandle
 var errQueueIsFull = errors.New("queue is full")
 
 // enqueue tries to enqueue an entry. It returns an error if the entry could not be enqueued.
-// errQueueIsFull when the queue is full and BlockOnOverflow is false, or context.Canceled when
-// endpoint is stopped.
+// errQueueIsFull when the queue is full and either BlockOnOverflow is false or the drain budget
+// is spent, or context.Canceled when endpoint is stopped.
 func (e *endpoint) enqueue(entry loki.Entry, segmentNum int) error {
 	defer e.backoff.Reset()
 
 	tenantID := getTenantID(e.cfg, entry)
 	for !e.shards.enqueue(tenantID, entry, segmentNum) {
-		if !e.cfg.QueueConfig.BlockOnOverflow {
+		if !e.cfg.QueueConfig.BlockOnOverflow || e.drainExpired() {
 			e.metrics.droppedEntries.WithLabelValues(e.cfg.URL.Host, tenantID, reasonQueueIsFull).Inc()
 			e.metrics.droppedBytes.WithLabelValues(e.cfg.URL.Host, tenantID, reasonQueueIsFull).Add(float64(entry.Size()))
 			return errQueueIsFull
@@ -79,7 +83,22 @@ func (e *endpoint) enqueue(entry loki.Entry, segmentNum int) error {
 	return nil
 }
 
+// stopAccepting gives enqueue DrainTimeout to place whatever it is still holding, after which it
+// drops on a full queue rather than blocking. Callers must run this before waiting on whichever
+// goroutine feeds enqueue, otherwise that wait deadlocks against the blocked enqueue.
+func (e *endpoint) stopAccepting() {
+	deadline := time.Now().Add(e.cfg.QueueConfig.DrainTimeout)
+	// Keep the deadline set by the first caller; shutdown paths call this more than once.
+	e.drainDeadline.CompareAndSwap(nil, &deadline)
+}
+
+func (e *endpoint) drainExpired() bool {
+	deadline := e.drainDeadline.Load()
+	return deadline != nil && !time.Now().Before(*deadline)
+}
+
 func (e *endpoint) stop() {
+	e.stopAccepting()
 	e.cancel()
 	e.shards.stop()
 }