diff --git a/docs/sources/reference/components/loki/loki.write.md b/docs/sources/reference/components/loki/loki.write.md index 0bb44414711..a9a65ad0091 100644 --- a/docs/sources/reference/components/loki/loki.write.md +++ b/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 | diff --git a/internal/component/common/loki/client/consumer.go b/internal/component/common/loki/client/consumer.go index 4aee54386bb..d7ed213c9a4 100644 --- a/internal/component/common/loki/client/consumer.go +++ b/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() } diff --git a/internal/component/common/loki/client/consumer_fanout.go b/internal/component/common/loki/client/consumer_fanout.go index 23d820e54a2..5805f067b8e 100644 --- a/internal/component/common/loki/client/consumer_fanout.go +++ b/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 { diff --git a/internal/component/common/loki/client/consumer_fanout_test.go b/internal/component/common/loki/client/consumer_fanout_test.go index 5cacf431c3f..6cde092616a 100644 --- a/internal/component/common/loki/client/consumer_fanout_test.go +++ b/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") + } +} diff --git a/internal/component/common/loki/client/consumer_wal.go b/internal/component/common/loki/client/consumer_wal.go index 184f438d931..524991c5363 100644 --- a/internal/component/common/loki/client/consumer_wal.go +++ b/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,8 +91,8 @@ 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() @@ -98,7 +100,7 @@ func (p endpointWatcherPair) Stop(drain bool) { 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() } diff --git a/internal/component/common/loki/client/consumer_wal_test.go b/internal/component/common/loki/client/consumer_wal_test.go index 838b823167c..afac0b90408 100644 --- a/internal/component/common/loki/client/consumer_wal_test.go +++ b/internal/component/common/loki/client/consumer_wal_test.go @@ -58,13 +58,19 @@ func TestWALConsumer(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 { @@ -125,13 +131,17 @@ func TestWALConsumer_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 endpoint being run @@ -286,7 +296,7 @@ func TestWALEndpoint(t *testing.T) { }, }, 0) - _ = adapter.AppendEntries(wal.RefEntries{ + _ = adapter.AppendEntries(t.Context(), wal.RefEntries{ Ref: chunks.HeadSeriesRef(mod), Created: time.Now().UnixMicro(), Entries: []push.Entry{{ @@ -305,7 +315,7 @@ func TestWALEndpoint(t *testing.T) { } // Stop the endpoint: it waits until the current batch is sent - adapter.Stop() + adapter.stop() close(receivedReqsChan) }) } @@ -431,7 +441,7 @@ func runWALEndpointBenchCase(b *testing.B, bc testCase, mhFactory func(t *testin }, }, 0) - _ = adapter.AppendEntries(wal.RefEntries{ + _ = adapter.AppendEntries(b.Context(), wal.RefEntries{ Ref: chunks.HeadSeriesRef(seriesId), Created: time.Now().UnixMicro(), Entries: []push.Entry{{ @@ -450,7 +460,7 @@ func runWALEndpointBenchCase(b *testing.B, bc testCase, mhFactory func(t *testin } // Stop the endpoint: it waits until the current batch is sent - adapter.Stop() + adapter.stop() close(receivedReqsChan) } @@ -512,7 +522,7 @@ func runEndpointBenchCase(b *testing.B, bc testCase) { // Send all the input log entries for j, l := range lines { seriesId := j % bc.numSeries - endpoint.enqueue(loki.Entry{ + endpoint.enqueue(b.Context(), loki.Entry{ Labels: model.LabelSet{ // take j module bc.numSeries to evenly distribute those numSeries across all sent entries "app": model.LabelValue(fmt.Sprintf("series-%d", seriesId)), @@ -536,3 +546,65 @@ func runEndpointBenchCase(b *testing.B, bc testCase) { endpoint.stop() close(receivedReqsChan) } + +func TestWALConsumer_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) + + walConfig := wal.Config{ + Dir: t.TempDir(), + Enabled: true, + MaxSegmentAge: time.Minute, + WatchConfig: wal.WatchConfig{ + MinReadFrequency: 10 * time.Millisecond, + MaxReadFrequency: 50 * time.Millisecond, + DrainTimeout: drainTimeout, + }, + } + + 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 := NewWALConsumer(logging.NewSlogNop(), prometheus.NewRegistry(), walConfig, endpointConfig) + require.NoError(t, err) + + feedUntilBlocked(t, blocked, consumer) + + stopped := make(chan struct{}) + go func() { + consumer.StopAndDrain() + close(stopped) + }() + + // A healthy shutdown costs at most the WAL drain timeout plus the queue drain timeout. + select { + case <-stopped: + case <-time.After(5 * (drainTimeout + drainTimeout)): + release() + t.Fatal("StopAndDrain did not finish in time") + } +} diff --git a/internal/component/common/loki/client/endpoint.go b/internal/component/common/loki/client/endpoint.go index c7e116f0501..520dff8af7b 100644 --- a/internal/component/common/loki/client/endpoint.go +++ b/internal/component/common/loki/client/endpoint.go @@ -18,13 +18,7 @@ type endpoint struct { cfg Config metrics *metrics logger *slog.Logger - entries chan loki.Entry - - ctx context.Context - cancel context.CancelFunc - shards *shards - backoff *backoff.Backoff } func newEndpoint(metrics *metrics, cfg Config, logger *slog.Logger, markerHandler marker.Tracker) (*endpoint, error) { @@ -35,52 +29,53 @@ func newEndpoint(metrics *metrics, cfg Config, logger *slog.Logger, markerHandle return nil, err } - ctx, cancel := context.WithCancel(context.Background()) c := &endpoint{ cfg: cfg, logger: logger, metrics: metrics, - entries: make(chan loki.Entry), - ctx: ctx, - cancel: cancel, shards: shards, - backoff: backoff.New(ctx, backoff.Config{ - MinBackoff: 5 * time.Millisecond, - MaxBackoff: 50 * time.Millisecond, - }), } c.shards.start(cfg.QueueConfig.MinShards) return c, nil } -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. -func (e *endpoint) enqueue(entry loki.Entry, segmentNum int) error { - defer e.backoff.Reset() +// enqueue tries to enqueue an entry. It waits for room while BlockOnOverflow +// is set and drops the entry when it is not, returning errQueueIsFull. It will +// return context error if caller cancels context or loki.ErrConsumerStopped +// if endpoint has been stopped. +func (e *endpoint) enqueue(ctx context.Context, entry loki.Entry, segmentNum int) error { + bo := backoff.New(ctx, backoff.Config{ + MinBackoff: 5 * time.Millisecond, + MaxBackoff: 50 * time.Millisecond, + }) tenantID := getTenantID(e.cfg, entry) - for !e.shards.enqueue(tenantID, entry, segmentNum) { - if !e.cfg.QueueConfig.BlockOnOverflow { + + for bo.Ongoing() { + err := e.shards.enqueue(tenantID, entry, segmentNum) + + if err == nil { + return nil + } + + if errors.Is(err, loki.ErrConsumerStopped) { + return err + } + + if errors.Is(err, errQueueIsFull) && !e.cfg.QueueConfig.BlockOnOverflow { 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 } - e.backoff.Wait() - if !e.backoff.Ongoing() { - return e.backoff.Err() - } + bo.Wait() } - return nil + return bo.Err() } func (e *endpoint) stop() { - e.cancel() e.shards.stop() } diff --git a/internal/component/common/loki/client/endpoint_test.go b/internal/component/common/loki/client/endpoint_test.go index e1b2d1b2ab7..9b039fdd057 100644 --- a/internal/component/common/loki/client/endpoint_test.go +++ b/internal/component/common/loki/client/endpoint_test.go @@ -1,8 +1,10 @@ package client import ( + "context" "errors" "net/http" + "net/http/httptest" "runtime" "strings" "testing" @@ -17,6 +19,7 @@ import ( "github.com/prometheus/common/config" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.uber.org/atomic" "github.com/grafana/alloy/internal/component/common/loki" "github.com/grafana/alloy/internal/component/common/loki/client/internal/marker" @@ -60,8 +63,8 @@ func TestEndpoint(t *testing.T) { loki_write_sent_entries_total{host="__HOST__",tenant=""} 3.0 # HELP loki_write_dropped_entries_total Number of log entries dropped because failed to be sent to the ingester after all retries. # TYPE loki_write_dropped_entries_total counter + loki_write_dropped_entries_total{host="__HOST__",reason="batch_too_large",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="ingester_error",tenant=""} 0 - loki_write_dropped_entries_total{host="__HOST__",reason="line_too_long",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="queue_is_full",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="rate_limited",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="stream_limited",tenant=""} 0 @@ -101,8 +104,8 @@ func TestEndpoint(t *testing.T) { loki_write_sent_entries_total{host="__HOST__",tenant=""} 2.0 # HELP loki_write_dropped_entries_total Number of log entries dropped because failed to be sent to the ingester after all retries. # TYPE loki_write_dropped_entries_total counter + loki_write_dropped_entries_total{host="__HOST__",reason="batch_too_large",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="ingester_error",tenant=""} 0 - loki_write_dropped_entries_total{host="__HOST__",reason="line_too_long",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="queue_is_full",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="rate_limited",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="stream_limited",tenant=""} 0 @@ -133,8 +136,8 @@ func TestEndpoint(t *testing.T) { expectedMetrics: ` # HELP loki_write_dropped_entries_total Number of log entries dropped because failed to be sent to the ingester after all retries. # TYPE loki_write_dropped_entries_total counter + loki_write_dropped_entries_total{host="__HOST__",reason="batch_too_large",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="ingester_error",tenant=""} 1 - loki_write_dropped_entries_total{host="__HOST__",reason="line_too_long",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="queue_is_full",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="rate_limited",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="stream_limited",tenant=""} 0 @@ -160,8 +163,8 @@ func TestEndpoint(t *testing.T) { expectedMetrics: ` # HELP loki_write_dropped_entries_total Number of log entries dropped because failed to be sent to the ingester after all retries. # TYPE loki_write_dropped_entries_total counter + loki_write_dropped_entries_total{host="__HOST__",reason="batch_too_large",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="ingester_error",tenant=""} 1 - loki_write_dropped_entries_total{host="__HOST__",reason="line_too_long",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="queue_is_full",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="rate_limited",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="stream_limited",tenant=""} 0 @@ -195,8 +198,8 @@ func TestEndpoint(t *testing.T) { expectedMetrics: ` # HELP loki_write_dropped_entries_total Number of log entries dropped because failed to be sent to the ingester after all retries. # TYPE loki_write_dropped_entries_total counter + loki_write_dropped_entries_total{host="__HOST__",reason="batch_too_large",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="ingester_error",tenant=""} 0 - loki_write_dropped_entries_total{host="__HOST__",reason="line_too_long",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="queue_is_full",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="rate_limited",tenant=""} 1 loki_write_dropped_entries_total{host="__HOST__",reason="stream_limited",tenant=""} 0 @@ -223,8 +226,8 @@ func TestEndpoint(t *testing.T) { expectedMetrics: ` # HELP loki_write_dropped_entries_total Number of log entries dropped because failed to be sent to the ingester after all retries. # TYPE loki_write_dropped_entries_total counter + loki_write_dropped_entries_total{host="__HOST__",reason="batch_too_large",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="ingester_error",tenant=""} 0 - loki_write_dropped_entries_total{host="__HOST__",reason="line_too_long",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="queue_is_full",tenant=""} 0 loki_write_dropped_entries_total{host="__HOST__",reason="rate_limited",tenant=""} 1 loki_write_dropped_entries_total{host="__HOST__",reason="stream_limited",tenant=""} 0 @@ -254,8 +257,8 @@ func TestEndpoint(t *testing.T) { loki_write_sent_entries_total{host="__HOST__",tenant="tenant-default"} 2.0 # HELP loki_write_dropped_entries_total Number of log entries dropped because failed to be sent to the ingester after all retries. # TYPE loki_write_dropped_entries_total counter + loki_write_dropped_entries_total{host="__HOST__", reason="batch_too_large", tenant="tenant-default"} 0 loki_write_dropped_entries_total{host="__HOST__", reason="ingester_error", tenant="tenant-default"} 0 - loki_write_dropped_entries_total{host="__HOST__",reason="line_too_long",tenant="tenant-default"} 0 loki_write_dropped_entries_total{host="__HOST__",reason="queue_is_full",tenant="tenant-default"} 0 loki_write_dropped_entries_total{host="__HOST__", reason="rate_limited", tenant="tenant-default"} 0 loki_write_dropped_entries_total{host="__HOST__",reason="stream_limited",tenant="tenant-default"} 0 @@ -292,12 +295,12 @@ func TestEndpoint(t *testing.T) { loki_write_sent_entries_total{host="__HOST__",tenant="tenant-default"} 1.0 # HELP loki_write_dropped_entries_total Number of log entries dropped because failed to be sent to the ingester after all retries. # TYPE loki_write_dropped_entries_total counter + loki_write_dropped_entries_total{host="__HOST__",reason="batch_too_large",tenant="tenant-1"} 0 + loki_write_dropped_entries_total{host="__HOST__",reason="batch_too_large",tenant="tenant-2"} 0 + loki_write_dropped_entries_total{host="__HOST__",reason="batch_too_large",tenant="tenant-default"} 0 loki_write_dropped_entries_total{host="__HOST__",reason="ingester_error",tenant="tenant-1"} 0 loki_write_dropped_entries_total{host="__HOST__",reason="ingester_error",tenant="tenant-2"} 0 loki_write_dropped_entries_total{host="__HOST__",reason="ingester_error",tenant="tenant-default"} 0 - loki_write_dropped_entries_total{host="__HOST__",reason="line_too_long",tenant="tenant-1"} 0 - loki_write_dropped_entries_total{host="__HOST__",reason="line_too_long",tenant="tenant-2"} 0 - loki_write_dropped_entries_total{host="__HOST__",reason="line_too_long",tenant="tenant-default"} 0 loki_write_dropped_entries_total{host="__HOST__",reason="queue_is_full",tenant="tenant-1"} 0 loki_write_dropped_entries_total{host="__HOST__",reason="queue_is_full",tenant="tenant-2"} 0 loki_write_dropped_entries_total{host="__HOST__",reason="queue_is_full",tenant="tenant-default"} 0 @@ -341,7 +344,7 @@ func TestEndpoint(t *testing.T) { // Send all the input log entries for i, logEntry := range tt.inputEntries { - c.enqueue(logEntry, 0) + c.enqueue(t.Context(), logEntry, 0) if tt.inputDelay > 0 && i < len(tt.inputEntries)-1 { time.Sleep(tt.inputDelay) @@ -402,14 +405,14 @@ func TestEndpointBlockOnOverflow(t *testing.T) { // NOTE: We have configured batch size to 1 so only one entry will fit in each batch. // To exceed the queue's capacity we need to pass 4 entries. We have one batch that we are actively trying // to send, one batch that is queued and one batch that we are currently working with filling up. - require.NoError(t, e.enqueue(entry, 0)) - require.NoError(t, e.enqueue(entry, 0)) + require.NoError(t, e.enqueue(t.Context(), entry, 0)) + require.NoError(t, e.enqueue(t.Context(), entry, 0)) // Which enqueue fails depends on whether the shard worker has already // consumed the queued batch after the third call. If the third call loses that race, // it returns errQueueIsFull, otherwise the fourth call does. - err3 := e.enqueue(entry, 0) - err4 := e.enqueue(entry, 0) + err3 := e.enqueue(t.Context(), entry, 0) + err4 := e.enqueue(t.Context(), entry, 0) queueIsFull := errors.Is(err3, errQueueIsFull) || errors.Is(err4, errQueueIsFull) require.True(t, queueIsFull, "expected either the third or fourth enqueue to fail with queue full") }) @@ -448,10 +451,10 @@ func TestEndpointBlockOnOverflow(t *testing.T) { // We just need to finish one request in order for all entries to be successfully enqueued. <-receivedReqsChan }() - require.NoError(t, e.enqueue(entry1, 0)) - require.NoError(t, e.enqueue(entry2, 0)) - require.NoError(t, e.enqueue(entry3, 0)) - require.NoError(t, e.enqueue(entry4, 0)) + require.NoError(t, e.enqueue(t.Context(), entry1, 0)) + require.NoError(t, e.enqueue(t.Context(), entry2, 0)) + require.NoError(t, e.enqueue(t.Context(), entry3, 0)) + require.NoError(t, e.enqueue(t.Context(), entry4, 0)) }) } @@ -483,7 +486,7 @@ func TestEndpointBatchSizeMetric(t *testing.T) { require.NoError(t, err) for _, entry := range entries { - require.NoError(t, e.enqueue(entry, 0)) + require.NoError(t, e.enqueue(t.Context(), entry, 0)) } // Stopping the endpoint waits until the current batch is sent. @@ -498,6 +501,98 @@ func TestEndpointBatchSizeMetric(t *testing.T) { assert.Equal(t, float64(entries[0].Size()+entries[1].Size()), sum) } +func TestEndpointCallerCancel(t *testing.T) { + t.Run("entry is not queued or sent if callers context is canceled", func(t *testing.T) { + called := atomic.NewBool(false) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called.Store(true) + })) + defer server.Close() + + var url flagext.URLValue + require.NoError(t, url.Set(server.URL)) + + e, err := newEndpoint(newMetrics(prometheus.NewRegistry()), Config{URL: url}, logging.NewSlogNop(), marker.NewNopTracker()) + require.NoError(t, err) + defer e.stop() + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + require.ErrorIs(t, e.enqueue(ctx, loki.Entry{Entry: push.Entry{Line: "my entry"}}, 0), context.Canceled) + require.Equal(t, false, called.Load()) + }) + + t.Run("when queue is full and callers context is done", func(t *testing.T) { + server, blocked, release := newBlockedServer() + defer server.Close() + defer release() + + var url flagext.URLValue + require.NoError(t, url.Set(server.URL)) + + e, err := newEndpoint(newMetrics(prometheus.NewRegistry()), Config{ + Name: "test-client", + URL: url, + Timeout: time.Minute, + BatchSize: 1, + BackoffConfig: backoff.Config{ + MinBackoff: time.Millisecond, + MaxBackoff: 10 * time.Millisecond, + MaxRetries: 0, + }, + QueueConfig: QueueConfig{ + Capacity: 1, + MinShards: 1, + DrainTimeout: 1 * time.Second, + BlockOnOverflow: true, + }, + }, logging.NewSlogNop(), marker.NewNopTracker()) + require.NoError(t, err) + + defer e.stop() + + entry := loki.Entry{Entry: push.Entry{Line: "my entry"}} + + require.NoError(t, e.enqueue(t.Context(), entry, 0)) + require.Eventually(t, func() bool { return blocked.Load() }, 3*time.Second, 100*time.Millisecond) + + require.NoError(t, e.enqueue(t.Context(), entry, 0)) + require.NoError(t, e.enqueue(t.Context(), entry, 0)) + + ctx, cancel := context.WithTimeout(t.Context(), 1*time.Second) + defer cancel() + + done := make(chan error, 1) + go func() { + done <- e.enqueue(ctx, entry, 0) + }() + + select { + case err := <-done: + require.ErrorIs(t, err, context.DeadlineExceeded) + case <-time.After(2 * time.Second): + t.Fatal("entry trying to be queued was not canceled in time") + } + }) +} + +func TestEndpointStopped(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) + var url flagext.URLValue + require.NoError(t, url.Set(srv.URL)) + + e, err := newEndpoint(newMetrics(prometheus.NewRegistry()), Config{ + URL: url, + }, logging.NewSlogNop(), marker.NewNopTracker()) + require.NoError(t, err) + e.stop() + + entry := loki.Entry{Entry: push.Entry{Line: "my entry"}} + require.ErrorIs(t, e.enqueue(t.Context(), entry, 0), loki.ErrConsumerStopped) +} + // histogramSumAndCount returns the sum and count of the single series of the // named histogram in reg. func histogramSumAndCount(t *testing.T, reg *prometheus.Registry, name string) (float64, uint64) { diff --git a/internal/component/common/loki/client/metrics.go b/internal/component/common/loki/client/metrics.go index 1470072ab1d..9f0e68d4782 100644 --- a/internal/component/common/loki/client/metrics.go +++ b/internal/component/common/loki/client/metrics.go @@ -15,12 +15,11 @@ const ( reasonGeneric = "ingester_error" reasonRateLimited = "rate_limited" reasonStreamLimited = "stream_limited" - reasonLineTooLong = "line_too_long" reasonQueueIsFull = "queue_is_full" reasonBatchTooLarge = "batch_too_large" ) -var reasons = []string{reasonGeneric, reasonRateLimited, reasonStreamLimited, reasonLineTooLong, reasonQueueIsFull} +var reasons = []string{reasonGeneric, reasonRateLimited, reasonStreamLimited, reasonQueueIsFull, reasonBatchTooLarge} type metrics struct { sentBytes *prometheus.CounterVec diff --git a/internal/component/common/loki/client/shards.go b/internal/component/common/loki/client/shards.go index 85321af8f00..96dda2aa546 100644 --- a/internal/component/common/loki/client/shards.go +++ b/internal/component/common/loki/client/shards.go @@ -361,13 +361,21 @@ func (s *shards) runShard(q *queue) { } } +var errQueueIsFull = errors.New("queue is full") + // enqueue routes a log entry to the appropriate shard based on its label fingerprint. -// Returns false if we could not enqueue the entry, either because the shard is shutting down or the queue is full. -// It is up to the caller to retry or drop the entry. -func (s *shards) enqueue(tenantID string, entry loki.Entry, segmentNum int) bool { +// Returns loki.ErrConsumerStopped if the shard is shutting down, in which case retrying will never +// succeed, and errQueueIsFull if the queue is full, which the caller may retry or drop the entry. +func (s *shards) enqueue(tenantID string, entry loki.Entry, segmentNum int) error { s.mut.Lock() defer s.mut.Unlock() + select { + case <-s.softShutdown: + return loki.ErrConsumerStopped + default: + } + if _, ok := s.tenants[tenantID]; !ok { s.tenants[tenantID] = struct{}{} s.initBatchMetrics(tenantID) @@ -376,12 +384,10 @@ func (s *shards) enqueue(tenantID string, entry loki.Entry, segmentNum int) bool fingerprint := entry.Labels.FastFingerprint() shard := uint64(fingerprint) % uint64(len(s.queues)) - select { - case <-s.softShutdown: - return false - default: - return s.queues[shard].append(tenantID, entry, segmentNum) + if !s.queues[shard].append(tenantID, entry, segmentNum) { + return errQueueIsFull } + return nil } func (s *shards) initBatchMetrics(tenantID string) { diff --git a/internal/component/common/loki/client/util_test.go b/internal/component/common/loki/client/util_test.go new file mode 100644 index 00000000000..0444d0b7c57 --- /dev/null +++ b/internal/component/common/loki/client/util_test.go @@ -0,0 +1,59 @@ +package client + +import ( + "context" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/grafana/loki/pkg/push" + "github.com/prometheus/common/model" + "go.uber.org/atomic" + + "github.com/grafana/alloy/internal/component/common/loki" +) + +func newBlockedServer() (*httptest.Server, *atomic.Bool, func()) { + var ( + done = make(chan struct{}) + doneOnce sync.Once + blocked = atomic.NewBool(false) + ) + + release := func() { doneOnce.Do(func() { close(done) }) } + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + blocked.Store(true) + select { + case <-done: + case <-r.Context().Done(): + } + })) + + return server, blocked, release +} + +func feedUntilBlocked(t *testing.T, blocked *atomic.Bool, consumer Consumer) { + t.Helper() + + const timeout = 10 * time.Second + + e := loki.NewEntry(model.LabelSet{"A": "b"}, push.Entry{ + Line: "test", + Timestamp: time.Now(), + }) + + deadline := time.Now().Add(timeout) + for !blocked.Load() { + if time.Now().After(deadline) { + t.Fatalf("endpoint did not block within %s", timeout) + } + + ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) + + _ = consumer.ConsumeEntry(ctx, e) + cancel() + } +} diff --git a/internal/component/common/loki/wal/internal/watcher_state.go b/internal/component/common/loki/wal/internal/watcher_state.go index b416a094445..56d64486b45 100644 --- a/internal/component/common/loki/wal/internal/watcher_state.go +++ b/internal/component/common/loki/wal/internal/watcher_state.go @@ -1,6 +1,7 @@ package internal import ( + "context" "log/slog" "sync" ) @@ -21,17 +22,23 @@ const ( // WatcherState is a holder for the state the Watcher is in. It provides handy methods for checking it it's stopping, getting // the current state, or blocking until it has stopped. type WatcherState struct { - current int - mut sync.RWMutex - stoppingSignal chan struct{} - logger *slog.Logger + current int + mut sync.RWMutex + logger *slog.Logger + + // ctx is canceled when the state transitions to StateStopping. + ctx context.Context + cancel context.CancelFunc } func NewWatcherState(logger *slog.Logger) *WatcherState { + ctx, cancel := context.WithCancel(context.Background()) + return &WatcherState{ - current: StateRunning, - stoppingSignal: make(chan struct{}), - logger: logger, + current: StateRunning, + logger: logger, + ctx: ctx, + cancel: cancel, } } @@ -42,10 +49,9 @@ func (s *WatcherState) Transition(next int) { s.logger.Debug("watcher transitioning state", "currentState", printState(s.current), "nextState", printState(next)) - // only perform channel close if the state is not already stopping - // expect s.s to be either draining ro running to perform a close + // only cancel context if the state is not already StateStopping. if next == StateStopping && s.current != next { - close(s.stoppingSignal) + s.cancel() } // update state @@ -66,9 +72,9 @@ func (s *WatcherState) IsStopping() bool { return s.current == StateStopping } -// WaitForStopping returns a channel in which the called can read, effectively waiting until the state changes to stopping. -func (s *WatcherState) WaitForStopping() <-chan struct{} { - return s.stoppingSignal +// StoppingContext returns a context that will be canceled when state is transitioned to StateStopping. +func (s *WatcherState) StoppingContext() context.Context { + return s.ctx } // printState prints a user-friendly name of the possible Watcher states. diff --git a/internal/component/common/loki/wal/watcher.go b/internal/component/common/loki/wal/watcher.go index 515218ed152..86a27d2ba91 100644 --- a/internal/component/common/loki/wal/watcher.go +++ b/internal/component/common/loki/wal/watcher.go @@ -1,6 +1,7 @@ package wal import ( + "context" "errors" "fmt" "io" @@ -53,7 +54,9 @@ type WriteTo interface { // found in. StoreSeries(series []record.RefSeries, segmentNum int) - AppendEntries(entries RefEntries, segmentNum int) error + // AppendEntries is called with the entries read from a segment. ctx is tied to the lifetime of the Watcher and is + // canceled when it stops, so an implementation that blocks must honour it or it will hold up shutdown. + AppendEntries(ctx context.Context, entries RefEntries, segmentNum int) error } // Marker allows the Watcher to start from a specific segment in the WAL. @@ -132,12 +135,12 @@ func (w *Watcher) mainLoop() { if w.state.IsDraining() && errors.Is(err, os.ErrNotExist) { w.logger.Info("reached non existing segment while draining, assuming end of WAL") // since we've reached the end of the WAL, and the Watcher is draining, promptly transition to stopping state - // so the watcher can stoppingSignal early + // so the watcher can stop early w.state.Transition(internal.StateStopping) } select { - case <-w.state.WaitForStopping(): + case <-w.state.StoppingContext().Done(): return case <-time.After(5 * time.Second): } @@ -222,7 +225,7 @@ func (w *Watcher) watch(segmentNum int, tail bool) error { for { select { - case <-w.state.WaitForStopping(): + case <-w.state.StoppingContext().Done(): return nil case <-segmentTicker.C: @@ -244,8 +247,14 @@ func (w *Watcher) watch(segmentNum int, tail bool) error { // We know that there's either a new segment (last > segmentNum), or we are draining the WAL. Either case, read // the remaining data from the segmentNum and return from `watch` to read the next one. + _, err = w.readSegment(reader, segmentNum) + // context.Canceled means the Watcher is stopping. + if errors.Is(err, context.Canceled) { + return nil + } + // io.EOF error are non-fatal since we are consuming the segment till the end if !errors.Is(err, io.EOF) { return err @@ -264,6 +273,12 @@ func (w *Watcher) watch(segmentNum int, tail bool) error { // read from open segment routine ok, err := w.readSegment(reader, segmentNum) + + // context.Canceled means the Watcher is stopping. + if errors.Is(err, context.Canceled) { + return nil + } + // Ignore all errors reading to end of segment whilst replaying the WAL. This is because when replaying not the // last segment, we assume that segment is not written anymore (closed), and the call to readSegment will read // to the end of it. If error, log a warning accordingly. After, error or no error, nil is returned so that the @@ -315,7 +330,7 @@ func (w *Watcher) readSegment(r *wlog.LiveReader, segmentNum int) (bool, error) } // decodeAndDispatch first decodes a WAL record. Upon reading either Series or Entries from the WAL record, call the -// appropriate callbacks in the writeTo. +// appropriate function in the writeTo. func (w *Watcher) decodeAndDispatch(b []byte, segmentNum int) (bool, error) { var readData bool @@ -333,7 +348,15 @@ func (w *Watcher) decodeAndDispatch(b []byte, segmentNum int) (bool, error) { readData = true for _, entries := range rec.RefEntries { - if err := w.actions.AppendEntries(entries, segmentNum); err != nil && firstErr == nil { + err := w.actions.AppendEntries(w.state.StoppingContext(), entries, segmentNum) + + // The WriteTo gave up because the Watcher is stopping, so there is no point handing it + // the rest of the record. The caller turns this into a quiet exit. + if errors.Is(err, context.Canceled) { + return readData, err + } + + if err != nil && firstErr == nil { firstErr = err } } @@ -351,12 +374,14 @@ func (w *Watcher) Drain() { select { case <-time.NewTimer(w.drainTimeout).C: w.logger.Warn("watcher drain timeout occurred, transitioning to Stopping") - case <-w.state.WaitForStopping(): + case <-w.state.StoppingContext().Done(): } } // Stop stops the Watcher, shutting down the main routine. func (w *Watcher) Stop() { + // Transitioning cancels the state's context, which releases a WriteTo blocked in + // AppendEntries so the wait below can return. w.state.Transition(internal.StateStopping) // upon calling stop, wait for main mainLoop execution to stop diff --git a/internal/component/common/loki/wal/watcher_test.go b/internal/component/common/loki/wal/watcher_test.go index 74f04b72a9c..ae82ee2b05a 100644 --- a/internal/component/common/loki/wal/watcher_test.go +++ b/internal/component/common/loki/wal/watcher_test.go @@ -2,6 +2,7 @@ package wal import ( "bytes" + "context" "encoding/binary" "fmt" "hash/crc32" @@ -40,7 +41,7 @@ func (t *testWriteTo) SeriesReset(segmentNum int) { t.ReceivedSeriesReset = append(t.ReceivedSeriesReset, segmentNum) } -func (t *testWriteTo) AppendEntries(entries RefEntries, _ int) error { +func (t *testWriteTo) AppendEntries(_ context.Context, entries RefEntries, _ int) error { var entry loki.Entry if l, ok := t.series[uint64(entries.Ref)]; ok { entry.Labels = l @@ -356,7 +357,7 @@ func TestWatcher(t *testing.T) { t, &watcherTestResources{ writeEntry: func(entry loki.Entry) { - _ = ew.WriteEntry(entry, wl) + _ = ew.writeEntry(entry, wl) }, notifyWrite: func() { watcher.NotifyWrite() @@ -432,7 +433,7 @@ func TestWatcher_Replay(t *testing.T) { ew := newEntryWriter() // First, write to segment 0. This will be the last "marked" segment - err = ew.WriteEntry(loki.Entry{ + err = ew.writeEntry(loki.Entry{ Labels: labels, Entry: push.Entry{ Timestamp: time.Now(), @@ -447,7 +448,7 @@ func TestWatcher_Replay(t *testing.T) { // Now, write to segment 1, this will be a segment not marked, hence replayed for _, line := range segment1Lines { - err = ew.WriteEntry(loki.Entry{ + err = ew.writeEntry(loki.Entry{ Labels: labels, Entry: push.Entry{ Timestamp: time.Now(), @@ -463,7 +464,7 @@ func TestWatcher_Replay(t *testing.T) { // Finally, write some data to the last segment, this will be the write head for _, line := range segment2Lines { - err = ew.WriteEntry(loki.Entry{ + err = ew.writeEntry(loki.Entry{ Labels: labels, Entry: push.Entry{ Timestamp: time.Now(), @@ -514,7 +515,7 @@ func TestWatcher_Replay(t *testing.T) { ew := newEntryWriter() // First, write to segment 0. This will be the last "marked" segment - err = ew.WriteEntry(loki.Entry{ + err = ew.writeEntry(loki.Entry{ Labels: labels, Entry: push.Entry{ Timestamp: time.Now(), @@ -529,7 +530,7 @@ func TestWatcher_Replay(t *testing.T) { // Now, write to segment 1, this will be a segment not marked, hence replayed for _, line := range segment1Lines { - err = ew.WriteEntry(loki.Entry{ + err = ew.writeEntry(loki.Entry{ Labels: labels, Entry: push.Entry{ Timestamp: time.Now(), @@ -551,7 +552,7 @@ func TestWatcher_Replay(t *testing.T) { // Write something after watcher started for _, line := range segment2Lines { - err = ew.WriteEntry(loki.Entry{ + err = ew.writeEntry(loki.Entry{ Labels: labels, Entry: push.Entry{ Timestamp: time.Now(), @@ -584,7 +585,7 @@ func (s *slowWriteTo) SeriesReset(segmentNum int) { func (s *slowWriteTo) StoreSeries(series []record.RefSeries, segmentNum int) { } -func (s *slowWriteTo) AppendEntries(entries RefEntries, segmentNum int) error { +func (s *slowWriteTo) AppendEntries(_ context.Context, entries RefEntries, segmentNum int) error { s.entriesReceived.Add(uint64(len(entries.Entries))) time.Sleep(s.sleepAfterAppendEntries) return nil @@ -641,7 +642,7 @@ func TestWatcher_StopAndDrainWAL(t *testing.T) { writeNLines := func(t *testing.T, n int) { for range n { // First, write to segment 0. This will be the last "marked" segment - err := ew.WriteEntry(loki.Entry{ + err := ew.writeEntry(loki.Entry{ Labels: labels, Entry: push.Entry{ Timestamp: time.Now(), @@ -693,7 +694,7 @@ func TestWatcher_StopAndDrainWAL(t *testing.T) { writeNLines := func(t *testing.T, n int) { for range n { // First, write to segment 0. This will be the last "marked" segment - err := ew.WriteEntry(loki.Entry{ + err := ew.writeEntry(loki.Entry{ Labels: labels, Entry: push.Entry{ Timestamp: time.Now(), @@ -746,7 +747,7 @@ func TestWatcher_StopAndDrainWAL(t *testing.T) { writeNLines := func(t *testing.T, n int) { for range n { // First, write to segment 0. This will be the last "marked" segment - err := ew.WriteEntry(loki.Entry{ + err := ew.writeEntry(loki.Entry{ Labels: labels, Entry: push.Entry{ Timestamp: time.Now(), diff --git a/internal/component/common/loki/wal/writer.go b/internal/component/common/loki/wal/writer.go index 4d1ff03436f..39d3ab11f87 100644 --- a/internal/component/common/loki/wal/writer.go +++ b/internal/component/common/loki/wal/writer.go @@ -37,16 +37,16 @@ type WriteEventSubscriber interface { NotifyWrite() } -// Writer implements loki.EntryHandler, exposing a channel were scraping targets can write to. Reading from there, it -// writes incoming entries to a WAL. -// Also, since Writer is responsible for all changing operations over the WAL, therefore a routine is run for cleaning -// old segments. +// Writer writes log entries to a write-ahead log. It is also responsible for the WAL's segments, +// removing those older than the configured maximum age and notifying subscribers when it does. type Writer struct { - entries chan loki.Entry - logger *slog.Logger - wg sync.WaitGroup - once sync.Once - wal WAL + logger *slog.Logger + cfg Config + wg sync.WaitGroup + wal WAL + + writeMut sync.Mutex + stopped bool entryWriter *entryWriter cleanupSubscribersLock sync.RWMutex @@ -57,7 +57,7 @@ type Writer struct { metrics *WriterMetrics - closeCleaner chan struct{} + done chan struct{} } // NewWriter creates a new Writer. @@ -72,76 +72,78 @@ func NewWriter(walCfg Config, logger *slog.Logger, reg prometheus.Registerer, me } wrt := &Writer{ - entries: make(chan loki.Entry), - logger: logger, - wg: sync.WaitGroup{}, - wal: wl, - entryWriter: newEntryWriter(), - closeCleaner: make(chan struct{}, 1), - metrics: metrics, + logger: logger, + entryWriter: newEntryWriter(), + wg: sync.WaitGroup{}, + cfg: walCfg, + wal: wl, + done: make(chan struct{}), + metrics: metrics, } return wrt, nil } -func (wrt *Writer) Start(maxSegmentAge time.Duration) { - // main WAL writer routine - wrt.wg.Go(func() { - for e := range wrt.entries { - if err := wrt.entryWriter.WriteEntry(e, wrt.wal); err != nil { - wrt.logger.Error("failed to write entry", "err", err) - // if an error occurred while writing the wal, go to next entry and don't notify write subscribers - continue - } - - // emit metric with latest written timestamp, to be able to track delay from writer to watcher - wrt.metrics.lastWrittenTimestamp.WithLabelValues().Set(float64(e.Timestamp.Unix())) - - wrt.writeSubscribersLock.RLock() - for _, s := range wrt.writeSubscribers { - s.NotifyWrite() - } - wrt.writeSubscribersLock.RUnlock() - } - }) - +func (wrt *Writer) Start() { // WAL cleanup routine that cleans old segments wrt.wg.Go(func() { // By cleaning every 10th of the configured threshold for considering a segment old, we are allowing a maximum slip // of 10%. If the configured time is 1 hour, that'd be 6 minutes. - triggerEvery := maxSegmentAge / 10 + triggerEvery := wrt.cfg.MaxSegmentAge / 10 if triggerEvery < minimumCleanSegmentsEvery { triggerEvery = minimumCleanSegmentsEvery } trigger := time.NewTicker(triggerEvery) + defer trigger.Stop() for { select { case <-trigger.C: wrt.logger.Debug("Running wal old segments cleanup") - if err := wrt.cleanSegments(maxSegmentAge); err != nil { + if err := wrt.cleanSegments(wrt.cfg.MaxSegmentAge); err != nil { wrt.logger.Error("Error cleaning old segments", "err", err) } - case <-wrt.closeCleaner: - trigger.Stop() + case <-wrt.done: return } } }) } -func (wrt *Writer) Chan() chan<- loki.Entry { - return wrt.entries +func (wrt *Writer) WriteEntry(entry loki.Entry) error { + wrt.writeMut.Lock() + defer wrt.writeMut.Unlock() + + if wrt.stopped { + return loki.ErrConsumerStopped + } + + if err := wrt.entryWriter.writeEntry(entry, wrt.wal); err != nil { + wrt.logger.Error("failed to write entry", "err", err) + return err + } + + // emit metric with latest written timestamp, to be able to track delay from writer to watcher + wrt.metrics.lastWrittenTimestamp.WithLabelValues().Set(float64(entry.Timestamp.Unix())) + + wrt.writeSubscribersLock.RLock() + for _, s := range wrt.writeSubscribers { + s.NotifyWrite() + } + wrt.writeSubscribersLock.RUnlock() + + return nil } func (wrt *Writer) Stop() { - wrt.once.Do(func() { - close(wrt.entries) - }) - // close cleaner routine - wrt.closeCleaner <- struct{}{} - // Wait for routine to write to wal all pending entries + wrt.writeMut.Lock() + defer wrt.writeMut.Unlock() + wrt.stopped = true + + // Stop cleaner routine and wait for it to stop. + close(wrt.done) wrt.wg.Wait() - // Close WAL to finalize all pending writes + + // Close WAL to finalize all pending writes. wrt.wal.Close() } @@ -224,17 +226,19 @@ func newEntryWriter() *entryWriter { } } -// WriteEntry writes a loki.Entry to a WAL. Note that since it's re-using the same Record object for every -// write, it first has to be reset, and then overwritten accordingly. Therefore, WriteEntry is not thread-safe. -func (ew *entryWriter) WriteEntry(entry loki.Entry, wl WAL) error { - // Reset wal record slices - ew.reusableWALRecord.Reset() +func (ew *entryWriter) writeEntry(entry loki.Entry, wl WAL) error { + defer ew.reusableWALRecord.Reset() var fp uint64 lbs := labels.FromMap(util.ModelLabelSetToMap(entry.Labels)) fp, _ = lbs.HashWithoutLabels(nil, []string(nil)...) ref := chunks.HeadSeriesRef(fp) + ew.reusableWALRecord.Series = append(ew.reusableWALRecord.Series, record.RefSeries{ + Ref: ref, + Labels: lbs, + }) + // Append the entry to an already existing stream (if any) ew.reusableWALRecord.RefEntries = append(ew.reusableWALRecord.RefEntries, RefEntries{ Ref: ref, @@ -243,10 +247,6 @@ func (ew *entryWriter) WriteEntry(entry loki.Entry, wl WAL) error { }, Created: entry.Created(), }) - ew.reusableWALRecord.Series = append(ew.reusableWALRecord.Series, record.RefSeries{ - Ref: ref, - Labels: lbs, - }) return wl.Log(ew.reusableWALRecord) } diff --git a/internal/component/common/loki/wal/writer_test.go b/internal/component/common/loki/wal/writer_test.go index c6a4e1d09be..0b2a19e0b22 100644 --- a/internal/component/common/loki/wal/writer_test.go +++ b/internal/component/common/loki/wal/writer_test.go @@ -22,50 +22,87 @@ import ( autil "github.com/grafana/alloy/internal/util" ) -func TestWriter_EntriesAreWrittenToWAL(t *testing.T) { - var ( - dir = t.TempDir() - reg = prometheus.NewRegistry() - ) +func TestWriter(t *testing.T) { + t.Run("entries are written to wal", func(t *testing.T) { + var ( + dir = t.TempDir() + reg = prometheus.NewRegistry() + ) + + writer, err := NewWriter(Config{ + Dir: dir, + Enabled: true, + MaxSegmentAge: time.Minute, + }, autil.TestAlloyLogger(t).Slog(), reg, NewWriterMetrics(reg)) + require.NoError(t, err) + defer writer.Stop() + + // write entries to wal and sync + writer.Start() + + var testLabels = model.LabelSet{ + "testing": "log", + } + var lines = []string{ + "some line", + "some other line", + "some other other line", + } - writer, err := NewWriter(Config{ - Dir: dir, - Enabled: true, - MaxSegmentAge: time.Minute, - }, autil.TestAlloyLogger(t).Slog(), reg, NewWriterMetrics(reg)) - require.NoError(t, err) - defer func() { - writer.Stop() - }() - // write entries to wal and sync - writer.Start(time.Minute) + for _, line := range lines { + require.NoError(t, writer.WriteEntry( + loki.Entry{ + Labels: testLabels, + Entry: push.Entry{ + Timestamp: time.Now(), + Line: line, + }, + }, + )) + } - var testLabels = model.LabelSet{ - "testing": "log", - } - var lines = []string{ - "some line", - "some other line", - "some other other line", - } + // accessing the WAL inside, just for testing! + require.NoError(t, writer.wal.Sync(), "failed to sync wal") - for _, line := range lines { - writer.Chan() <- loki.Entry{ - Labels: testLabels, + // assert over WAL entries + readEntries := eventuallyReadWAL(t, len(lines), dir) + require.NotNil(t, readEntries) + require.Equal(t, testLabels, readEntries[0].Labels) + }) + + t.Run("stopped writer returns error", func(t *testing.T) { + var ( + dir = t.TempDir() + reg = prometheus.NewRegistry() + ) + + writer, err := NewWriter(Config{ + Dir: dir, + Enabled: true, + MaxSegmentAge: time.Minute, + }, logging.NewSlogNop(), reg, NewWriterMetrics(reg)) + + require.NoError(t, err) + writer.Start() + + require.NoError(t, writer.WriteEntry(loki.Entry{ + Labels: model.LabelSet{"key": "value"}, Entry: push.Entry{ Timestamp: time.Now(), - Line: line, + Line: "lin", }, - } - } + })) - // accessing the WAL inside, just for testing! - require.NoError(t, writer.wal.Sync(), "failed to sync wal") + writer.Stop() - // assert over WAL entries - readEntries := eventuallyReadWAL(t, len(lines), dir) - require.NotNil(t, readEntries) - require.Equal(t, testLabels, readEntries[0].Labels) + require.ErrorIs(t, writer.WriteEntry(loki.Entry{ + Labels: model.LabelSet{"key": "value"}, + Entry: push.Entry{ + Timestamp: time.Now(), + Line: "lin", + }, + }), loki.ErrConsumerStopped) + }) } func TestWriter_MetricsWorkAfterRecreation(t *testing.T) { @@ -81,9 +118,9 @@ func TestWriter_MetricsWorkAfterRecreation(t *testing.T) { }, logging.NewSlogNop(), reg, NewWriterMetrics(reg)) require.NoError(t, err) - writer.Start(time.Minute) + writer.Start() entry := loki.NewEntry(model.LabelSet{"foo": "bar"}, push.Entry{Timestamp: time.Now(), Line: "line"}) - writer.Chan() <- entry + writer.WriteEntry(entry) expected := fmt.Sprintf(` # HELP loki_write_wal_writer_last_written_timestamp Latest timestamp that was written to the WAL @@ -103,11 +140,11 @@ func TestWriter_MetricsWorkAfterRecreation(t *testing.T) { MaxSegmentAge: time.Minute, }, logging.NewSlogNop(), reg, NewWriterMetrics(reg)) require.NoError(t, err) - writer.Start(time.Minute) + writer.Start() defer writer.Stop() newEntry := loki.NewEntry(model.LabelSet{"foo": "bar"}, push.Entry{Timestamp: time.Now().Add(1 * time.Second), Line: "line"}) - writer.Chan() <- newEntry + writer.WriteEntry(newEntry) expected = fmt.Sprintf(` # HELP loki_write_wal_writer_last_written_timestamp Latest timestamp that was written to the WAL @@ -144,10 +181,8 @@ func TestWriter_OldSegmentsAreCleanedUp(t *testing.T) { MaxSegmentAge: maxSegmentAge, }, autil.TestAlloyLogger(t).Slog(), reg, NewWriterMetrics(reg)) require.NoError(t, err) - writer.Start(maxSegmentAge) - defer func() { - writer.Stop() - }() + defer writer.Stop() + writer.Start() notificationMutex := sync.Mutex{} // add writer events subscriber. Add multiple to test fanout @@ -173,13 +208,15 @@ func TestWriter_OldSegmentsAreCleanedUp(t *testing.T) { } for _, line := range lines { - writer.Chan() <- loki.Entry{ - Labels: testLabels, - Entry: push.Entry{ - Timestamp: time.Now(), - Line: line, + require.NoError(t, writer.WriteEntry( + loki.Entry{ + Labels: testLabels, + Entry: push.Entry{ + Timestamp: time.Now(), + Line: line, + }, }, - } + )) } // accessing the WAL inside, just for testing! @@ -228,7 +265,7 @@ func TestWriter_NoSegmentIsCleanedUpIfTheresOnlyOne(t *testing.T) { var ( dir = t.TempDir() reg = prometheus.NewRegistry() - maxSegmentAge = time.Second * 2 + maxSegmentAge = 2 * time.Second segmentsReclaimedNotificationsReceived = []int{} ) @@ -238,7 +275,7 @@ func TestWriter_NoSegmentIsCleanedUpIfTheresOnlyOne(t *testing.T) { MaxSegmentAge: maxSegmentAge, }, autil.TestAlloyLogger(t).Slog(), reg, NewWriterMetrics(reg)) require.NoError(t, err) - writer.Start(maxSegmentAge) + writer.Start() defer func() { writer.Stop() }() @@ -257,13 +294,15 @@ func TestWriter_NoSegmentIsCleanedUpIfTheresOnlyOne(t *testing.T) { } for _, line := range lines { - writer.Chan() <- loki.Entry{ - Labels: testLabels, - Entry: push.Entry{ - Timestamp: time.Now(), - Line: line, + require.NoError(t, writer.WriteEntry( + loki.Entry{ + Labels: testLabels, + Entry: push.Entry{ + Timestamp: time.Now(), + Line: line, + }, }, - } + )) } // accessing the WAL inside, just for testing! @@ -413,20 +452,22 @@ func benchWriteEntries(b *testing.B, lines, labelSetCount int) { MaxSegmentAge: time.Minute, }, logging.NewSlogNop(), reg, NewWriterMetrics(reg)) require.NoError(b, err) - writer.Start(time.Minute) + writer.Start() defer func() { writer.Stop() }() for i := 0; i < lines; i++ { - writer.Chan() <- loki.Entry{ - Labels: model.LabelSet{ - "someLabel": model.LabelValue(fmt.Sprint(i % labelSetCount)), + require.NoError(b, writer.WriteEntry( + loki.Entry{ + Labels: model.LabelSet{ + "someLabel": model.LabelValue(fmt.Sprint(i % labelSetCount)), + }, + Entry: push.Entry{ + Timestamp: time.Now(), + Line: fmt.Sprintf("some line being written %d", i), + }, }, - Entry: push.Entry{ - Timestamp: time.Now(), - Line: fmt.Sprintf("some line being written %d", i), - }, - } + )) } } diff --git a/internal/component/loki/source/api/api_test.go b/internal/component/loki/source/api/api_test.go index f0f305f6e45..79d760b289f 100644 --- a/internal/component/loki/source/api/api_test.go +++ b/internal/component/loki/source/api/api_test.go @@ -110,14 +110,13 @@ func TestLokiSourceAPI_Simple(t *testing.T) { defer lokiClient.Stop() now := time.Now() - select { - case lokiClient.Chan() <- loki.Entry{ - Labels: map[model.LabelName]model.LabelValue{"source": "test"}, - Entry: push.Entry{Timestamp: now, Line: "hello world!"}, - }: - case <-ctx.Done(): - t.Fatalf("timed out while sending test entries via loki client") - } + require.NoError(t, lokiClient.ConsumeEntry( + ctx, + loki.NewEntry( + model.LabelSet{"source": "test"}, + push.Entry{Timestamp: now, Line: "hello world!"}, + ), + )) require.Eventually( t, @@ -156,14 +155,13 @@ func TestLokiSourceAPI_Update(t *testing.T) { defer lokiClient.Stop() now := time.Now() - select { - case lokiClient.Chan() <- loki.Entry{ - Labels: map[model.LabelName]model.LabelValue{"source": "test"}, - Entry: push.Entry{Timestamp: now, Line: "hello world!"}, - }: - case <-ctx.Done(): - t.Fatalf("timed out while sending test entries via loki client") - } + require.NoError(t, lokiClient.ConsumeEntry( + ctx, + loki.NewEntry( + model.LabelSet{"source": "test"}, + push.Entry{Timestamp: now, Line: "hello world!"}, + ), + )) require.Eventually( t, @@ -186,14 +184,14 @@ func TestLokiSourceAPI_Update(t *testing.T) { receiver.Clear() - select { - case lokiClient.Chan() <- loki.Entry{ - Labels: map[model.LabelName]model.LabelValue{"source": "test"}, - Entry: push.Entry{Timestamp: now, Line: "hello brave new world!"}, - }: - case <-ctx.Done(): - t.Fatalf("timed out while sending test entries via loki client") - } + require.NoError(t, lokiClient.ConsumeEntry( + ctx, + loki.NewEntry( + model.LabelSet{"source": "test"}, + push.Entry{Timestamp: now, Line: "hello brave new world!"}, + ), + )) + require.Eventually( t, func() bool { return len(receiver.Received()) == 1 }, @@ -238,15 +236,13 @@ func TestLokiSourceAPI_FanOut(t *testing.T) { const messagesCount = 100 for i := range messagesCount { - entry := loki.Entry{ - Labels: map[model.LabelName]model.LabelValue{"source": "test"}, - Entry: push.Entry{Line: fmt.Sprintf("test message #%d", i)}, - } - select { - case lokiClient.Chan() <- entry: - case <-ctx.Done(): - t.Log("timed out while sending test entries via loki client") - } + require.NoError(t, lokiClient.ConsumeEntry( + ctx, + loki.NewEntry( + model.LabelSet{"source": "test"}, + push.Entry{Line: fmt.Sprintf("test message #%d", i)}, + ), + )) } require.Eventually( @@ -380,14 +376,17 @@ func TestLokiSourceAPI_TLS(t *testing.T) { defer lokiClient.Stop() now := time.Now() - select { - case lokiClient.Chan() <- loki.Entry{ - Labels: map[model.LabelName]model.LabelValue{"source": "test"}, - Entry: push.Entry{Timestamp: now, Line: "hello world over TLS!"}, - }: - case <-ctx.Done(): - t.Fatalf("timed out while sending test entries via TLS loki client") - } + + require.NoError( + t, + lokiClient.ConsumeEntry( + ctx, + loki.NewEntry( + model.LabelSet{"source": "test"}, + push.Entry{Timestamp: now, Line: "hello world over TLS!"}, + ), + ), + ) require.Eventually( t, diff --git a/internal/component/loki/write/write.go b/internal/component/loki/write/write.go index 48e1e186be3..79693f9d370 100644 --- a/internal/component/loki/write/write.go +++ b/internal/component/loki/write/write.go @@ -80,18 +80,14 @@ var ( // Component implements the loki.write component. type Component struct { - opts component.Options - - mut sync.RWMutex - args Arguments + opts component.Options receiver loki.LogsReceiver + mut sync.RWMutex + externalLabels model.LabelSet + // remote write consumer consumer client.Consumer - - // sink is the place where log entries received by this component should be written to. - // It will in turn write to client.Consumer. - sink loki.EntryHandler } // New creates a new loki.write component. @@ -116,10 +112,8 @@ func New(o component.Options, args Arguments) (*Component, error) { // Run implements component.Component. func (c *Component) Run(ctx context.Context) error { defer func() { - // First we need to stop the sink. Stopping the sink will not stop the wrapped handler. - if c.sink != nil { - c.sink.Stop() - } + c.mut.Lock() + defer c.mut.Unlock() if c.consumer != nil { if d, ok := c.consumer.(client.DrainableConsumer); ok { @@ -135,15 +129,8 @@ func (c *Component) Run(ctx context.Context) error { select { case <-ctx.Done(): return nil - case entry := <-c.receiver.Chan(): - c.mut.RLock() - select { - case <-ctx.Done(): - c.mut.RUnlock() - return nil - case c.sink.Chan() <- entry: - } - c.mut.RUnlock() + case e := <-c.receiver.Chan(): + c.consumeEntry(ctx, e) } } } @@ -158,17 +145,14 @@ func (c *Component) Update(args component.Arguments) error { c.mut.Lock() defer c.mut.Unlock() - c.args = newArgs - - if c.sink != nil { - c.sink.Stop() - } if c.consumer != nil { // only drain on component shutdown c.consumer.Stop() } + c.externalLabels = util.MapToModelLabelSet(newArgs.ExternalLabels) + cfgs := newArgs.convertEndpointConfigs() uid := alloyseed.Get().UID @@ -199,22 +183,24 @@ func (c *Component) Update(args component.Arguments) error { } if err != nil { - return fmt.Errorf("failed to create cliens: %w", err) + return fmt.Errorf("failed to create clients: %w", err) } - c.sink = newEntryHandler(c.consumer, util.MapToModelLabelSet(c.args.ExternalLabels)) - return nil } -func newEntryHandler(handler loki.EntryHandler, externalLabels model.LabelSet) loki.EntryHandler { - return loki.NewEntryMutatorHandler(handler, func(e loki.Entry) loki.Entry { - if len(externalLabels) == 0 { - return e - } - e.Labels = externalLabels.Merge(e.Labels) - return e - }) +func (c *Component) consumeEntry(ctx context.Context, e loki.Entry) { + c.mut.RLock() + consumer := c.consumer + + if len(c.externalLabels) > 0 { + e.Labels = c.externalLabels.Merge(e.Labels) + } + c.mut.RUnlock() + + // NOTE: For now it's ok to ignore error here. Error mean the consumer is going away, + // either because ctx was canceled or because it has been stopped by a shutdown or an update. + _ = consumer.ConsumeEntry(ctx, e) } func validateConfigStabilityLevel(o component.Options, args Arguments) error { diff --git a/internal/component/loki/write/write_test.go b/internal/component/loki/write/write_test.go index 65fe70e0ac7..ed1b8a02711 100644 --- a/internal/component/loki/write/write_test.go +++ b/internal/component/loki/write/write_test.go @@ -1,6 +1,7 @@ package write import ( + "context" "fmt" "math" "net/http" @@ -309,14 +310,14 @@ func testMultipleEndpoint(t *testing.T, alterArgs func(arguments *Arguments)) { require.FailNow(t, "failed waiting for logs") case req := <-ch1: require.Len(t, req.Streams, 1) - require.Equal(t, req.Streams[0].Labels, wantLabelSet.Clone().Merge(model.LabelSet{"lbl": "foo"}).String()) + require.Equal(t, wantLabelSet.Clone().Merge(model.LabelSet{"lbl": "foo"}).String(), req.Streams[0].Labels) require.Len(t, req.Streams[0].Entries, 1) - require.Equal(t, req.Streams[0].Entries[0].Line, "writing some text") + require.Equal(t, "writing some text", req.Streams[0].Entries[0].Line) case req := <-ch2: require.Len(t, req.Streams, 1) - require.Equal(t, req.Streams[0].Labels, wantLabelSet.Clone().Merge(model.LabelSet{"lbl": "bar"}).String()) + require.Equal(t, wantLabelSet.Clone().Merge(model.LabelSet{"lbl": "bar"}).String(), req.Streams[0].Labels) require.Len(t, req.Streams[0].Entries, 1) - require.Equal(t, req.Streams[0].Entries[0].Line, "writing some text") + require.Equal(t, "writing some text", req.Streams[0].Entries[0].Line) } } } @@ -476,3 +477,110 @@ func benchSingleEndpoint(b *testing.B, tc testCase, alterConfig func(arguments * }, time.Minute, time.Second, "haven't seen expected number of lines") } } + +func TestUpdateWhileSendingIsBlocked(t *testing.T) { + for _, walEnabled := range []bool{false, true} { + t.Run(fmt.Sprintf("wal_enabled=%t", walEnabled), func(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + + ctrl, _ := blockedComponent(t, ctx, walEnabled) + + updated := make(chan error, 1) + go func() { updated <- ctrl.Update(blockedEndpointArgs(t, "http://localhost:1", walEnabled)) }() + + select { + case err := <-updated: + require.NoError(t, err) + case <-time.After(30 * time.Second): + t.Fatal("Update did not return while sending was blocked") + } + }) + } +} + +func TestStopWhileSendingIsBlocked(t *testing.T) { + for _, walEnabled := range []bool{false, true} { + t.Run(fmt.Sprintf("wal_enabled=%t", walEnabled), func(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + _, stopped := blockedComponent(t, ctx, walEnabled) + cancel() + + select { + case err := <-stopped: + require.NoError(t, err) + case <-time.After(30 * time.Second): + t.Fatal("Run did not return while sending was blocked") + } + }) + } +} + +func blockedComponent(t *testing.T, ctx context.Context, walEnabled bool) (*componenttest.Controller, <-chan error) { + t.Helper() + + blocked := atomic.NewBool(false) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + blocked.Store(true) + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(srv.Close) + + ctrl, err := componenttest.NewControllerFromID(logging.NewSlogNop(), "loki.write") + require.NoError(t, err) + + stopped := make(chan error, 1) + go func() { + stopped <- ctrl.Run(ctx, blockedEndpointArgs(t, srv.URL, walEnabled), func(o component.Options) component.Options { + o.MinStability = featuregate.StabilityExperimental + return o + }) + }() + + require.NoError(t, ctrl.WaitExports(5*time.Second)) + + ch := ctrl.Exports().(Exports).Receiver.Chan() + + go func() { + entry := loki.NewEntry(model.LabelSet{"foo": "bar"}, push.Entry{Timestamp: time.Now(), Line: "very important log"}) + for { + select { + case ch <- entry: + case <-ctx.Done(): + return + } + } + }() + + require.Eventually(t, blocked.Load, 10*time.Second, 10*time.Millisecond) + return ctrl, stopped +} +func blockedEndpointArgs(t *testing.T, url string, walEnabled bool) Arguments { + t.Helper() + + cfg := fmt.Sprintf(` + endpoint { + url = "%s" + batch_size = "1B" + min_backoff_period = "1ms" + max_backoff_period = "5ms" + max_backoff_retries = 0 + + queue_config { + capacity = "1B" + min_shards = 1 + drain_timeout = "1s" + block_on_overflow = true + } + } + + wal { + enabled = %t + drain_timeout = "1s" + } + `, url, walEnabled) + + var args Arguments + require.NoError(t, syntax.Unmarshal([]byte(cfg), &args)) + return args +}