Make test deterministic.

Author: tcard

pkg/ingester/ingester_test.go

@@ -64,6 +64,7 @@ import (
 	asmodel "github.com/grafana/mimir/pkg/ingester/activeseries/model"
 	"github.com/grafana/mimir/pkg/ingester/client"
 	"github.com/grafana/mimir/pkg/mimirpb"
+	mimirpbtest "github.com/grafana/mimir/pkg/mimirpb/testutil"
 	"github.com/grafana/mimir/pkg/storage/chunk"
 	"github.com/grafana/mimir/pkg/storage/sharding"
 	mimir_tsdb "github.com/grafana/mimir/pkg/storage/tsdb"
@@ -9156,16 +9157,19 @@ func TestIngesterNoRW2MetadataRefLeaks(t *testing.T) {
 	err = mimirpb.Unmarshal(buf, &wr)
 	require.NoError(t, err)
 
-	nextLeak := mimirpb.NextRefLeakChannel(t.Context())
+	nextLeak := mimirpbtest.NextRefLeakCheck(t.Context(), wr.Buffer().ReadOnlyData())
 
 	ctx := user.InjectOrgID(context.Background(), userID)
 	_, err = ing.Push(ctx, &wr.WriteRequest)
 	require.NoError(t, err)
 
 	select {
-	case addr := <-nextLeak:
-		require.Fail(t, "reference leak detected", "addr: %x", addr)
-	case <-time.After(10 * time.Millisecond):
+	case leaked := <-nextLeak:
+		if leaked {
+			require.Fail(t, "reference leak detected", "addr: %p", buf)
+		}
+	case <-time.After(1 * time.Second):
+		require.Fail(t, "expected reference leak check", "addr: %p", buf)
 	}
 }
 

pkg/mimirpb/custom.go

@@ -4,7 +4,6 @@ package mimirpb
 
 import (
 	"bytes"
-	"context"
 	"fmt"
 	"math"
 	"runtime"
@@ -13,6 +12,7 @@ import (
 	"weak"
 
 	gogoproto "github.com/gogo/protobuf/proto"
+	"github.com/grafana/mimir/pkg/mimirpb/internal"
 	"github.com/prometheus/prometheus/model/histogram"
 	"go.uber.org/atomic"
 	"google.golang.org/grpc/encoding"
@@ -215,6 +215,10 @@ type BufferHolder struct {
 	buffer mem.Buffer
 }
 
+func (m *BufferHolder) Buffer() mem.Buffer {
+	return m.buffer
+}
+
 func (m *BufferHolder) SetBuffer(buf mem.Buffer) {
 	m.buffer = buf
 }
@@ -252,46 +256,29 @@ func (b *instrumentLeaksBuf) Free() {
 		}
 		// Remove our ref; no (strong) refs should remain (except in the stack).
 		b.Buffer = nil
-		weakRef := weak.Make(unsafe.SliceData(buf))
+		ptr := unsafe.SliceData(buf)
+		addr := uintptr(unsafe.Pointer(ptr))
+		weakRef := weak.Make(ptr)
 
 		// Check in a separate goroutine, because this stack still has locals
 		// pointing to the buffer.
 		go func() {
 			runtime.GC()
 			runtime.GC()
 
-			if v := weakRef.Value(); v != nil {
-				select {
-				case leaks <- v:
-				default:
-					panic(fmt.Sprintf("reference leak for object %p", v))
+			leaked := weakRef.Value() != nil
+
+			select {
+			case internal.RefLeakChecks <- internal.RefLeakCheck{Addr: addr, Leaked: leaked}:
+			default:
+				if leaked {
+					panic(fmt.Sprintf("reference leak for object 0x%x", addr))
 				}
 			}
 		}()
 	}
 }
 
-var leaks = make(chan *byte)
-
-// NextRefLeakChannel returns a channel where the next detected reference leak's
-// address will be sent, instead of the default behavior of panicking.
-//
-// Canceling the context immediately closes the channel and restores the default
-// behavior.
-//
-// Intended for use in tests.
-func NextRefLeakChannel(ctx context.Context) <-chan uintptr {
-	ch := make(chan uintptr, 1)
-	go func() {
-		defer close(ch)
-		select {
-		case ch <- uintptr(unsafe.Pointer(<-leaks)):
-		case <-ctx.Done():
-		}
-	}()
-	return ch
-}
-
 // MinTimestamp returns the minimum timestamp (milliseconds) among all series
 // in the WriteRequest. Returns math.MaxInt64 if the request is empty.
 func (m *WriteRequest) MinTimestamp() int64 {

pkg/mimirpb/custom_test.go

@@ -8,7 +8,6 @@ package mimirpb
 import (
 	"testing"
 	"time"
-	"unsafe"
 
 	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/stretchr/testify/assert"
@@ -17,6 +16,7 @@ import (
 	"google.golang.org/grpc/encoding/proto"
 	"google.golang.org/grpc/mem"
 
+	"github.com/grafana/mimir/pkg/mimirpb/internal"
 	"github.com/grafana/mimir/pkg/util/test"
 )
 
@@ -208,7 +208,7 @@ func TestCodecV2_Unmarshal(t *testing.T) {
 
 	require.True(t, origReq.Equal(req))
 
-	require.NotNil(t, req.buffer)
+	require.NotNil(t, req.Buffer)
 	req.FreeBuffer()
 }
 
@@ -258,44 +258,36 @@ func TestInstrumentRefLeaks(t *testing.T) {
 	buf, err := src.Marshal()
 	require.NoError(t, err)
 
-	var req WriteRequest
-	err = Unmarshal(buf, &req)
-	require.NoError(t, err)
-
-	bufAddr := uintptr(unsafe.Pointer(unsafe.SliceData(req.buffer.ReadOnlyData())))
-	// Label names are UnsafeMutableStrings pointing to buf. They shouldn't outlive
-	// the call to req.FreeBuffer.
-	leakingLabelName := req.Timeseries[0].Labels[0].Name
-
-	leaks := NextRefLeakChannel(t.Context())
+	var leaks <-chan bool
+	var leakingLabelName UnsafeMutableString
+	func() {
+		var req WriteRequest
+		err = Unmarshal(buf, &req)
+		require.NoError(t, err)
 
-	recvLeak := func() (uintptr, bool) {
-		select {
-		case detectedAddr := <-leaks:
-			return detectedAddr, true
-		case <-time.After(10 * time.Millisecond):
-			return 0, false
-		}
-	}
+		// Label names are UnsafeMutableStrings pointing to buf. They shouldn't outlive
+		// the call to req.FreeBuffer.
+		leakingLabelName = req.Timeseries[0].Labels[0].Name
 
-	req.FreeBuffer() // leakingLabelName becomes a leak here
+		leaks = internal.NextRefLeakCheck(t.Context(), req.Buffer.ReadOnlyData())
+		req.FreeBuffer() // leakingLabelName becomes a leak here
+	}()
 
 	// Expect to receive a leak detection.
-	detectedAddr, ok := recvLeak()
-	require.True(t, ok, "expected a reference leak")
-	require.Equal(t, bufAddr, detectedAddr)
+	require.Eventually(t, func() bool { return <-leaks }, 10*time.Millisecond, 1*time.Second, "expected a reference leak")
 	// Expect the label name contents to have been replaced with a taint word.
 	// Keep this check last, because we need to extend the lifespan of leakingLabelName
 	// to avoid Go optimizing the leak away.
 	require.Contains(t, leakingLabelName, "KAEL")
 
 	// Now let's check a non-leak doesn't get falsely detected.
-	dataNoLeak, err := src.Marshal()
-	require.NoError(t, err)
-	var reqNoLeak WriteRequest
-	err = Unmarshal(dataNoLeak, &reqNoLeak)
-	require.NoError(t, err)
-	leaks = NextRefLeakChannel(t.Context())
-	detectedAddr, ok = recvLeak()
-	require.False(t, ok, "unexpected a reference leak %x", detectedAddr)
+	func() {
+		var reqNoLeak WriteRequest
+		err = Unmarshal(buf, &reqNoLeak)
+		require.NoError(t, err)
+
+		leaks = internal.NextRefLeakCheck(t.Context(), reqNoLeak.Buffer.ReadOnlyData())
+		reqNoLeak.FreeBuffer()
+	}()
+	require.Eventually(t, func() bool { return !<-leaks }, 10*time.Millisecond, 1*time.Second, "expected no reference leaks")
 }

pkg/mimirpb/internal/refleaks.go

@@ -0,0 +1,48 @@
+package internal
+
+import (
+	"context"
+	"reflect"
+	"testing"
+)
+
+type RefLeakCheck struct {
+	Addr   uintptr
+	Leaked bool
+}
+
+var RefLeakChecks chan RefLeakCheck
+
+func init() {
+	if testing.Testing() {
+		RefLeakChecks = make(chan RefLeakCheck)
+	}
+}
+
+// NextRefLeakCheck is documented in testutil.NextRefLeakCheck.
+func NextRefLeakCheck(ctx context.Context, object any) <-chan bool {
+	if !testing.Testing() {
+		panic("NextRefLeakCheck can only be used in tests")
+	}
+	var addr uintptr
+	switch o := reflect.ValueOf(object); o.Type().Kind() {
+	case reflect.Pointer, reflect.Slice, reflect.String:
+		addr = uintptr(o.UnsafePointer())
+	}
+
+	ch := make(chan bool, 1)
+	go func() {
+		defer close(ch)
+		for {
+			select {
+			case leak := <-RefLeakChecks:
+				if leak.Addr == addr {
+					ch <- leak.Leaked
+					return
+				}
+			case <-ctx.Done():
+			}
+		}
+	}()
+	return ch
+}

pkg/mimirpb/testutil/refleaks.go

@@ -0,0 +1,22 @@
+package testutil
+
+import (
+	"context"
+
+	"github.com/grafana/mimir/pkg/mimirpb/internal"
+)
+
+// NextRefLeakChannel returns a channel where the next check for reference leaks
+// to the given object will be reported, whether a leak is detected
+// or not. It replaces the default behavior of panicking when a leak is detected.
+//
+// The object must be identified with a reference value: a pointer, a slice, or
+// a string.
+//
+// Canceling the context immediately closes the channel and restores the default
+// behavior.
+//
+// Only works when [testing.Testing].
+func NextRefLeakCheck(ctx context.Context, object any) <-chan bool {
+	return internal.NextRefLeakCheck(ctx, object)
+}