chore: Use in-memory connection for scheduler tests (#4233)

Author: callpraths

pkg/scheduler/scheduler.go

@@ -97,6 +97,10 @@ type Config struct {
 	QuerierForgetDelay      time.Duration             `yaml:"querier_forget_delay" category:"experimental"`
 	GRPCClientConfig        grpcclient.Config         `yaml:"grpc_client_config" doc:"description=This configures the gRPC client used to report errors back to the query-frontend."`
 	ServiceDiscovery        schedulerdiscovery.Config `yaml:",inline"`
+
+	// Dial options used to initiate outgoing gRPC connections.
+	// Intended to be used by tests to use in-memory network connections.
+	DialOpts []grpc.DialOption `yaml:"-"`
 }
 
 func (cfg *Config) RegisterFlags(f *flag.FlagSet, logger log.Logger) {
@@ -536,6 +540,7 @@ func (s *Scheduler) forwardErrorToFrontend(ctx context.Context, req *schedulerRe
 		return
 	}
 
+	opts = append(opts, s.cfg.DialOpts...)
 	conn, err := grpc.DialContext(ctx, req.frontendAddress, opts...)
 	if err != nil {
 		level.Warn(s.log).Log("msg", "failed to create gRPC connection to frontend to report error", "frontend", req.frontendAddress, "err", err, "requestErr", requestErr)

pkg/scheduler/scheduler_test.go

@@ -12,7 +12,6 @@ import (
 	"net"
 	"net/http"
 	"net/http/httptest"
-	"net/url"
 	"strings"
 	"sync"
 	"testing"
@@ -33,6 +32,7 @@ import (
 	"golang.org/x/net/http2/h2c"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/test/bufconn"
 
 	"github.com/grafana/pyroscope/pkg/frontend/frontendpb"
 	"github.com/grafana/pyroscope/pkg/scheduler/schedulerpb"
@@ -44,30 +44,50 @@ import (
 
 const testMaxOutstandingPerTenant = 5
 
-func setupScheduler(t *testing.T, reg prometheus.Registerer, opts ...connect.HandlerOption) (*Scheduler, schedulerpb.SchedulerForFrontendClient, schedulerpb.SchedulerForQuerierClient) {
-	cfg := Config{}
+type schedulerArgs struct {
+	reg         prometheus.Registerer
+	handlerOpts []connect.HandlerOption
+	dialOpts    []grpc.DialOption
+}
+
+func setupScheduler(t *testing.T, args schedulerArgs) (*Scheduler, schedulerpb.SchedulerForFrontendClient, schedulerpb.SchedulerForQuerierClient) {
+	cfg := Config{
+		DialOpts: args.dialOpts,
+	}
 	flagext.DefaultValues(&cfg)
 	cfg.MaxOutstandingPerTenant = testMaxOutstandingPerTenant
 
-	s, err := NewScheduler(cfg, &limits{queriers: 2}, log.NewNopLogger(), reg)
+	s, err := NewScheduler(cfg, &limits{queriers: 2}, log.NewNopLogger(), args.reg)
+
 	require.NoError(t, err)
 
 	server := httptest.NewUnstartedServer(nil)
 	mux := mux.NewRouter()
 	server.Config.Handler = h2c.NewHandler(mux, &http2.Server{})
+
+	// Use an in-memory network connection to avoid test flake from network access.
+	listener := bufconn.Listen(256 << 10)
+	server.Listener = listener
+
 	server.Start()
-	u, err := url.Parse(server.URL)
 	require.NoError(t, err)
-	schedulerpbconnect.RegisterSchedulerForFrontendHandler(mux, s, opts...)
-	schedulerpbconnect.RegisterSchedulerForQuerierHandler(mux, s, opts...)
+	schedulerpbconnect.RegisterSchedulerForFrontendHandler(mux, s, args.handlerOpts...)
+	schedulerpbconnect.RegisterSchedulerForQuerierHandler(mux, s, args.handlerOpts...)
 
 	require.NoError(t, services.StartAndAwaitRunning(context.Background(), s))
 	t.Cleanup(func() {
 		_ = services.StopAndAwaitTerminated(context.Background(), s)
 		server.Close()
 	})
 
-	c, err := grpc.Dial(u.Hostname()+":"+u.Port(), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	c, err := grpc.NewClient(
+		// Target address is irrelevant as we're using an in-memory connection.
+		// We simply need the DNS resolution to succeed.
+		"localhost:3030",
+		grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return listener.Dial() }),
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+	)
+
 	require.NoError(t, err)
 
 	t.Cleanup(func() {
@@ -77,9 +97,14 @@ func setupScheduler(t *testing.T, reg prometheus.Registerer, opts ...connect.Han
 	return s, schedulerpb.NewSchedulerForFrontendClient(c), schedulerpb.NewSchedulerForQuerierClient(c)
 }
 
-func Test_Timeout(t *testing.T) {
-	s, _, querierClient := setupScheduler(t, nil, connect.WithInterceptors(util.WithTimeout(1*time.Second)))
+func setupSchedulerWithHandlerOpts(t *testing.T, handlerOpts ...connect.HandlerOption) (*Scheduler, schedulerpb.SchedulerForFrontendClient, schedulerpb.SchedulerForQuerierClient) {
+	return setupScheduler(t, schedulerArgs{
+		handlerOpts: handlerOpts,
+	})
+}
 
+func Test_Timeout(t *testing.T) {
+	s, _, querierClient := setupSchedulerWithHandlerOpts(t, connect.WithInterceptors(util.WithTimeout(1*time.Second)))
 	ql, err := querierClient.QuerierLoop(context.Background())
 	require.NoError(t, err)
 	require.NoError(t, ql.Send(&schedulerpb.QuerierToScheduler{QuerierID: "querier-1"}))
@@ -88,7 +113,7 @@ func Test_Timeout(t *testing.T) {
 }
 
 func TestSchedulerBasicEnqueue(t *testing.T) {
-	scheduler, frontendClient, querierClient := setupScheduler(t, nil)
+	scheduler, frontendClient, querierClient := setupScheduler(t, schedulerArgs{})
 
 	frontendLoop := initFrontendLoop(t, frontendClient, "frontend-12345")
 	frontendToScheduler(t, frontendLoop, &schedulerpb.FrontendToScheduler{
@@ -116,7 +141,7 @@ func TestSchedulerBasicEnqueue(t *testing.T) {
 }
 
 func TestSchedulerEnqueueWithCancel(t *testing.T) {
-	scheduler, frontendClient, querierClient := setupScheduler(t, nil)
+	scheduler, frontendClient, querierClient := setupScheduler(t, schedulerArgs{})
 
 	frontendLoop := initFrontendLoop(t, frontendClient, "frontend-12345")
 	frontendToScheduler(t, frontendLoop, &schedulerpb.FrontendToScheduler{
@@ -146,7 +171,7 @@ func initQuerierLoop(t *testing.T, querierClient schedulerpb.SchedulerForQuerier
 }
 
 func TestSchedulerEnqueueByMultipleFrontendsWithCancel(t *testing.T) {
-	scheduler, frontendClient, querierClient := setupScheduler(t, nil)
+	scheduler, frontendClient, querierClient := setupScheduler(t, schedulerArgs{})
 
 	frontendLoop1 := initFrontendLoop(t, frontendClient, "frontend-1")
 	frontendLoop2 := initFrontendLoop(t, frontendClient, "frontend-2")
@@ -187,7 +212,7 @@ func TestSchedulerEnqueueByMultipleFrontendsWithCancel(t *testing.T) {
 }
 
 func TestSchedulerEnqueueWithFrontendDisconnect(t *testing.T) {
-	scheduler, frontendClient, querierClient := setupScheduler(t, nil)
+	scheduler, frontendClient, querierClient := setupScheduler(t, schedulerArgs{})
 
 	frontendLoop := initFrontendLoop(t, frontendClient, "frontend-12345")
 	frontendToScheduler(t, frontendLoop, &schedulerpb.FrontendToScheduler{
@@ -217,7 +242,7 @@ func TestSchedulerEnqueueWithFrontendDisconnect(t *testing.T) {
 }
 
 func TestCancelRequestInProgress(t *testing.T) {
-	scheduler, frontendClient, querierClient := setupScheduler(t, nil)
+	scheduler, frontendClient, querierClient := setupScheduler(t, schedulerArgs{})
 
 	frontendLoop := initFrontendLoop(t, frontendClient, "frontend-12345")
 	frontendToScheduler(t, frontendLoop, &schedulerpb.FrontendToScheduler{
@@ -249,7 +274,7 @@ func TestCancelRequestInProgress(t *testing.T) {
 }
 
 func TestTracingContext(t *testing.T) {
-	scheduler, frontendClient, _ := setupScheduler(t, nil)
+	scheduler, frontendClient, _ := setupScheduler(t, schedulerArgs{})
 
 	frontendLoop := initFrontendLoop(t, frontendClient, "frontend-12345")
 
@@ -280,7 +305,7 @@ func TestTracingContext(t *testing.T) {
 }
 
 func TestSchedulerShutdown_FrontendLoop(t *testing.T) {
-	scheduler, frontendClient, _ := setupScheduler(t, nil)
+	scheduler, frontendClient, _ := setupScheduler(t, schedulerArgs{})
 
 	frontendLoop := initFrontendLoop(t, frontendClient, "frontend-12345")
 
@@ -301,7 +326,7 @@ func TestSchedulerShutdown_FrontendLoop(t *testing.T) {
 }
 
 func TestSchedulerShutdown_QuerierLoop(t *testing.T) {
-	scheduler, frontendClient, querierClient := setupScheduler(t, nil)
+	scheduler, frontendClient, querierClient := setupScheduler(t, schedulerArgs{})
 
 	frontendLoop := initFrontendLoop(t, frontendClient, "frontend-12345")
 	frontendToScheduler(t, frontendLoop, &schedulerpb.FrontendToScheduler{
@@ -333,7 +358,7 @@ func TestSchedulerShutdown_QuerierLoop(t *testing.T) {
 }
 
 func TestSchedulerMaxOutstandingRequests(t *testing.T) {
-	_, frontendClient, _ := setupScheduler(t, nil)
+	_, frontendClient, _ := setupScheduler(t, schedulerArgs{})
 
 	for i := 0; i < testMaxOutstandingPerTenant; i++ {
 		// coming from different frontends
@@ -365,32 +390,32 @@ func TestSchedulerMaxOutstandingRequests(t *testing.T) {
 }
 
 func TestSchedulerForwardsErrorToFrontend(t *testing.T) {
-	_, frontendClient, querierClient := setupScheduler(t, nil)
+
+	l := bufconn.Listen(256 << 10)
+	_, frontendClient, querierClient := setupScheduler(t, schedulerArgs{
+		// Have the scheduler use the in-memory connection to call back into the frontend.
+		dialOpts: []grpc.DialOption{
+			grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return l.Dial() }),
+		},
+	})
 
 	fm := &frontendMock{resp: map[uint64]*httpgrpc.HTTPResponse{}}
-	frontendAddress := ""
 
 	// Setup frontend grpc server
 	{
 		frontendGrpcServer := grpc.NewServer()
 		frontendpb.RegisterFrontendForQuerierServer(frontendGrpcServer, fm)
 
-		l, err := net.Listen("tcp", "127.0.0.1:")
-		require.NoError(t, err)
-
-		frontendAddress = l.Addr().String()
-
 		go func() {
 			_ = frontendGrpcServer.Serve(l)
 		}()
-
 		t.Cleanup(func() {
 			_ = l.Close()
 		})
 	}
 
 	// After preparations, start frontend and querier.
-	frontendLoop := initFrontendLoop(t, frontendClient, frontendAddress)
+	frontendLoop := initFrontendLoop(t, frontendClient, "irrelevant://because-we-use-in-memory-connection")
 	frontendToScheduler(t, frontendLoop, &schedulerpb.FrontendToScheduler{
 		Type:        schedulerpb.FrontendToSchedulerType_ENQUEUE,
 		QueryID:     100,
@@ -427,7 +452,7 @@ func TestSchedulerForwardsErrorToFrontend(t *testing.T) {
 func TestSchedulerMetrics(t *testing.T) {
 	reg := prometheus.NewPedanticRegistry()
 
-	scheduler, frontendClient, _ := setupScheduler(t, reg)
+	scheduler, frontendClient, _ := setupScheduler(t, schedulerArgs{reg: reg})
 
 	frontendLoop := initFrontendLoop(t, frontendClient, "frontend-12345")
 	frontendToScheduler(t, frontendLoop, &schedulerpb.FrontendToScheduler{