fix: De-flake several flaky unit and integration tests

integration/labelaccess_test.go

@@ -442,27 +442,45 @@ func (tc *testQueryAndLabelResults) metricsRangeQuery(t *testing.T, qc *client.C
 		require.NoError(t, err)
 		require.Equal(t, "success", output.Status)
 		require.Equal(t, "matrix", output.Data.ResultType)
-		require.EqualValues(t, len(tc.lines), sumMatrixValues(t, output.Data.Matrix))
+		require.EqualValues(t, len(tc.lines), maxStepValue(t, output.Data.Matrix))
 	})
 }
 
-func sumMatrixValues(t *testing.T, matrix []client.MatrixValues) float64 {
+// maxStepValue returns the largest per-step sum across all series in the matrix.
+//
+// The query is count_over_time(...[2h]) evaluated as a range query. Because the
+// lookback window (2h) is much larger than the range step, the recently pushed
+// lines fall inside several consecutive steps' windows, so each of those steps
+// reports the full count. Summing every sample would therefore multiply the true
+// count by the (timing-dependent) number of overlapping steps. The step whose
+// window fully covers the pushed data reports exactly the total we want, so take
+// the largest per-step sum.
+func maxStepValue(t *testing.T, matrix []client.MatrixValues) float64 {
 	t.Helper()
 
-	var total float64
+	perStep := map[float64]float64{}
 	for _, row := range matrix {
 		for _, sample := range row.Values {
 			require.Len(t, sample, 2)
+			ts, ok := sample[0].(float64)
+			require.True(t, ok)
 			value, ok := sample[1].(string)
 			require.True(t, ok)
 
 			parsed, err := strconv.ParseFloat(value, 64)
 			require.NoError(t, err)
-			total += parsed
+			perStep[ts] += parsed
+		}
+	}
+
+	var maxValue float64
+	for _, total := range perStep {
+		if total > maxValue {
+			maxValue = total
 		}
 	}
 
-	return total
+	return maxValue
 }
 
 func normalizeLabelsData(labels map[string][]string) map[string][]string {

pkg/canary/comparator/comparator_test.go

@@ -20,6 +20,7 @@ func TestComparatorEntryReceivedOutOfOrder(t *testing.T) {
 
 	actual := &bytes.Buffer{}
 	c := NewComparator(actual, 1*time.Hour, 1*time.Hour, 1*time.Hour, 15*time.Minute, 4*time.Hour, 4*time.Hour, 0, 1*time.Minute, 0, 1*time.Hour, 3*time.Hour, 30*time.Minute, 0, 1, make(chan time.Time), make(chan time.Time), nil, false)
+	defer c.Stop()
 
 	t1 := time.Now()
 	t2 := t1.Add(1 * time.Second)
@@ -61,6 +62,7 @@ func TestComparatorEntryReceivedNotExpected(t *testing.T) {
 
 	actual := &bytes.Buffer{}
 	c := NewComparator(actual, 1*time.Hour, 1*time.Hour, 1*time.Hour, 15*time.Minute, 4*time.Hour, 4*time.Hour, 0, 1*time.Minute, 0, 1*time.Hour, 3*time.Hour, 30*time.Minute, 0, 1, make(chan time.Time), make(chan time.Time), nil, false)
+	defer c.Stop()
 
 	t1 := time.Now()
 	t2 := t1.Add(1 * time.Second)
@@ -102,6 +104,7 @@ func TestComparatorEntryReceivedDuplicate(t *testing.T) {
 
 	actual := &bytes.Buffer{}
 	c := NewComparator(actual, 1*time.Hour, 1*time.Hour, 1*time.Hour, 15*time.Minute, 4*time.Hour, 4*time.Hour, 0, 1*time.Minute, 0, 1*time.Hour, 3*time.Hour, 30*time.Minute, 0, 1, make(chan time.Time), make(chan time.Time), nil, false)
+	defer c.Stop()
 
 	t1 := time.Unix(0, 0)
 	t2 := t1.Add(1 * time.Second)
@@ -160,6 +163,7 @@ func TestEntryNeverReceived(t *testing.T) {
 	maxWait := 300 * time.Second
 	//We set the prune interval timer to a huge value here so that it never runs, instead we call pruneEntries manually below
 	c := NewComparator(actual, wait, maxWait, 50*time.Hour, 15*time.Minute, 4*time.Hour, 4*time.Hour, 0, 1*time.Minute, 0, 1*time.Hour, 3*time.Hour, 30*time.Minute, 0, 1, make(chan time.Time), make(chan time.Time), mr, false)
+	defer c.Stop()
 
 	c.entrySent(t1)
 	c.entrySent(t2)
@@ -233,6 +237,7 @@ func TestConcurrentConfirmMissing(t *testing.T) {
 	maxWait := 30 * time.Millisecond
 
 	c := NewComparator(output, wait, maxWait, 50*time.Hour, 15*time.Minute, 4*time.Hour, 4*time.Hour, 0, 1*time.Minute, 0, 1*time.Hour, 3*time.Hour, 30*time.Minute, 0, 1, make(chan time.Time), make(chan time.Time), mr, false)
+	defer c.Stop()
 
 	for _, t := range found {
 		tCopy := t
@@ -264,6 +269,7 @@ func TestPruneAckdEntires(t *testing.T) {
 	maxWait := 30 * time.Millisecond
 	//We set the prune interval timer to a huge value here so that it never runs, instead we call pruneEntries manually below
 	c := NewComparator(actual, wait, maxWait, 50*time.Hour, 15*time.Minute, 4*time.Hour, 4*time.Hour, 0, 1*time.Minute, 0, 1*time.Hour, 3*time.Hour, 30*time.Minute, 0, 1, make(chan time.Time), make(chan time.Time), nil, false)
+	defer c.Stop()
 
 	t1 := time.Unix(0, 0)
 	t2 := t1.Add(1 * time.Millisecond)
@@ -321,6 +327,7 @@ func TestSpotCheck(t *testing.T) {
 	spotCheckMax := 20 * time.Millisecond
 	//We set the prune interval timer to a huge value here so that it never runs, instead we call spotCheckEntries manually below
 	c := NewComparator(actual, 1*time.Hour, 1*time.Hour, 50*time.Hour, spotCheck, spotCheckMax, 4*time.Hour, 3*time.Millisecond, 1*time.Minute, 0, 1*time.Hour, 3*time.Hour, 30*time.Minute, 0, 1, make(chan time.Time), make(chan time.Time), mr, false)
+	defer c.Stop()
 
 	// Send all the entries
 	for i := range entries {
@@ -369,6 +376,7 @@ func TestCacheTest(t *testing.T) {
 	cacheTestNow := 2 * time.Second
 
 	c := NewComparator(actual, 1*time.Hour, 1*time.Hour, 50*time.Hour, 0, 0, 4*time.Hour, 0, 10*time.Minute, 0, cacheTestInterval, cacheTestRange, cacheTestNow, 1*time.Hour, 1, make(chan time.Time), make(chan time.Time), mr, false)
+	defer c.Stop()
 	// Force the start time to a known value
 	c.startTime = time.Unix(10, 0)
 
@@ -408,6 +416,7 @@ func TestMetricTest(t *testing.T) {
 	metricTestRange := 30 * time.Second
 	//We set the prune interval timer to a huge value here so that it never runs, instead we call spotCheckEntries manually below
 	c := NewComparator(actual, 1*time.Hour, 1*time.Hour, 50*time.Hour, 0, 0, 4*time.Hour, 0, 10*time.Minute, metricTestRange, 1*time.Hour, 3*time.Hour, 30*time.Minute, writeInterval, 1, make(chan time.Time), make(chan time.Time), mr, false)
+	defer c.Stop()
 	// Force the start time to a known value
 	c.startTime = time.Unix(10, 0)
 

pkg/distributor/rendezvous/partition_ring_watcher_test.go

@@ -2,6 +2,7 @@ package rendezvous
 
 import (
 	"context"
+	"slices"
 	"testing"
 	"time"
 
@@ -102,8 +103,12 @@ func TestPartitionWatcher_PicksUpChanges(t *testing.T) {
 	// Update KV store to a ring with only partition 2.
 	writePartitionRing(t, kvClient, activePartitionRing(2))
 
+	// Wait until the watcher reflects the new ring. Checking identity alone is
+	// racy: ShuffleSharder() can return a fresh object that still holds the
+	// pre-update partitions, so wait for the partitions themselves to change.
 	require.Eventually(t, func() bool {
-		return watcher.ShuffleSharder() != initial
+		ss := watcher.ShuffleSharder()
+		return ss != initial && slices.Equal(ss.partitions, []int32{2})
 	}, 5*time.Second, 10*time.Millisecond, "watcher did not pick up KV change")
 
 	updated := watcher.ShuffleSharder()

pkg/storage/chunk/client/gcp/gcs_object_client.go

@@ -291,7 +291,10 @@ func IsStorageTimeoutErr(err error) bool {
 	if isContextErr(err) {
 		// Go 1.23 changed the type of the error returned by the http client when a timeout occurs
 		// while waiting for headers.  This is a server side timeout.
-		return strings.Contains(err.Error(), "Client.Timeout")
+		// The transport's ResponseHeaderTimeout also wraps a context deadline but is
+		// unambiguously a server-side slowness (never a caller cancellation), so retry it too.
+		return strings.Contains(err.Error(), "Client.Timeout") ||
+			strings.Contains(err.Error(), "timeout awaiting response headers")
 	}
 
 	// connection misconfiguration, or writing on a closed connection

pkg/storage/chunk/client/gcp/gcs_object_client_test.go

@@ -156,15 +156,16 @@ func TestUpstreamRetryableErrs(t *testing.T) {
 func TestTCPErrs(t *testing.T) {
 
 	tests := []struct {
-		name           string
-		responseSleep  time.Duration
-		connectSleep   time.Duration
-		clientTimeout  time.Duration
-		serverTimeout  time.Duration
-		connectTimeout time.Duration
-		closeOnNew     bool
-		closeOnActive  bool
-		retryable      bool
+		name                  string
+		responseSleep         time.Duration
+		connectSleep          time.Duration
+		clientTimeout         time.Duration
+		serverTimeout         time.Duration
+		connectTimeout        time.Duration
+		responseHeaderTimeout time.Duration
+		closeOnNew            bool
+		closeOnActive         bool
+		retryable             bool
 	}{
 		{
 			name:          "request took longer than client timeout, not retryable",
@@ -179,13 +180,20 @@ func TestTCPErrs(t *testing.T) {
 			retryable:     false,
 		},
 		{
-			// there are retryable because it's a server-side timeout
-			name:           "transport connect timeout exceeded, retryable",
-			connectSleep:   time.Millisecond * 40,
-			connectTimeout: time.Millisecond * 20,
-			// even though the client timeout is set, the connect timeout will be hit first
-			clientTimeout: time.Millisecond * 100,
-			retryable:     true,
+			// A transport timeout awaiting response headers is retryable: it's
+			// a server-side slowness, surfaced as a net timeout. We use the
+			// transport's ResponseHeaderTimeout rather than http.Client.Timeout
+			// because Go surfaces the latter inconsistently under load (as a
+			// "Client.Timeout" error or a bare "context deadline exceeded"),
+			// whereas ResponseHeaderTimeout is a deterministic net timeout.
+			// responseSleep only needs to exceed responseHeaderTimeout; it is
+			// interruptible so it doesn't slow teardown, and the client-context
+			// timeout is set high so it never wins the race.
+			name:                  "transport response header timeout exceeded, retryable",
+			responseSleep:         5 * time.Second,
+			responseHeaderTimeout: 100 * time.Millisecond,
+			clientTimeout:         30 * time.Second,
+			retryable:             true,
 		},
 		{
 			name:          "connection is closed server-side before being established",
@@ -211,6 +219,9 @@ func TestTCPErrs(t *testing.T) {
 
 			client := http.DefaultClient
 			transport := http.DefaultTransport.(*http.Transport).Clone()
+			if tc.responseHeaderTimeout != 0 {
+				transport.ResponseHeaderTimeout = tc.responseHeaderTimeout
+			}
 			client.Transport = transport
 			client.Timeout = tc.connectTimeout
 
@@ -245,14 +256,24 @@ func fakeHTTPRespondingServer(t *testing.T, code int) *httptest.Server {
 }
 
 func fakeSleepingServer(t *testing.T, responseSleep, connectSleep time.Duration, closeOnNew, closeOnActive bool) *httptest.Server {
+	// done unblocks the sleeps below at teardown, so a sleep set longer than the
+	// client's timeout (to reliably win the timeout race under load) does not
+	// stall server.Close().
+	done := make(chan struct{})
+	sleep := func(d time.Duration) {
+		select {
+		case <-time.After(d):
+		case <-done:
+		}
+	}
 	server := httptest.NewUnstartedServer(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {
 		// sleep on response to mimic server overload
-		time.Sleep(responseSleep)
+		sleep(responseSleep)
 	}))
 	server.Config.ConnState = func(conn net.Conn, state http.ConnState) {
 		// sleep on initial connection attempt to mimic server non-responsiveness
 		if state == http.StateNew {
-			time.Sleep(connectSleep)
+			sleep(connectSleep)
 			if closeOnNew {
 				require.NoError(t, conn.Close())
 			}
@@ -262,7 +283,10 @@ func fakeSleepingServer(t *testing.T, responseSleep, connectSleep time.Duration,
 			require.NoError(t, conn.Close())
 		}
 	}
-	t.Cleanup(server.Close)
+	t.Cleanup(func() {
+		close(done)
+		server.Close()
+	})
 	server.Start()
 	return server
 }

pkg/storage/stores/shipper/indexshipper/uploads/table_manager_test.go

@@ -33,6 +33,12 @@ func buildTestTableManager(t *testing.T, testDir string) (TableManager, stopFunc
 
 	cfg := Config{
 		UploadInterval: time.Hour,
+		// The manager's background loop runs an upload on startup, which races
+		// with tests that add indexes and then inspect the in-memory set. With
+		// a zero retain period, Cleanup drops indexes as soon as they're
+		// uploaded, so ForEach could miss just-added indexes. Retain them for
+		// the duration of the test.
+		DBRetainPeriod: time.Hour,
 	}
 	tm, err := NewTableManager(cfg, storageClient, nil, log.NewNopLogger())
 	require.NoError(t, err)

pkg/util/cfg/dynamic_test.go

@@ -21,10 +21,11 @@ server:
 		data := NewDynamicConfig(mockApplyDynamicConfig)
 		fs := flag.NewFlagSet(t.Name(), flag.PanicOnError)
 
-		file, err := os.CreateTemp("", "config.yaml")
+		file, err := os.CreateTemp(t.TempDir(), "config.yaml")
 		require.NoError(t, err)
 		_, err = file.WriteString(config)
 		require.NoError(t, err)
+		require.NoError(t, file.Close())
 
 		configFileArgs := []string{"-config.file", file.Name()}
 		if args == nil {