Deflake pkg/trace/writer by isolating tests from each other

pkg/trace/writer/sender.go

@@ -245,11 +245,14 @@ type sender struct {
 	closed  bool         // closed reports if the loop is stopped
 	statsd  statsd.ClientInterface
 	enabled *atomic.Bool // false on inactive MRF senders. True otherwise
+
+	wg       sync.WaitGroup // tracks the worker goroutines started by newSender
+	stopOnce sync.Once      // guards Stop, so it can safely be called more than once
 }
 
 // newSender returns a new sender based on the given config cfg.
 func newSender(cfg *senderConfig, apiKeyManager *apiKeyManager, statsd statsd.ClientInterface) *sender {
-	s := sender{
+	s := &sender{
 		cfg:           cfg,
 		apiKeyManager: apiKeyManager,
 		queue:         make(chan *payload, cfg.maxQueued),
@@ -258,10 +261,14 @@ func newSender(cfg *senderConfig, apiKeyManager *apiKeyManager, statsd statsd.Cl
 		statsd:        statsd,
 		enabled:       atomic.NewBool(true),
 	}
+	s.wg.Add(cfg.maxConns)
 	for i := 0; i < cfg.maxConns; i++ {
-		go s.loop()
+		go func() {
+			defer s.wg.Done()
+			s.loop()
+		}()
 	}
-	return &s
+	return s
 }
 
 // loop runs the main sender loop.
@@ -281,13 +288,43 @@ func (s *sender) backoff(attempt int) {
 }
 
 // Stop stops the sender. It attempts to wait for all inflight payloads to complete
-// with a timeout of 5 seconds.
+// with a timeout of 5 seconds, and then waits for the worker goroutines to exit.
+// Stop is idempotent.
 func (s *sender) Stop() {
-	s.WaitForInflight()
-	s.mu.Lock()
-	s.closed = true
-	s.mu.Unlock()
-	close(s.queue)
+	s.stopOnce.Do(func() {
+		// Wait once before closing, so that payloads already inflight keep their
+		// full retry budget.
+		s.WaitForInflight()
+		s.mu.Lock()
+		s.closed = true
+		s.mu.Unlock()
+		// Push increments inflight while holding a read lock, so any Push that
+		// observed !closed has already been counted by the time the write lock
+		// above is acquired. Waiting again here therefore cannot miss a payload
+		// that raced with the close, and guarantees the channel is only closed
+		// once no producer can still be sending on it.
+		s.WaitForInflight()
+		close(s.queue)
+		// Wait for the workers to exit so they cannot outlive their sender, but
+		// keep the same 5 second budget the inflight wait uses: a worker asleep
+		// in backoff must not hold up agent shutdown indefinitely.
+		waitTimeout(&s.wg, 5*time.Second)
+	})
+}
+
+// waitTimeout waits for wg, giving up after d. It reports whether wg completed.
+func waitTimeout(wg *sync.WaitGroup, d time.Duration) bool {
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+	select {
+	case <-done:
+		return true
+	case <-time.After(d):
+		return false
+	}
 }
 
 // WaitForInflight blocks until all in progress payloads are sent,
@@ -315,14 +352,18 @@ func (s *sender) Push(p *payload) {
 		s.mu.RUnlock()
 		return
 	}
+	// Count the payload before it becomes visible to a worker: incrementing
+	// after the enqueue lets a concurrent WaitForInflight observe zero while p
+	// is already queued, so callers such as FlushSync and Stop could return
+	// before p was ever sent.
+	s.inflight.Inc()
 	s.mu.RUnlock()
 	select {
 	case s.queue <- p:
 	default:
 		_ = s.statsd.Count("datadog.trace_agent.sender.push_blocked", 1, nil, 1)
 		s.queue <- p
 	}
-	s.inflight.Inc()
 }
 
 // sendPayload sends the payload p to the destination URL.

pkg/trace/writer/sender_test.go

@@ -7,12 +7,15 @@ package writer
 
 import (
 	"bytes"
+	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"os"
+	"runtime"
 	"strconv"
+	"strings"
 	"sync"
 	"testing"
 	"testing/synctest"
@@ -31,7 +34,46 @@ const testAPIKey = "123"
 
 func TestMain(m *testing.M) {
 	log.SetLogger(log.NoopLogger)
-	os.Exit(m.Run())
+	code := m.Run()
+	// Only check for leaks on an otherwise green run: a failing test may well
+	// have skipped its own cleanup, and the leak would be a red herring.
+	if code == 0 {
+		if n, sample := leakedSenderGoroutines(); n > 0 {
+			fmt.Fprintf(os.Stderr,
+				"%d sender goroutine(s) outlived the tests that created them. A leaked sender\n"+
+					"keeps worker goroutines and an HTTP client alive, and its test server keeps a\n"+
+					"listener bound, so both can interfere with later tests in this package. Stop\n"+
+					"every writer and sender you create (see cleanupWriter).\nSample stack:\n%s\n",
+				n, sample)
+			code = 1
+		}
+	}
+	os.Exit(code)
+}
+
+// leakedSenderGoroutines reports how many sender worker goroutines are still
+// running, along with one representative stack. It gives them a short grace
+// period, since Stop closes the queue before the workers observe it.
+func leakedSenderGoroutines() (int, string) {
+	const marker = "trace/writer.(*sender).loop"
+	deadline := time.Now().Add(2 * time.Second)
+	for {
+		var found []string
+		buf := make([]byte, 1<<22)
+		dump := string(buf[:runtime.Stack(buf, true)])
+		for g := range strings.SplitSeq(dump, "\n\n") {
+			if strings.Contains(g, marker) {
+				found = append(found, g)
+			}
+		}
+		if len(found) == 0 {
+			return 0, ""
+		}
+		if time.Now().After(deadline) {
+			return len(found), found[0]
+		}
+		time.Sleep(20 * time.Millisecond)
+	}
 }
 
 func TestMaxConns(t *testing.T) {

pkg/trace/writer/stats_test.go

@@ -39,20 +39,29 @@ const (
 	testEnv      = "testing"
 )
 
-func assertPayload(t *testing.T, testSets []*pb.StatsPayload, payloads []*payload) {
+func assertPayload(t *testing.T, testSets []*pb.StatsPayload, srv *testServer) {
 	t.Helper()
 	expectedHeaders := map[string]string{
 		"X-Datadog-Reported-Languages": strings.Join(info.Languages(), "|"),
 		"Content-Type":                 "application/msgpack",
 		"Content-Encoding":             "gzip",
 		"Dd-Api-Key":                   "123",
 	}
+	received := srv.Received()
+	// A body that is not valid gzip has historically shown up here, and the count
+	// is the first thing that tells apart "the writer under test misbehaved" from
+	// "a request that belongs to another test reached this server". Check it
+	// before decoding, and describe every recorded request when it is wrong.
+	require.Len(t, received, len(testSets),
+		"unexpected number of requests reached the test server; recorded: %s", describeReceived(received))
+	require.Empty(t, srv.ReadErrors(),
+		"the test server failed to read a request body, which means a request was in flight when it closed")
 	var decoded []*pb.StatsPayload
-	for _, p := range payloads {
+	for _, p := range received {
 		var statsPayload pb.StatsPayload
 		r, err := gzip.NewReader(p.body)
-		require.NoError(t, err, "payload body is not valid gzip")
-		require.NoError(t, msgp.Decode(r, &statsPayload))
+		require.NoError(t, err, "payload body is not valid gzip: %s", p.describe())
+		require.NoError(t, msgp.Decode(r, &statsPayload), "payload body is not decodable: %s", p.describe())
 		require.NoError(t, r.Close())
 		for k, v := range expectedHeaders {
 			assert.Equal(t, v, p.headers[k])
@@ -70,7 +79,7 @@ func assertPayload(t *testing.T, testSets []*pb.StatsPayload, payloads []*payloa
 
 func TestStatsWriter(t *testing.T) {
 	t.Run("ok", func(t *testing.T) {
-		sw, srv := testStatsWriter()
+		sw, srv := testStatsWriter(t)
 		go sw.Run()
 
 		testSets := []*pb.StatsPayload{
@@ -106,11 +115,11 @@ func TestStatsWriter(t *testing.T) {
 		sw.Write(testSets[0])
 		sw.Write(testSets[1])
 		sw.Stop()
-		assertPayload(t, testSets, srv.Payloads())
+		assertPayload(t, testSets, srv)
 	})
 
-	t.Run("race", func(_ *testing.T) {
-		sw, _ := testStatsWriter()
+	t.Run("race", func(t *testing.T) {
+		sw, _ := testStatsWriter(t)
 		// Don't start the writer as we're going to call send ourselves to test for a race
 		stopChan := make(chan struct{})
 		wg := sync.WaitGroup{}
@@ -179,7 +188,7 @@ func TestStatsWriter(t *testing.T) {
 
 	t.Run("buildPayloads", func(t *testing.T) {
 		assert := assert.New(t)
-		sw, srv := testStatsWriter()
+		sw, srv := testStatsWriter(t)
 		srv.Close()
 		// This gives us a total of 45 entries. 3 per span, 5
 		// spans per stat bucket. Each buckets have the same
@@ -251,7 +260,7 @@ func TestStatsWriter(t *testing.T) {
 	t.Run("no-split", func(t *testing.T) {
 		assert := assert.New(t)
 
-		sw, srv := testStatsWriter()
+		sw, srv := testStatsWriter(t)
 		srv.Close()
 		// This gives us a total of 45 entries. 3 per span, 5 spans per
 		// stat bucket. Each bucket has the same time window (start:
@@ -278,7 +287,7 @@ func TestStatsWriter(t *testing.T) {
 
 	t.Run("container-tags", func(t *testing.T) {
 		assert := assert.New(t)
-		sw, srv := testStatsWriter()
+		sw, srv := testStatsWriter(t)
 		srv.Close()
 		stats := &pb.StatsPayload{
 			AgentHostname: "agenthost",
@@ -312,7 +321,7 @@ func TestStatsWriter(t *testing.T) {
 }
 
 func TestStatsResetBuffer(t *testing.T) {
-	w, _ := testStatsSyncWriter()
+	w, _ := testStatsSyncWriter(t)
 
 	runtime.GC()
 	var m runtime.MemStats
@@ -342,7 +351,7 @@ func TestStatsSyncWriter(t *testing.T) {
 
 	t.Run("ok", func(t *testing.T) {
 		assert := assert.New(t)
-		sw, srv := testStatsSyncWriter()
+		sw, srv := testStatsSyncWriter(t)
 		go sw.Run()
 		testSets := []*pb.StatsPayload{
 			{
@@ -378,11 +387,11 @@ func TestStatsSyncWriter(t *testing.T) {
 		assert.Nil(err)
 		sw.Stop()
 		srv.Close()
-		assertPayload(t, testSets, srv.Payloads())
+		assertPayload(t, testSets, srv)
 	})
 
 	t.Run("stop", func(t *testing.T) {
-		sw, srv := testStatsSyncWriter()
+		sw, srv := testStatsSyncWriter(t)
 		go sw.Run()
 
 		testSets := []*pb.StatsPayload{
@@ -413,13 +422,16 @@ func TestStatsSyncWriter(t *testing.T) {
 		sw.Write(testSets[1])
 		sw.Stop()
 		srv.Close()
-		assertPayload(t, testSets, srv.Payloads())
+		// Unlike the sibling subtests, nothing is sent here: in sync mode Write
+		// only buffers and Stop does not flush that buffer.
+		require.Empty(t, srv.Received(), "Stop is not expected to flush in sync mode")
+		require.Empty(t, srv.ReadErrors())
 	})
 }
 
 func TestStatsWriterUpdateAPIKey(t *testing.T) {
 	assert := assert.New(t)
-	sw, srv := testStatsSyncWriter()
+	sw, srv := testStatsSyncWriter(t)
 	go sw.Run()
 	defer sw.Stop()
 
@@ -444,7 +456,7 @@ func TestStatsWriterInfo(t *testing.T) {
 	assert := assert.New(t)
 	// statsLastMinute updates depend on StatsWriter internal ticker, but are also triggered
 	// with sync mode. We will use sync writer to test the stats info updates.
-	sw, srv := testStatsSyncWriter()
+	sw, srv := testStatsSyncWriter(t)
 	go sw.Run()
 
 	time.Sleep(200 * time.Millisecond) // allow stats to be initialized
@@ -484,7 +496,7 @@ func TestStatsWriterInfo(t *testing.T) {
 	err := sw.FlushSync()
 	assert.Nil(err)
 
-	assertPayload(t, testSets, srv.Payloads())
+	assertPayload(t, testSets, srv)
 
 	assert.NotEmpty(sw.statsLastMinute.Bytes.Load())
 	assert.Empty(sw.statsLastMinute.Errors.Load())
@@ -568,7 +580,7 @@ func TestContainerTagsBufferManyTracerPayload(t *testing.T) {
 				},
 			}
 
-			sw, srv := testStatsWriterWithBuffer(mockBuf)
+			sw, srv := testStatsWriterWithBuffer(t, mockBuf)
 			go sw.Run()
 			defer sw.Stop()
 
@@ -619,30 +631,47 @@ func (m *mockContainerTagsBuffer) AsyncEnrichment(containerID string, cb func([]
 	return m.pending
 }
 
-func testStatsWriterWithBuffer(buffer containertagsbuffer.ContainerTagsBuffer) (*DatadogStatsWriter, *testServer) {
-	writer, srv := testStatsWriter()
+func testStatsWriterWithBuffer(t *testing.T, buffer containertagsbuffer.ContainerTagsBuffer) (*DatadogStatsWriter, *testServer) {
+	writer, srv := testStatsWriter(t)
 	writer.containerTagsBuffer = buffer
 	return writer, srv
 }
 
-func testStatsWriter() (*DatadogStatsWriter, *testServer) {
+func testStatsWriter(t *testing.T) (*DatadogStatsWriter, *testServer) {
 	srv := newTestServer()
 	cfg := &config.AgentConfig{
 		Endpoints:     []*config.Endpoint{{Host: srv.URL, APIKey: "123"}},
 		StatsWriter:   &config.WriterConfig{ConnectionLimit: 20, QueueSize: 20},
 		ContainerTags: func(_ string) ([]string, error) { return nil, nil },
 	}
-	return NewStatsWriter(cfg, telemetry.NewNoopCollector(), &statsd.NoOpClient{}, &timing.NoopReporter{}, &containertagsbuffer.NoOpTagsBuffer{}), srv
+	w := NewStatsWriter(cfg, telemetry.NewNoopCollector(), &statsd.NoOpClient{}, &timing.NoopReporter{}, &containertagsbuffer.NoOpTagsBuffer{})
+	cleanupWriter(t, w, srv)
+	return w, srv
 }
 
-func testStatsSyncWriter() (*DatadogStatsWriter, *testServer) {
+func testStatsSyncWriter(t *testing.T) (*DatadogStatsWriter, *testServer) {
 	srv := newTestServer()
 	cfg := &config.AgentConfig{
 		Endpoints:           []*config.Endpoint{{Host: srv.URL, APIKey: "123"}},
 		StatsWriter:         &config.WriterConfig{ConnectionLimit: 20, QueueSize: 20},
 		SynchronousFlushing: true,
 	}
-	return NewStatsWriter(cfg, telemetry.NewNoopCollector(), &statsd.NoOpClient{}, &timing.NoopReporter{}, &containertagsbuffer.NoOpTagsBuffer{}), srv
+	w := NewStatsWriter(cfg, telemetry.NewNoopCollector(), &statsd.NoOpClient{}, &timing.NoopReporter{}, &containertagsbuffer.NoOpTagsBuffer{})
+	cleanupWriter(t, w, srv)
+	return w, srv
+}
+
+// cleanupWriter guarantees that a writer's senders and its test server are torn
+// down when the test ends, in that order. NewStatsWriter starts ConnectionLimit
+// sender goroutines immediately, so a test that forgets to stop its writer used
+// to leave workers, an HTTP client and a listener alive for every later test in
+// the package to interact with. Both sender.Stop and testServer.Close are
+// idempotent, so a test may still stop or close explicitly.
+func cleanupWriter(t *testing.T, w *DatadogStatsWriter, srv *testServer) {
+	t.Cleanup(func() {
+		stopSenders(w.senders)
+		srv.Close()
+	})
 }
 
 type key struct {