rpc, server: add loopback dialer support for DRPC

Observed a performance regression in the sysbench `oltp_read_only`
microbenchmark when DRPC was enabled by default (ref: #151312). No such
regression was seen in `oltp_read_write`. Root cause analysis traced it to
missing loopback dialer support in DRPC.

Epic: CRDB-51459
Fixes: #148493
Release note: None

Author: cthumuluru-crdb

pkg/rpc/BUILD.bazel

@@ -183,6 +183,8 @@ go_test(
         "@com_github_stretchr_testify//require",
         "@io_storj_drpc//drpcctx",
         "@io_storj_drpc//drpcmetadata",
+        "@io_storj_drpc//drpcmux",
+        "@io_storj_drpc//drpcserver",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//credentials",

pkg/rpc/context.go

@@ -54,6 +54,7 @@ import (
 	"google.golang.org/grpc/stats"
 	"storj.io/drpc"
 	"storj.io/drpc/drpcclient"
+	"storj.io/drpc/drpcmigrate"
 )
 
 // NewServer sets up an RPC server. Depending on the ServerOptions, the Server
@@ -287,6 +288,9 @@ type Context struct {
 	// the gRPC protocol over an in-memory pipe.
 	loopbackDialFn func(context.Context) (net.Conn, error)
 
+	// This is similar to the loopbackDialFn above, but for DRPC connections.
+	loopbackDRPCDialFn func(context.Context) (net.Conn, error)
+
 	// clientCreds is used to pass additional headers to called RPCs.
 	clientCreds credentials.PerRPCCredentials
 
@@ -295,7 +299,7 @@ type Context struct {
 	windowSizeSettings
 }
 
-// SetLoopbackDialer configures the loopback dialer function.
+// SetLoopbackDialer configures the loopback dialer function to dial gRPC connections.
 func (c *Context) SetLoopbackDialer(loopbackDialFn func(context.Context) (net.Conn, error)) {
 	if c.ContextOptions.Knobs.NoLoopbackDialer {
 		// A test has decided it is opting out of the special loopback
@@ -306,6 +310,17 @@ func (c *Context) SetLoopbackDialer(loopbackDialFn func(context.Context) (net.Co
 	c.loopbackDialFn = loopbackDialFn
 }
 
+// SetLoopbackDRPCDialer configures the loopback dialer function to dial DRPC connections.
+func (c *Context) SetLoopbackDRPCDialer(loopbackDialFn func(context.Context) (net.Conn, error)) {
+	if c.ContextOptions.Knobs.NoLoopbackDialer {
+		// A test has decided it is opting out of the special loopback
+		// dialing mechanism. Obey it. We already have defined
+		// loopbackDialFn in that case in NewContext().
+		return
+	}
+	c.loopbackDRPCDialFn = loopbackDialFn
+}
+
 // StoreLivenessGracePeriod computes the grace period after a store restarts before which it will
 // not withdraw support from other stores.
 func (c *Context) StoreLivenessWithdrawalGracePeriod() time.Duration {
@@ -596,6 +611,9 @@ func NewContext(ctx context.Context, opts ContextOptions) *Context {
 			d := onlyOnceDialer{}
 			return d.dial(ctx, opts.AdvertiseAddr)
 		}
+		rpcCtx.loopbackDRPCDialFn = func(ctx context.Context) (net.Conn, error) {
+			return drpcmigrate.DialWithHeader(ctx, "tcp", opts.AdvertiseAddr, drpcmigrate.DRPCHeader)
+		}
 	}
 
 	// We only monitor remote clocks in server-to-server connections.

pkg/rpc/context_test.go

@@ -8,6 +8,7 @@ package rpc
 import (
 	"context"
 	"crypto/rand"
+	"crypto/tls"
 	"fmt"
 	"io"
 	"net"
@@ -50,6 +51,8 @@ import (
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
+	"storj.io/drpc/drpcmux"
+	"storj.io/drpc/drpcserver"
 )
 
 // TestingConnHealth returns nil if we have an open connection to the given
@@ -807,6 +810,86 @@ func TestConnectLoopback(t *testing.T) {
 	require.Equal(t, codes.DataLoss, gogostatus.Code(errors.UnwrapAll(err)), "%+v", err)
 }
 
+// TestDRPCConnectLoopback verifies that we correctly go through the internal
+// server when dialing the local address.
+func TestDRPCConnectLoopback(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	stopper := stop.NewStopper()
+	defer stopper.Stop(context.Background())
+
+	clock := timeutil.NewManualTime(timeutil.Unix(0, 1))
+	maxOffset := time.Duration(0)
+
+	clusterID := uuid.MakeV4()
+	const nodeID = 1
+
+	ctx := context.Background()
+	clientCtx := newTestContext(clusterID, clock, maxOffset, stopper)
+	clientCtx.NodeID.Set(context.Background(), nodeID)
+	drpcLoopbackL := netutil.NewLoopbackListener(ctx, stopper)
+	clientCtx.loopbackDRPCDialFn = func(ctx context.Context) (net.Conn, error) {
+		return drpcLoopbackL.Connect(ctx)
+	}
+	// Ensure that there's no error connecting to the local node when an internal
+	// server has been registered.
+	clientCtx.SetLocalInternalServer(
+		&internalServer{},
+		ServerInterceptorInfo{}, ClientInterceptorInfo{})
+
+	// Create a DRPC server to listen on loopback.
+	newTestServer := func() *drpcServer {
+		d := &drpcServer{}
+		mux := drpcmux.New()
+		d.Server = drpcserver.NewWithOptions(mux, drpcserver.Options{})
+		d.Mux = mux
+
+		return d
+	}
+	drpcS := newTestServer()
+
+	// Register a heartbeat service that always returns an error.
+	errLoopback := gogostatus.Newf(codes.Unknown, "loopback!").Err() // TODO(server): replace code with DRPC error codes
+	require.NoError(t, DRPCRegisterHeartbeat(drpcS, &ManualHeartbeatService{readyFn: func() error {
+		return errLoopback
+	}}))
+
+	tlsConfig, err := clientCtx.GetServerTLSConfig()
+	require.NoError(t, err)
+
+	// Start the DRPC server on loopback.
+	require.NoError(t, stopper.RunAsyncTask(ctx, "listen-server-loopback", func(ctx context.Context) {
+		drpcLoopbackTLSL := tls.NewListener(drpcLoopbackL, tlsConfig)
+		netutil.FatalIfUnexpected(drpcS.Serve(ctx, drpcLoopbackTLSL))
+	}))
+
+	serveDRPC := func(stopper *stop.Stopper, drpcS *drpcServer, drpcTLSL net.Listener) error {
+		waitQuiesce := func(context.Context) {
+			<-stopper.ShouldQuiesce()
+			netutil.FatalIfUnexpected(drpcTLSL.Close())
+		}
+		if err := stopper.RunAsyncTask(ctx, "listen-quiesce", waitQuiesce); err != nil {
+			waitQuiesce(ctx)
+			return err
+		}
+
+		return stopper.RunAsyncTask(ctx, "listen-serve", func(context.Context) {
+			netutil.FatalIfUnexpected(drpcS.Serve(ctx, drpcTLSL))
+		})
+	}
+	drpcL, err := net.Listen("tcp", util.TestAddr.String())
+	require.NoError(t, err)
+	drpcTLSL := tls.NewListener(drpcL, tlsConfig)
+	require.NoError(t, serveDRPC(stopper, drpcS, drpcTLSL))
+
+	addr := drpcL.Addr().String()
+	clientCtx.AdvertiseAddr = addr
+
+	// Connect and get the error that comes from loopbackLn, proving that we were routed there.
+	_, err = clientCtx.DRPCDialNode(addr, nodeID, roachpb.Locality{}, rpcbase.DefaultClass).Connect(ctx)
+	require.Equal(t, codes.Unknown, gogostatus.Code(errors.UnwrapAll(err)), "%+v", err)
+}
+
 func TestOffsetMeasurement(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 

pkg/rpc/drpc.go

@@ -53,7 +53,12 @@ func DialDRPC(
 		pooledConn := pool.Get(ctx /* unused */, struct{}{}, func(ctx context.Context,
 			_ struct{}) (drpcpool.Conn, error) {
 
-			netConn, err := drpcmigrate.DialWithHeader(ctx, "tcp", target, drpcmigrate.DRPCHeader)
+			netConn, err := func(ctx context.Context) (net.Conn, error) {
+				if rpcCtx.ContextOptions.AdvertiseAddr == target && !rpcCtx.ClientOnly {
+					return rpcCtx.loopbackDRPCDialFn(ctx)
+				}
+				return drpcmigrate.DialWithHeader(ctx, "tcp", target, drpcmigrate.DRPCHeader)
+			}(ctx)
 			if err != nil {
 				return nil, err
 			}

pkg/server/drpc_test.go

@@ -220,16 +220,13 @@ func TestDialDRPC_InterceptorsAreSet(t *testing.T) {
 	args := base.TestClusterArgs{
 		ReplicationMode: base.ReplicationManual,
 		ServerArgs: base.TestServerArgs{
-			Settings: cluster.MakeClusterSettings(),
-			Insecure: true,
+			Settings:          cluster.MakeClusterSettings(),
+			DefaultDRPCOption: base.TestDRPCEnabled,
 		},
 	}
 
-	rpc.ExperimentalDRPCEnabled.Override(ctx, &args.ServerArgs.Settings.SV, true)
 	c := testcluster.StartTestCluster(t, numNodes, args)
 	defer c.Stopper().Stop(ctx)
-	require.True(t, c.Server(0).RPCContext().Insecure)
-
 	rpcAddr := c.Server(0).RPCAddr()
 
 	// Client setup
@@ -238,9 +235,10 @@ func TestDialDRPC_InterceptorsAreSet(t *testing.T) {
 	rpcContextOptions.Stopper = c.Stopper()
 	rpcContextOptions.Settings = c.Server(0).ClusterSettings()
 	rpcCtx := rpc.NewContext(ctx, rpcContextOptions)
-	rpcCtx.ContextOptions = rpc.ContextOptions{Insecure: true}
+
 	// Adding test interceptors
 	rpcCtx.Knobs = rpc.ContextTestingKnobs{
+		NoLoopbackDialer: true,
 		UnaryClientInterceptorDRPC: func(target string, class rpcbase.ConnectionClass) drpcclient.UnaryClientInterceptor {
 			return mockUnaryInterceptor
 		},

pkg/server/server.go

@@ -1659,7 +1659,7 @@ func (s *topLevelServer) PreStart(ctx context.Context) error {
 	// and dispatches the server worker for the RPC.
 	// The SQL listener is returned, to start the SQL server later
 	// below when the server has initialized.
-	pgL, loopbackPgL, rpcLoopbackDialFn, startRPCServer, err := startListenRPCAndSQL(
+	pgL, loopbackPgL, grpcLoopbackDialFn, drpcLoopbackDialFn, startRPCServer, err := startListenRPCAndSQL(
 		ctx, workersCtx, s.cfg.BaseConfig,
 		s.stopper, s.grpc, s.drpc, ListenAndUpdateAddrs, true /* enableSQLListener */, s.cfg.AcceptProxyProtocolHeaders)
 	if err != nil {
@@ -1669,7 +1669,8 @@ func (s *topLevelServer) PreStart(ctx context.Context) error {
 	s.loopbackPgL = loopbackPgL
 
 	// Tell the RPC context how to connect in-memory.
-	s.rpcContext.SetLoopbackDialer(rpcLoopbackDialFn)
+	s.rpcContext.SetLoopbackDialer(grpcLoopbackDialFn)
+	s.rpcContext.SetLoopbackDRPCDialer(drpcLoopbackDialFn)
 
 	if s.cfg.TestingKnobs.Server != nil {
 		knobs := s.cfg.TestingKnobs.Server.(*TestingKnobs)

pkg/server/start_listen.go

@@ -49,7 +49,8 @@ func startListenRPCAndSQL(
 ) (
 	sqlListener net.Listener,
 	pgLoopbackListener *netutil.LoopbackListener,
-	rpcLoopbackDial func(context.Context) (net.Conn, error),
+	grpcLoopbackDial func(context.Context) (net.Conn, error),
+	drpcLoopbackDial func(context.Context) (net.Conn, error),
 	startRPCServer func(ctx context.Context),
 	err error,
 ) {
@@ -66,7 +67,7 @@ func startListenRPCAndSQL(
 		var err error
 		ln, err = rpcListenerFactory(ctx, &cfg.Addr, &cfg.AdvertiseAddr, rpcChanName, acceptProxyProtocolHeaders)
 		if err != nil {
-			return nil, nil, nil, nil, err
+			return nil, nil, nil, nil, nil, err
 		}
 		log.Eventf(ctx, "listening on port %s", cfg.Addr)
 	}
@@ -79,7 +80,7 @@ func startListenRPCAndSQL(
 			pgL = cfg.SQLAddrListener
 		}
 		if err != nil {
-			return nil, nil, nil, nil, err
+			return nil, nil, nil, nil, nil, err
 		}
 		// The SQL listener shutdown worker, which closes everything under
 		// the SQL port when the stopper indicates we are shutting down.
@@ -95,7 +96,7 @@ func startListenRPCAndSQL(
 		}
 		if err := stopper.RunAsyncTask(workersCtx, "wait-quiesce", waitQuiesce); err != nil {
 			waitQuiesce(workersCtx)
-			return nil, nil, nil, nil, err
+			return nil, nil, nil, nil, nil, err
 		}
 		log.Eventf(ctx, "listening on sql port %s", cfg.SQLAddr)
 	}
@@ -129,7 +130,7 @@ func startListenRPCAndSQL(
 		// Then we update the advertised addr with the right port, if
 		// the port had been auto-allocated.
 		if err := UpdateAddrs(ctx, &cfg.SQLAddr, &cfg.SQLAdvertiseAddr, ln.Addr()); err != nil {
-			return nil, nil, nil, nil, errors.Wrapf(err, "internal error")
+			return nil, nil, nil, nil, nil, errors.Wrapf(err, "internal error")
 		}
 	}
 
@@ -158,18 +159,22 @@ func startListenRPCAndSQL(
 		}
 	}
 
-	rpcLoopbackL := netutil.NewLoopbackListener(ctx, stopper)
+	grpcLoopbackL := netutil.NewLoopbackListener(ctx, stopper)
 	sqlLoopbackL := netutil.NewLoopbackListener(ctx, stopper)
 	drpcCtx, drpcCancel := context.WithCancel(workersCtx)
+	// Create a dedicated DRPC loopback listener. Since this listener is exclusively for DRPC,
+	// no protocol header inspection is needed.
+	drpcLoopbackL := netutil.NewLoopbackListener(ctx, stopper)
 
 	// The remainder shutdown worker.
 	waitForQuiesce := func(context.Context) {
 		<-stopper.ShouldQuiesce()
 		drpcCancel()
 		// TODO(bdarnell): Do we need to also close the other listeners?
 		netutil.FatalIfUnexpected(grpcL.Close())
-		netutil.FatalIfUnexpected(drpcL.Close())
-		netutil.FatalIfUnexpected(rpcLoopbackL.Close())
+		netutil.FatalIfUnexpected(grpcLoopbackL.Close())
+		netutil.FatalIfUnexpected(drpcL.Close())         // Closing this listener is as good as closing drpcTLSL
+		netutil.FatalIfUnexpected(drpcLoopbackL.Close()) // Closing this listener is as good as closing drpcLoopbackTLSL
 		netutil.FatalIfUnexpected(sqlLoopbackL.Close())
 		netutil.FatalIfUnexpected(ln.Close())
 	}
@@ -191,7 +196,7 @@ func startListenRPCAndSQL(
 		waitForQuiesce(ctx)
 		stopGRPC()
 		drpcCancel()
-		return nil, nil, nil, nil, err
+		return nil, nil, nil, nil, nil, err
 	}
 	stopper.AddCloser(stop.CloserFn(stopGRPC))
 
@@ -213,7 +218,15 @@ func startListenRPCAndSQL(
 			}
 		})
 		_ = stopper.RunAsyncTask(workersCtx, "serve-loopback-grpc", func(context.Context) {
-			netutil.FatalIfUnexpected(grpc.Serve(rpcLoopbackL))
+			netutil.FatalIfUnexpected(grpc.Serve(grpcLoopbackL))
+		})
+		_ = stopper.RunAsyncTask(workersCtx, "serve-loopback-drpc", func(context.Context) {
+			if cfg := drpc.tlsCfg; cfg != nil {
+				drpcdrpcLoopbackTLSL := tls.NewListener(drpcLoopbackL, cfg)
+				netutil.FatalIfUnexpected(drpc.Serve(ctx, drpcdrpcLoopbackTLSL))
+			} else {
+				netutil.FatalIfUnexpected(drpc.Serve(ctx, drpcLoopbackL))
+			}
 		})
 
 		_ = stopper.RunAsyncTask(ctx, "serve-mux", func(context.Context) {
@@ -223,5 +236,5 @@ func startListenRPCAndSQL(
 		})
 	}
 
-	return pgL, sqlLoopbackL, rpcLoopbackL.Connect, startRPCServer, nil
+	return pgL, sqlLoopbackL, grpcLoopbackL.Connect, drpcLoopbackL.Connect, startRPCServer, nil
 }

pkg/server/tenant.go

@@ -604,7 +604,7 @@ func (s *SQLServerWrapper) PreStart(ctx context.Context) error {
 		lf = s.sqlServer.cfg.RPCListenerFactory
 	}
 
-	pgL, loopbackPgL, rpcLoopbackDialFn, startRPCServer, err := startListenRPCAndSQL(
+	pgL, loopbackPgL, grpcLoopbackDialFn, drpcLoopbackDialFn, startRPCServer, err := startListenRPCAndSQL(
 		ctx, workersCtx, *s.sqlServer.cfg, s.stopper,
 		s.grpc, s.drpc, lf, enableSQLListener, s.cfg.AcceptProxyProtocolHeaders)
 	if err != nil {
@@ -616,7 +616,8 @@ func (s *SQLServerWrapper) PreStart(ctx context.Context) error {
 	s.loopbackPgL = loopbackPgL
 
 	// Tell the RPC context how to connect in-memory.
-	s.rpcContext.SetLoopbackDialer(rpcLoopbackDialFn)
+	s.rpcContext.SetLoopbackDialer(grpcLoopbackDialFn)
+	s.rpcContext.SetLoopbackDRPCDialer(drpcLoopbackDialFn)
 
 	// NB: This is where (*Server).PreStart() reports the listener readiness
 	// via testing knobs: PauseAfterGettingRPCAddress, SignalAfterGettingRPCAddress.