[libbeat] Fix a shutdown race in the memory queue (#47248)

Fix a race in the memory queue where an event that was reported as unpublished during queue shutdown could still be enqueued, ingested, and have its acknowledgment callback triggered. In specific circumstances this could cause a panic during Filebeat shutdown, see #47246.

A deterministic test is impossible since the behavior depended on precise timing of an internal select statement with multiple valid paths, but the accompanying unit test fails 90% of the time on the old code, and passes 100% of the time with this change.

Author: faec

changelog/fragments/1761334594-memory-queue-shutdown-fix.yaml

@@ -0,0 +1,32 @@
+# Kind can be one of:
+# - breaking-change: a change to previously-documented behavior
+# - deprecation: functionality that is being removed in a later release
+# - bug-fix: fixes a problem in a previous version
+# - enhancement: extends functionality but does not break or fix existing behavior
+# - feature: new functionality
+# - known-issue: problems that we are aware of in a given version
+# - security: impacts on the security of a product or a user’s deployment.
+# - upgrade: important information for someone upgrading from a prior version
+# - other: does not fit into any of the other categories
+kind: bug-fix
+
+# Change summary; a 80ish characters long description of the change.
+summary: Fix potential Filebeat panic during memory queue shutdown
+
+# Long description; in case the summary is not enough to describe the change
+# this field accommodate a description without length limits.
+# NOTE: This field will be rendered only for breaking-change and known-issue kinds at the moment.
+#description:
+
+# Affected component; usually one of "elastic-agent", "fleet-server", "filebeat", "metricbeat", "auditbeat", "all", etc.
+component: filebeat
+
+# PR URL; optional; the PR number that added the changeset.
+# If not present is automatically filled by the tooling finding the PR where this changelog fragment has been added.
+# NOTE: the tooling supports backports, so it's able to fill the original PR number instead of the backport PR number.
+# Please provide it if you are adding a fragment for a different PR.
+#pr: https://github.com/owner/repo/1234
+
+# Issue URL; optional; the GitHub issue related to this changeset (either closes or is part of).
+# If not present is automatically filled by the tooling with the issue linked to the PR number.
+#issue: https://github.com/owner/repo/1234

libbeat/publisher/queue/memqueue/produce.go

@@ -134,17 +134,7 @@ func (st *openState) publish(req pushRequest) (queue.EntryID, bool) {
 	}
 	select {
 	case st.events <- req:
-		// The events channel is buffered, which means we may successfully
-		// write to it even if the queue is shutting down. To avoid blocking
-		// forever during shutdown, we also have to wait on the queue's
-		// shutdown channel.
-		select {
-		case resp := <-req.resp:
-			return resp, true
-		case <-st.queueClosing:
-			st.events = nil
-			return 0, false
-		}
+		return st.handlePendingResponse(req.resp)
 	case <-st.done:
 		st.events = nil
 		return 0, false
@@ -162,17 +152,7 @@ func (st *openState) tryPublish(req pushRequest) (queue.EntryID, bool) {
 	}
 	select {
 	case st.events <- req:
-		// The events channel is buffered, which means we may successfully
-		// write to it even if the queue is shutting down. To avoid blocking
-		// forever during shutdown, we also have to wait on the queue's
-		// shutdown channel.
-		select {
-		case resp := <-req.resp:
-			return resp, true
-		case <-st.queueClosing:
-			st.events = nil
-			return 0, false
-		}
+		return st.handlePendingResponse(req.resp)
 	case <-st.done:
 		st.events = nil
 		return 0, false
@@ -181,3 +161,31 @@ func (st *openState) tryPublish(req pushRequest) (queue.EntryID, bool) {
 		return 0, false
 	}
 }
+
+func (st *openState) handlePendingResponse(respChan chan queue.EntryID) (queue.EntryID, bool) {
+	// The events channel is buffered, which means we may successfully
+	// write to it even if the queue is shutting down. To avoid blocking
+	// forever during shutdown, we also have to wait on the queue's
+	// shutdown channel.
+	select {
+	case resp := <-respChan:
+		return resp, true
+	case <-st.queueClosing:
+	}
+
+	// Clear the request channel so we can't write to it again
+	st.events = nil
+
+	// Once the queue starts closing, it will not handle any more push
+	// requests, however it may have handled ours before the closing
+	// channel was triggered (and both may have arrived concurrently
+	// at the select statement above). So to know whether our entry was
+	// accepted we also need to check if there's a buffered response in
+	// our channel.
+	select {
+	case resp := <-respChan:
+		return resp, true
+	default:
+	}
+	return 0, false
+}

libbeat/publisher/queue/memqueue/queue_test.go

@@ -18,10 +18,10 @@
 package memqueue
 
 import (
+	"context"
 	"encoding/binary"
 	"flag"
 	"fmt"
-	"math"
 	"math/rand/v2"
 	"sync"
 	"sync/atomic"
@@ -278,32 +278,6 @@ func makeTestQueue(sz, minEvents int, flushTimeout time.Duration) queuetest.Queu
 	}
 }
 
-func TestAdjustInputQueueSize(t *testing.T) {
-	t.Run("zero yields default value (main queue size=0)", func(t *testing.T) {
-		assert.Equal(t, minInputQueueSize, AdjustInputQueueSize(0, 0))
-	})
-	t.Run("zero yields default value (main queue size=10)", func(t *testing.T) {
-		assert.Equal(t, minInputQueueSize, AdjustInputQueueSize(0, 10))
-	})
-	t.Run("can't go below min", func(t *testing.T) {
-		assert.Equal(t, minInputQueueSize, AdjustInputQueueSize(1, 0))
-	})
-	t.Run("can set any value within bounds", func(t *testing.T) {
-		for q, mainQueue := minInputQueueSize+1, 4096; q < int(float64(mainQueue)*maxInputQueueSizeRatio); q += 10 {
-			assert.Equal(t, q, AdjustInputQueueSize(q, mainQueue))
-		}
-	})
-	t.Run("can set any value if no upper bound", func(t *testing.T) {
-		for q := minInputQueueSize + 1; q < math.MaxInt32; q *= 2 {
-			assert.Equal(t, q, AdjustInputQueueSize(q, 0))
-		}
-	})
-	t.Run("can't go above upper bound", func(t *testing.T) {
-		mainQueue := 4096
-		assert.Equal(t, int(float64(mainQueue)*maxInputQueueSizeRatio), AdjustInputQueueSize(mainQueue, mainQueue))
-	})
-}
-
 func TestBatchFreeEntries(t *testing.T) {
 	const queueSize = 10
 	const batchSize = 5
@@ -341,3 +315,130 @@ func TestBatchFreeEntries(t *testing.T) {
 		require.Nilf(t, testQueue.buf[i].event, "Queue index %v: all events should be nil after calling FreeEntries on both batches")
 	}
 }
+
+func TestProducerShutdown(t *testing.T) {
+	// Test that the number of acknowledgment callbacks exactly matches the
+	// number of published events when many goroutines are publishing during
+	// queue shutdown.
+	//
+	// The numbers here (queue size, number of publisher workers, etc.) are
+	// kind of magic since there's no deterministic way to verify this, but they
+	// were chosen so that, when there _was_ a race in the queue shutdown that
+	// could send an extra acknowledgment
+	// (https://github.com/elastic/beats/issues/47246), this test failed about
+	// 90% of the time.
+	const queueSize = 1000
+	const publishWorkers = 50
+	var ackedCount atomic.Int64
+	var publishedCount atomic.Int64
+	testQueue := NewQueue(
+		logp.NewNopLogger(),
+		nil,
+		Settings{
+			Events:        queueSize,
+			MaxGetRequest: queueSize,
+			FlushTimeout:  time.Second},
+		0,
+		nil)
+
+	var wg sync.WaitGroup
+	// Start workers to continuously publish events to the queue
+	publishWorker := func() {
+		defer wg.Done()
+		// Continuously publish events until Publish returns false indicating queue
+		// shutdown.
+		producer := testQueue.Producer(
+			queue.ProducerConfig{
+				ACK: func(count int) { ackedCount.Add(int64(count)) },
+			})
+		for {
+			_, published := producer.Publish(0)
+			if published {
+				publishedCount.Add(1)
+			} else {
+				return
+			}
+		}
+	}
+	for range publishWorkers {
+		wg.Add(1)
+		go publishWorker()
+	}
+	// Start a reader to continuously drain the queue and acknowledge the events
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		// Continuously read and acknowledge events from the queue
+		for {
+			batch, err := testQueue.Get(queueSize)
+			if err == nil {
+				batch.Done()
+			} else {
+				return
+			}
+		}
+	}()
+
+	// Wait for the queue to go through at least one full rotation
+	require.Eventually(
+		t,
+		func() bool { return publishedCount.Load() > queueSize },
+		time.Second,
+		time.Millisecond,
+		"events are not flowing through the queue")
+
+	// Trigger queue shutdown
+	testQueue.Close(false)
+
+	// Wait for queue context to finish
+	select {
+	case <-testQueue.Done():
+	case <-time.After(5 * time.Second):
+		require.Fail(t, "queue never shut down")
+	}
+
+	// Wait for helper routines to finish
+	wg.Wait()
+
+	// Wait for the ack loop to finish processing callbacks
+	testQueue.wg.Wait()
+
+	require.Equal(t, publishedCount.Load(), ackedCount.Load(), "published and acknowledged event counts should match")
+}
+
+func BenchmarkProducerThroughput(b *testing.B) {
+	const queueSize = 10000
+	const publishWorkers = 10
+	testQueue := NewQueue(
+		logp.NewNopLogger(),
+		nil,
+		Settings{
+			Events:        queueSize,
+			MaxGetRequest: queueSize,
+			FlushTimeout:  time.Second},
+		0,
+		nil)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	publishWorker := func() {
+		producer := testQueue.Producer(queue.ProducerConfig{})
+		for ctx.Err() == nil {
+			producer.Publish(0)
+		}
+	}
+	for range publishWorkers {
+		go publishWorker()
+	}
+	for b.Loop() {
+		// With a flush timeout of a second, we can confidently expect we'll get
+		// a full batch each time, so each iteration is measuring the time for the
+		// publish workers to fill the queue.
+		batch, err := testQueue.Get(queueSize)
+		if err != nil {
+			b.Fatal("Fetching queue batch should succeed")
+		}
+		batch.Done()
+	}
+	cancel()
+	testQueue.Close(true)
+}