Merge #145907

145907: rpc, server: properly handle DRPC stream context cancellations r=cthumuluru-crdb a=cthumuluru-crdb

The graceful way to cancel a stream in DRPC is to close the stream. When stream context is canceled, DRPC closes the underlying transport by default as we see `<manager closed>` errors as a result. storj/drpc@a5d487a/drpcmanager/manager.go#L363-L378

`SoftClose` option on the DRPC manager allows for reusing the underlying transport as long as it's not busy (active write). In our case, since the stream is idle and at any point in time only one stream can use a connection we can use that option to avoid too many connection reopens. storj/drpc@a5d487a/drpcmanager/manager.go#L348-L363

Fixes: #140670
Epic: CRDB-48929
Release note: none

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

Author: craig[bot]

pkg/rpc/drpc.go

@@ -10,6 +10,7 @@ import (
 	"crypto/tls"
 	"math"
 	"net"
+	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
@@ -28,6 +29,9 @@ import (
 // have the DRPC server enabled.
 var ErrDRPCDisabled = errors.New("DRPC is not enabled")
 
+// Default idle connection timeout for DRPC connections in the pool.
+var defaultDRPCConnIdleTimeout = 5 * time.Minute
+
 type drpcServerI interface {
 	Serve(ctx context.Context, lis net.Listener) error
 }
@@ -93,7 +97,9 @@ func newDRPCServer(_ context.Context, rpcCtx *Context) (*DRPCServer, error) {
 func dialDRPC(rpcCtx *Context) func(ctx context.Context, target string) (drpcpool.Conn, error) {
 	return func(ctx context.Context, target string) (drpcpool.Conn, error) {
 		// TODO(server): could use connection class instead of empty key here.
-		pool := drpcpool.New[struct{}, drpcpool.Conn](drpcpool.Options{})
+		pool := drpcpool.New[struct{}, drpcpool.Conn](drpcpool.Options{
+			Expiration: defaultDRPCConnIdleTimeout,
+		})
 		pooledConn := pool.Get(ctx /* unused */, struct{}{}, func(ctx context.Context,
 			_ struct{}) (drpcpool.Conn, error) {
 
@@ -110,6 +116,7 @@ func dialDRPC(rpcCtx *Context) func(ctx context.Context, target string) (drpcpoo
 					Stream: drpcstream.Options{
 						MaximumBufferSize: 0, // unlimited
 					},
+					SoftCancel: true, // don't close the transport when stream context is canceled
 				},
 			}
 			var conn *drpcconn.Conn

pkg/server/BUILD.bazel

@@ -618,6 +618,7 @@ go_test(
         "@in_gopkg_yaml_v2//:yaml_v2",
         "@io_opentelemetry_go_otel//attribute",
         "@io_storj_drpc//drpcconn",
+        "@io_storj_drpc//drpcmanager",
         "@io_storj_drpc//drpcmigrate",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//codes",

pkg/server/drpc_test.go

@@ -20,6 +20,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/stretchr/testify/require"
 	"storj.io/drpc/drpcconn"
+	"storj.io/drpc/drpcmanager"
 	"storj.io/drpc/drpcmigrate"
 )
 
@@ -79,3 +80,82 @@ func TestDRPCBatchServer(t *testing.T) {
 		require.NoError(t, err)
 	})
 }
+
+func TestStreamContextCancel(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	const numNodes = 1
+	args := base.TestClusterArgs{
+		ReplicationMode: base.ReplicationManual,
+		ServerArgs: base.TestServerArgs{
+			Settings: cluster.MakeClusterSettings(),
+		},
+	}
+
+	ctx := context.Background()
+	rpc.ExperimentalDRPCEnabled.Override(ctx, &args.ServerArgs.Settings.SV, true)
+	c := testcluster.StartTestCluster(t, numNodes, args)
+	defer c.Stopper().Stop(ctx)
+
+	rpcAddr := c.Server(0).RPCAddr()
+
+	// Dial the drpc server with the drpc connection header.
+	rawconn, err := drpcmigrate.DialWithHeader(ctx, "tcp", rpcAddr, drpcmigrate.DRPCHeader)
+	require.NoError(t, err)
+
+	cm, err := c.Server(0).RPCContext().GetCertificateManager()
+	require.NoError(t, err)
+	tlsCfg, err := cm.GetNodeClientTLSConfig()
+	require.NoError(t, err)
+	tlsCfg = tlsCfg.Clone()
+	tlsCfg.ServerName = "*.local"
+	tlsConn := tls.Client(rawconn, tlsCfg)
+	conn := drpcconn.NewWithOptions(tlsConn, drpcconn.Options{
+		Manager: drpcmanager.Options{
+			SoftCancel: true, // don't close the transport when stream context is canceled
+		},
+	})
+	defer func() {
+		require.NoError(t, conn.Close())
+	}()
+
+	desc := c.LookupRangeOrFatal(t, c.ScratchRange(t))
+	client := kvpb.NewDRPCBatchClient(conn)
+
+	singleRequest := func() {
+		streamCtx, streamCtxCancel := context.WithCancel(ctx)
+		defer streamCtxCancel()
+
+		s, err := client.BatchStream(streamCtx)
+		require.NoError(t, err)
+
+		ba := &kvpb.BatchRequest{}
+		ba.RangeID = desc.RangeID
+
+		var ok bool
+		ba.Replica, ok = desc.GetReplicaDescriptor(1)
+		require.True(t, ok)
+
+		req := &kvpb.LeaseInfoRequest{}
+		req.Key = desc.StartKey.AsRawKey()
+		ba.Add(req)
+
+		err = s.Send(ba)
+		require.NoError(t, err)
+
+		_, err = s.Recv()
+		require.NoError(t, err)
+	}
+
+	// Make two consecutive stream requests using the same connection.
+	for i := 0; i < 2; i++ {
+		select {
+		case <-conn.Closed():
+			t.Fatal("connection closed unexpectedly")
+		default:
+		}
+
+		singleRequest()
+	}
+}