chore: Use in-memory network connection for metastore client tests (#4235)

Author: callpraths

pkg/experiment/metastore/client/client.go

@@ -31,6 +31,7 @@ type Client struct {
 	stopped          bool
 	logger           log.Logger
 	grpcClientConfig grpcclient.Config
+	dialOpts         []grpc.DialOption
 }
 
 type client struct {
@@ -53,14 +54,16 @@ type instance interface {
 	raftnodepb.RaftNodeServiceClient
 }
 
-func New(logger log.Logger, grpcClientConfig grpcclient.Config, d discovery.Discovery) *Client {
+func New(logger log.Logger, grpcClientConfig grpcclient.Config, d discovery.Discovery, dialOpts ...grpc.DialOption) *Client {
 	var c Client
 	logger = log.With(logger, "component", "metastore-client")
 	c.service = services.NewIdleService(c.starting, c.stopping)
 	c.logger = logger
 	c.grpcClientConfig = grpcClientConfig
 	c.servers = make(map[raft.ServerID]*client)
 	c.discovery = d
+	c.dialOpts = dialOpts
+
 	c.discovery.Subscribe(discovery.UpdateFunc(func(servers []discovery.Server) {
 		c.updateServers(servers)
 	}))
@@ -117,7 +120,7 @@ func (c *Client) updateServers(servers []discovery.Server) {
 				continue
 			}
 		}
-		cl, err := newClient(s[0], c.grpcClientConfig, c.logger)
+		cl, err := newClient(s[0], c.grpcClientConfig, c.dialOpts...)
 		if err != nil {
 			level.Error(c.logger).Log("msg", "failed to create client", "err", err)
 			continue
@@ -139,12 +142,12 @@ func (c *Client) updateServers(servers []discovery.Server) {
 	c.servers = newServers
 }
 
-func newClient(s discovery.Server, config grpcclient.Config, logger log.Logger) (*client, error) {
+func newClient(s discovery.Server, config grpcclient.Config, dialOpts ...grpc.DialOption) (*client, error) {
 	address := s.Raft.Address
 	if s.ResolvedAddress != "" {
 		address = raft.ServerAddress(s.ResolvedAddress)
 	}
-	conn, err := dial(string(address), config, logger)
+	conn, err := dial(string(address), config, dialOpts...)
 	if err != nil {
 		return nil, err
 	}
@@ -159,13 +162,14 @@ func newClient(s discovery.Server, config grpcclient.Config, logger log.Logger)
 	}, nil
 }
 
-func dial(address string, grpcClientConfig grpcclient.Config, _ log.Logger) (*grpc.ClientConn, error) {
+func dial(address string, grpcClientConfig grpcclient.Config, dialOpts ...grpc.DialOption) (*grpc.ClientConn, error) {
 	options, err := grpcClientConfig.DialOption(nil, nil)
 	if err != nil {
 		return nil, err
 	}
 	// TODO: https://github.com/grpc/grpc-proto/blob/master/grpc/service_config/service_config.proto
 	options = append(options, grpc.WithDefaultServiceConfig(grpcServiceConfig))
+	options = append(options, dialOpts...)
 	return grpc.Dial(address, options...)
 }
 

pkg/experiment/metastore/client/client_test.go

@@ -7,7 +7,6 @@ import (
 
 	"github.com/grafana/dskit/flagext"
 	"github.com/grafana/dskit/grpcclient"
-	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 
@@ -16,20 +15,16 @@ import (
 	"github.com/grafana/pyroscope/pkg/test/mocks/mockdiscovery"
 )
 
-const nServers = 3
-
 func TestUnavailable(t *testing.T) {
 	d := mockdiscovery.NewMockDiscovery(t)
 	d.On("Subscribe", mock.Anything).Return()
 	l := test.NewTestingLogger(t)
 	c := New(l, grpcclient.Config{}, d)
-	ports, err := test.GetFreePorts(nServers)
-	assert.NoError(t, err)
 
 	d.On("Rediscover").Run(func(args mock.Arguments) {
 	}).Return()
 
-	c.updateServers(createServers(ports))
+	c.updateServers(createServers([]int{30030, 30031, 30032}))
 	res, err := c.AddBlock(context.Background(), &metastorev1.AddBlockRequest{})
 	require.Error(t, err)
 	require.Nil(t, res)
@@ -73,32 +68,30 @@ func testRediscoverWrongLeader(t *testing.T, f func(c *Client)) {
 	l := test.NewTestingLogger(t)
 	config := &grpcclient.Config{}
 	flagext.DefaultValues(config)
-	c := New(l, *config, d)
-	ports, err := test.GetFreePorts(nServers * 2)
-	assert.NoError(t, err)
 
-	p1 := ports[:nServers]
-	p2 := ports[nServers:]
-	m := sync.Mutex{}
-	var servers *mockServers
-	defer servers.Close()
+	dServers1 := createServers([]int{30031, 30032, 30033})
+
+	dServers2 := createServers([]int{40031, 40032, 40033})
+	mockServers2, dialOpts := createMockServers(t, l, dServers2)
+	defer mockServers2.Close()
 
+	c := New(l, *config, d, dialOpts...)
+	m := sync.Mutex{}
 	verify := func() {}
+	initWrongLeaderCalled := false
 	d.On("Rediscover", mock.Anything).Run(func(args mock.Arguments) {
 		m.Lock()
 		defer m.Unlock()
-		if servers == nil {
-			srvInfo := createServers(p2)
-			servers = createMockServers(t, l, p2)
-			verify = servers.InitWrongLeader()
-
+		if !initWrongLeaderCalled {
+			initWrongLeaderCalled = true
+			verify = mockServers2.InitWrongLeader()
 			// call updateServers twice
-			c.updateServers(srvInfo)
-			c.updateServers(srvInfo)
+			c.updateServers(dServers2)
+			c.updateServers(dServers2)
 		}
 	}).Return()
 
-	c.updateServers(createServers(p1))
+	c.updateServers(dServers1)
 	f(c)
 	verify()
 }

pkg/experiment/metastore/client/server_mock_test.go

@@ -14,6 +14,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"google.golang.org/grpc/test/bufconn"
 
 	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
 	"github.com/grafana/pyroscope/pkg/experiment/metastore/discovery"
@@ -95,13 +96,13 @@ func (m *mockServer) PromoteToLeader(ctx context.Context, request *raftnodepb.Pr
 
 func createServers(ports []int) []discovery.Server {
 	var servers []discovery.Server
-	for i := 0; i < nServers; i++ {
+	for i, p := range ports {
 		servers = append(servers, discovery.Server{
 			Raft: raft.Server{
 				ID:      testServerId(i),
 				Address: raft.ServerAddress(fmt.Sprintf("server-%d", i)),
 			},
-			ResolvedAddress: fmt.Sprintf("127.0.0.1:%d", ports[i]),
+			ResolvedAddress: fmt.Sprintf("127.0.0.1:%d", p),
 		})
 	}
 	return servers
@@ -139,6 +140,7 @@ func (m *mockServers) InitWrongLeader() func() {
 	s := new(wrongLeaderState)
 	s.leaderIndex = -1
 
+	nServers := len(m.servers)
 	for _, srv := range m.servers {
 		srv := srv
 		errf := func() error {
@@ -196,30 +198,38 @@ func errOrT[T any](t *T, f func() error) (*T, error) {
 	return t, nil
 }
 
-func createMockServers(t *testing.T, l log.Logger, ports []int) *mockServers {
+// Returns the grpc.DialOptions needed for a client connection to the created mock servers.
+func createMockServers(t *testing.T, l log.Logger, dServers []discovery.Server) (*mockServers, []grpc.DialOption) {
 	var servers []*mockServer
-	for idx, port := range ports {
+	listeners := make(map[string]*bufconn.Listener)
+	for idx, dserv := range dServers {
 		s := newMockServer(t)
 		s.index = idx
 		s.id = testServerId(idx)
-		s.address = fmt.Sprintf(":%d", port)
-		lis, err := net.Listen("tcp", s.address)
-		if err != nil {
-			assert.NoError(t, err)
-		}
+		s.address = dserv.ResolvedAddress
+		lis := bufconn.Listen(256 << 10)
+		listeners[s.address] = lis
 		go func() {
 			if err := s.srv.Serve(lis); err != nil {
 				assert.NoError(t, err)
 			}
 		}()
 		servers = append(servers, s)
 	}
-	return &mockServers{
+
+	ms := &mockServers{
 		servers: servers,
 		t:       t,
 		l:       l,
 	}
-
+	dialer := func(_ context.Context, address string) (net.Conn, error) {
+		el := listeners[address]
+		if el != nil {
+			return el.Dial()
+		}
+		return net.Dial("tcp", address)
+	}
+	return ms, []grpc.DialOption{grpc.WithContextDialer(dialer)}
 }
 
 func newMockServer(t *testing.T) *mockServer {