Merge #152717 #152921 #152923

152717: rpc, server: add loopback dialer support for DRPC r=cthumuluru-crdb a=cthumuluru-crdb

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

152921: kvcoord: correctly initialize txnWriteBuffer r=yuzefovich a=stevendanna

Previously we were initializing the pipelineEnabler on a copy of the txnWriteBuffer that would later be overwritten. As a result, we were never re-enabling pipelining after a mid-transaction flush. This is OK from a correctness perspective since this was only intended as a potential performance improvement for some workloads under buffered writes.

There is a potentially strong argument that we shouldn't do this at all and that if someone sees a regression of their workload under buffered writes they should turn it off rather than relying on this fallback.

Epic: None
Release note: None

152923: kvpb: always use t.ExpectExclusionSince in EarliestActiveTimestamp r=miraradeva a=stevendanna

Previously, some cases used ExpectExclusionSince while others used ExpectExclusionSince.Next().

I think this was OK in practice because any write we needed to observe would need to be at ExpectExclusionSince.Next() since we set that field based on the ReadTimestamp of a locking Get. Such Get's scan with FailOnMoreRecent set which, despite the name, also fails if the current value equals the read timestamp. Thus, any new write violating our exclusion would have to be at .Next() or higher.

If the above analysis doesn't convince your or is wrong, also note that this is only required for weak isolation buffered writes which can only be enabled via a non-public cluster setting (SSI transactions refresh their reads).

In case the argument above is wrong, I've chosen to normalize these to the earlier timestamp.

I've forgone a release note given the above argument.

Epic: none
Release note: None

Co-authored-by: Chandra Thumuluru <[email protected]>
Co-authored-by: Steven Danna <[email protected]>

Author: 3 people

pkg/kv/kvclient/kvcoord/txn_coord_sender.go

@@ -281,9 +281,6 @@ func newRootTxnCoordSender(
 		timeSource: timeutil.DefaultTimeSource{},
 		txn:        &tcs.mu.txn,
 	}
-	tcs.interceptorAlloc.txnWriteBuffer.init(
-		&tcs.interceptorAlloc.txnPipeliner,
-	)
 
 	tcs.initCommonInterceptors(tcf, txn, kv.RootTxn)
 
@@ -363,6 +360,7 @@ func (tc *TxnCoordSender) initCommonInterceptors(
 		allowConcurrentRequests: typ == kv.LeafTxn,
 	}
 	tc.interceptorAlloc.txnSeqNumAllocator.writeSeq = txn.Sequence
+	tc.interceptorAlloc.txnWriteBuffer.init(&tc.interceptorAlloc.txnPipeliner)
 }
 
 func (tc *TxnCoordSender) connectInterceptors() {

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer_client_test.go

@@ -92,3 +92,71 @@ func TestTxnCoordSenderWriteBufferingDisablesPipelining(t *testing.T) {
 		kvpb.Scan, kvpb.Put, kvpb.EndTxn,
 	}, calls)
 }
+
+// TestTxnCoordSenderWriteBufferingReEnablesPipelining verifies that pipelining
+// is re-enabled after a mid-transaction flush.
+func TestTxnCoordSenderWriteBufferingReEnablesPipelining(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	s := serverutils.StartServerOnly(t, base.TestServerArgs{})
+	defer s.Stopper().Stop(ctx)
+
+	distSender := s.DistSenderI().(*DistSender)
+	batchCount := 0
+	var calls []kvpb.Method
+	var senderFn kv.SenderFunc = func(
+		ctx context.Context, ba *kvpb.BatchRequest,
+	) (*kvpb.BatchResponse, *kvpb.Error) {
+		batchCount++
+		t.Logf("Batch: %#+v", ba.Methods())
+		calls = append(calls, ba.Methods()...)
+		if et, ok := ba.GetArg(kvpb.EndTxn); ok {
+			// Ensure that no transactions enter a STAGING state.
+			et.(*kvpb.EndTxnRequest).InFlightWrites = nil
+		}
+		return distSender.Send(ctx, ba)
+	}
+
+	st := s.ClusterSettings()
+	BufferedWritesMaxBufferSize.Override(ctx, &st.SV, defaultBufferSize)
+
+	tsf := NewTxnCoordSenderFactory(TxnCoordSenderFactoryConfig{
+		AmbientCtx: s.AmbientCtx(),
+		Settings:   st,
+		Clock:      s.Clock(),
+		Stopper:    s.Stopper(),
+		// Disable transaction heartbeats so that they don't disrupt our attempt to
+		// track the requests issued by the transactions.
+		HeartbeatInterval: -1,
+	}, senderFn)
+	db := kv.NewDB(s.AmbientCtx(), tsf, s.Clock(), s.Stopper())
+
+	require.NoError(t, db.Put(ctx, "test-key-a", "hello"))
+	require.NoError(t, db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+		txn.SetBufferedWritesEnabled(true)
+		if err := txn.Put(ctx, "test-key-c", "hello"); err != nil {
+			return err
+		}
+		if _, err := txn.DelRange(ctx, "test-key", "test-key-d", true); err != nil {
+			return err
+		}
+		if err := txn.Put(ctx, "test-key-a", "hello"); err != nil {
+			return err
+		}
+		return nil
+	}))
+
+	require.Equal(t, 4, batchCount)
+	require.Equal(t, []kvpb.Method{
+		// The initial setup
+		kvpb.Put,
+		// The first (buffered) Put and the DeleteRange that flushes the buffer.
+		kvpb.Put, kvpb.DeleteRange,
+		// The second (pipelined) Put
+		kvpb.Put,
+		// EndTxn with the QueryIntent because pipelining was turned back on.
+		kvpb.QueryIntent, kvpb.EndTxn,
+	}, calls)
+}

pkg/kv/kvpb/batch.go

@@ -101,7 +101,7 @@ func (ba *BatchRequest) EarliestActiveTimestamp() hlc.Timestamp {
 			//
 			// See the example in RefreshRequest for more details.
 			if !t.ExpectExclusionSince.IsEmpty() {
-				ts.Backward(t.ExpectExclusionSince.Next())
+				ts.Backward(t.ExpectExclusionSince)
 			}
 		case *ExportRequest:
 			if !t.StartTime.IsEmpty() {
@@ -115,7 +115,7 @@ func (ba *BatchRequest) EarliestActiveTimestamp() hlc.Timestamp {
 			//
 			// See the example in RefreshRequest for more details.
 			if !t.ExpectExclusionSince.IsEmpty() {
-				ts.Backward(t.ExpectExclusionSince.Next())
+				ts.Backward(t.ExpectExclusionSince)
 			}
 		case *PutRequest:
 			// A PutRequest with ExpectExclusionSince set need to be able to observe MVCC

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)