Merge #147287

147287: server: refactor drpc server out of grpcServer r=cthumuluru-crdb a=shubhamdhama

This change extracts the DRPC server from grpcServer. It also introduces a `serve` mode, mirroring the gRPC implementation.

Fixes: #147285
Epic: CRDB-48925
Release note: None

Co-authored-by: Shubham Dhama <[email protected]>

Author: craig[bot]

pkg/rpc/BUILD.bazel

@@ -81,9 +81,7 @@ go_library(
         "@io_storj_drpc//drpcconn",
         "@io_storj_drpc//drpcmanager",
         "@io_storj_drpc//drpcmigrate",
-        "@io_storj_drpc//drpcmux",
         "@io_storj_drpc//drpcpool",
-        "@io_storj_drpc//drpcserver",
         "@io_storj_drpc//drpcstream",
         "@io_storj_drpc//drpcwire",
         "@org_golang_google_grpc//:grpc",

pkg/rpc/context.go

@@ -56,7 +56,7 @@ import (
 // either expects incoming connections from KV nodes, or from tenant SQL
 // servers.
 func NewServer(ctx context.Context, rpcCtx *Context, opts ...ServerOption) (*grpc.Server, error) {
-	srv, _, _, err := NewServerEx(ctx, rpcCtx, opts...)
+	srv, _, err := NewServerEx(ctx, rpcCtx, opts...)
 	return srv, err
 }
 
@@ -84,7 +84,7 @@ type ClientInterceptorInfo struct {
 // internalClientAdapter does).
 func NewServerEx(
 	ctx context.Context, rpcCtx *Context, opts ...ServerOption,
-) (s *grpc.Server, d *DRPCServer, sii ServerInterceptorInfo, err error) {
+) (s *grpc.Server, sii ServerInterceptorInfo, err error) {
 	var o serverOpts
 	for _, f := range opts {
 		f(&o)
@@ -113,7 +113,7 @@ func NewServerEx(
 	if !rpcCtx.ContextOptions.Insecure {
 		tlsConfig, err := rpcCtx.GetServerTLSConfig()
 		if err != nil {
-			return nil, nil, sii, err
+			return nil, sii, err
 		}
 		grpcOpts = append(grpcOpts, grpc.Creds(newTLSCipherRestrictCred(tlsConfig)))
 	}
@@ -194,13 +194,10 @@ func NewServerEx(
 	grpcOpts = append(grpcOpts, grpc.ChainStreamInterceptor(streamInterceptor...))
 
 	s = grpc.NewServer(grpcOpts...)
-	d, err = newDRPCServer(ctx, rpcCtx)
-	if err != nil {
-		return nil, nil, ServerInterceptorInfo{}, err
-	}
+
 	RegisterHeartbeatServer(s, rpcCtx.NewHeartbeatService())
 
-	return s, d, ServerInterceptorInfo{
+	return s, ServerInterceptorInfo{
 		UnaryInterceptors:  unaryInterceptor,
 		StreamInterceptors: streamInterceptor,
 	}, nil

pkg/rpc/context_test.go

@@ -435,7 +435,7 @@ func TestInternalClientAdapterRunsInterceptors(t *testing.T) {
 	serverCtx.AdvertiseAddr = "127.0.0.1:8888"
 	serverCtx.NodeID.Set(context.Background(), 1)
 
-	_ /* gRPC server */, _ /* drpc server */, serverInterceptors, err := NewServerEx(ctx, serverCtx)
+	_ /* gRPC server */, serverInterceptors, err := NewServerEx(ctx, serverCtx)
 	require.NoError(t, err)
 
 	// Pile on one more interceptor to make sure it's called.
@@ -536,7 +536,7 @@ func TestInternalClientAdapterWithClientStreamInterceptors(t *testing.T) {
 	serverCtx.AdvertiseAddr = "127.0.0.1:8888"
 	serverCtx.NodeID.Set(context.Background(), 1)
 
-	_ /* gRPC server */, _ /* drpc server */, serverInterceptors, err := NewServerEx(ctx, serverCtx)
+	_ /* gRPC server */, serverInterceptors, err := NewServerEx(ctx, serverCtx)
 	require.NoError(t, err)
 	var clientInterceptors ClientInterceptorInfo
 	var s *testClientStream
@@ -599,7 +599,7 @@ func TestInternalClientAdapterWithServerStreamInterceptors(t *testing.T) {
 	serverCtx.AdvertiseAddr = "127.0.0.1:8888"
 	serverCtx.NodeID.Set(context.Background(), 1)
 
-	_ /* gRPC server */, _ /* drpc server */, serverInterceptors, err := NewServerEx(ctx, serverCtx)
+	_ /* gRPC server */, serverInterceptors, err := NewServerEx(ctx, serverCtx)
 	require.NoError(t, err)
 
 	const int1Name = "interceptor 1"
@@ -737,7 +737,7 @@ func BenchmarkInternalClientAdapter(b *testing.B) {
 	serverCtx.AdvertiseAddr = "127.0.0.1:8888"
 	serverCtx.NodeID.Set(context.Background(), 1)
 
-	_ /* gRPC server */, _ /* drpc server */, interceptors, err := NewServerEx(ctx, serverCtx)
+	_ /* gRPC server */, interceptors, err := NewServerEx(ctx, serverCtx)
 	require.NoError(b, err)
 
 	internal := &internalServer{}
@@ -2369,7 +2369,7 @@ func TestMetricsInterceptor(t *testing.T) {
 		return nil, nil
 	}
 
-	_, _, serverInterceptors, err := NewServerEx(ctx, serverCtx, WithMetricsServerInterceptor(interceptor))
+	_, serverInterceptors, err := NewServerEx(ctx, serverCtx, WithMetricsServerInterceptor(interceptor))
 	require.NoError(t, err)
 	require.GreaterOrEqual(t, len(serverInterceptors.UnaryInterceptors), 2)
 	// make sure that the RequestMetricsInterceptor is the second registered

pkg/rpc/drpc.go

@@ -9,91 +9,20 @@ import (
 	"context"
 	"crypto/tls"
 	"math"
-	"net"
 	"time"
 
-	"github.com/cockroachdb/cockroach/pkg/util/log"
-	"github.com/cockroachdb/errors"
 	"storj.io/drpc"
 	"storj.io/drpc/drpcconn"
 	"storj.io/drpc/drpcmanager"
 	"storj.io/drpc/drpcmigrate"
-	"storj.io/drpc/drpcmux"
 	"storj.io/drpc/drpcpool"
-	"storj.io/drpc/drpcserver"
 	"storj.io/drpc/drpcstream"
 	"storj.io/drpc/drpcwire"
 )
 
-// ErrDRPCDisabled is returned from hosts that in principle could but do not
-// 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
-}
-
-type drpcMuxI interface {
-	Register(srv interface{}, desc drpc.Description) error
-}
-
-type DRPCServer struct {
-	Srv    drpcServerI
-	Mux    drpcMuxI
-	TLSCfg *tls.Config
-}
-
-var _ drpcServerI = (*drpcserver.Server)(nil)
-var _ drpcServerI = (*drpcOffServer)(nil)
-
-func newDRPCServer(_ context.Context, rpcCtx *Context) (*DRPCServer, error) {
-	var dmux drpcMuxI = &drpcOffServer{}
-	var dsrv drpcServerI = &drpcOffServer{}
-	var tlsCfg *tls.Config
-
-	if ExperimentalDRPCEnabled.Get(&rpcCtx.Settings.SV) {
-		mux := drpcmux.New()
-		dsrv = drpcserver.NewWithOptions(mux, drpcserver.Options{
-			Log: func(err error) {
-				log.Warningf(context.Background(), "drpc server error %v", err)
-			},
-			// The reader's max buffer size defaults to 4mb, and if it is exceeded (such
-			// as happens with AddSSTable) the RPCs fail.
-			Manager: drpcmanager.Options{Reader: drpcwire.ReaderOptions{MaximumBufferSize: math.MaxInt}},
-		})
-		dmux = mux
-
-		var err error
-		tlsCfg, err = rpcCtx.GetServerTLSConfig()
-		if err != nil {
-			return nil, err
-		}
-
-		// NB: any server middleware (server interceptors in gRPC parlance) would go
-		// here:
-		//     dmux = whateverMiddleware1(dmux)
-		//     dmux = whateverMiddleware2(dmux)
-		//     ...
-		//
-		// Each middleware must implement the Handler interface:
-		//
-		//   HandleRPC(stream Stream, rpc string) error
-		//
-		// where Stream
-		// See here for an example:
-		// https://github.com/bryk-io/pkg/blob/4da5fbfef47770be376e4022eab5c6c324984bf7/net/drpc/server.go#L91-L101
-	}
-
-	return &DRPCServer{
-		Srv:    dsrv,
-		Mux:    dmux,
-		TLSCfg: tlsCfg,
-	}, nil
-}
-
 func dialDRPC(
 	rpcCtx *Context,
 ) func(ctx context.Context, target string, _ ConnectionClass) (drpc.Conn, error) {
@@ -162,22 +91,4 @@ func (c *closeEntirePoolConn) Close() error {
 	return c.pool.Close()
 }
 
-// drpcOffServer is used for drpcServerI and drpcMuxI if the DRPC server is
-// disabled. It immediately closes accepted connections and returns
-// ErrDRPCDisabled.
-type drpcOffServer struct{}
-
-func (srv *drpcOffServer) Serve(_ context.Context, lis net.Listener) error {
-	conn, err := lis.Accept()
-	if err != nil {
-		return err
-	}
-	_ = conn.Close()
-	return ErrDRPCDisabled
-}
-
-func (srv *drpcOffServer) Register(interface{}, drpc.Description) error {
-	return nil
-}
-
 type DRPCConnection = Connection[drpc.Conn]

pkg/server/BUILD.bazel

@@ -24,6 +24,7 @@ go_library(
         "distsql_flows.go",
         "doc.go",
         "drain.go",
+        "drpc_server.go",
         "env_sampler.go",
         "external_storage_builder.go",
         "fanout_clients.go",
@@ -51,6 +52,7 @@ go_library(
         "rlimit_bsd.go",
         "rlimit_darwin.go",
         "rlimit_unix.go",
+        "serve_mode.go",
         "server.go",
         "server_controller.go",
         "server_controller_channel_orchestrator.go",
@@ -377,7 +379,13 @@ go_library(
         "@com_github_prometheus_client_model//go",
         "@com_github_prometheus_common//expfmt",
         "@in_gopkg_yaml_v2//:yaml_v2",
+        "@io_storj_drpc//:drpc",
+        "@io_storj_drpc//drpcerr",
+        "@io_storj_drpc//drpcmanager",
         "@io_storj_drpc//drpcmigrate",
+        "@io_storj_drpc//drpcmux",
+        "@io_storj_drpc//drpcserver",
+        "@io_storj_drpc//drpcwire",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//metadata",

pkg/server/admin.go

@@ -117,6 +117,7 @@ type adminServer struct {
 	rpcContext       *rpc.Context
 	clock            *hlc.Clock
 	grpc             *grpcServer
+	drpc             *drpcServer
 	db               *kv.DB
 	drainServer      *drainServer
 }
@@ -156,6 +157,7 @@ func newSystemAdminServer(
 	clock *hlc.Clock,
 	distSender *kvcoord.DistSender,
 	grpc *grpcServer,
+	drpc *drpcServer,
 	drainServer *drainServer,
 	s *topLevelServer,
 ) *systemAdminServer {
@@ -172,6 +174,7 @@ func newSystemAdminServer(
 		clock,
 		distSender,
 		grpc,
+		drpc,
 		drainServer,
 	)
 	return &systemAdminServer{
@@ -199,6 +202,7 @@ func newAdminServer(
 	clock *hlc.Clock,
 	distSender *kvcoord.DistSender,
 	grpc *grpcServer,
+	drpc *drpcServer,
 	drainServer *drainServer,
 ) *adminServer {
 	server := &adminServer{
@@ -217,6 +221,7 @@ func newAdminServer(
 		rpcContext:     rpcCtx,
 		clock:          clock,
 		grpc:           grpc,
+		drpc:           drpc,
 		db:             db,
 		drainServer:    drainServer,
 	}
@@ -2100,10 +2105,18 @@ func (s *adminServer) Health(
 
 // checkReadinessForHealthCheck returns a gRPC error.
 func (s *adminServer) checkReadinessForHealthCheck(ctx context.Context) error {
+	// A gRPC server will always be running, so ensure that we check its health
+	// until it is completely removed after the DRPC to gRPC migration.
 	if err := s.grpc.health(ctx); err != nil {
 		return err
 	}
 
+	if s.drpc.enabled {
+		if err := s.drpc.health(ctx); err != nil {
+			return err
+		}
+	}
+
 	if !s.sqlServer.isReady.Load() {
 		return grpcstatus.Errorf(codes.Unavailable, "node is not accepting SQL clients")
 	}
@@ -2138,10 +2151,18 @@ func (s *systemAdminServer) Health(
 
 // checkReadinessForHealthCheck returns a gRPC error.
 func (s *systemAdminServer) checkReadinessForHealthCheck(ctx context.Context) error {
+	// A gRPC server will always be running, so ensure that we check its health
+	// until it is completely removed after the DRPC to gRPC migration.
 	if err := s.grpc.health(ctx); err != nil {
 		return err
 	}
 
+	if s.drpc.enabled {
+		if err := s.drpc.health(ctx); err != nil {
+			return err
+		}
+	}
+
 	status := s.nodeLiveness.GetNodeVitalityFromCache(roachpb.NodeID(s.serverIterator.getID()))
 	if !status.IsLive(livenesspb.AdminHealthCheck) {
 		return grpcstatus.Errorf(codes.Unavailable, "node is not healthy")

pkg/server/drain.go

@@ -119,6 +119,7 @@ type drainServer struct {
 	// stopTrigger is used to request that the server is shut down.
 	stopTrigger  *stopTrigger
 	grpc         *grpcServer
+	drpc         *drpcServer
 	sqlServer    *SQLServer
 	drainSleepFn func(time.Duration)
 	serverCtl    *serverController
@@ -135,6 +136,7 @@ func newDrainServer(
 	stopper *stop.Stopper,
 	stopTrigger *stopTrigger,
 	grpc *grpcServer,
+	drpc *drpcServer,
 	sqlServer *SQLServer,
 ) *drainServer {
 	var drainSleepFn = time.Sleep
@@ -147,6 +149,7 @@ func newDrainServer(
 		stopper:      stopper,
 		stopTrigger:  stopTrigger,
 		grpc:         grpc,
+		drpc:         drpc,
 		sqlServer:    sqlServer,
 		drainSleepFn: drainSleepFn,
 	}
@@ -384,6 +387,7 @@ func (s *drainServer) drainClients(
 	// Set the gRPC mode of the node to "draining" and mark the node as "not ready".
 	// Probes to /health?ready=1 will now notice the change in the node's readiness.
 	s.grpc.setMode(modeDraining)
+	s.drpc.setMode(modeDraining)
 	s.sqlServer.isReady.Store(false)
 
 	// Log the number of connections periodically.