Merge #147822

147822: drpc: move drpc server back to pkg/rpc r=cthumuluru-crdb a=shubhamdhama

`NewRaftTransport` needs a `drpc.Mux` to register its server
implementations, so the straightforward approach is to pass
`drpcServer.Mux`. However, some tests within the `kvserver` package (not
`kvserver_test` itself) also require this `drpcServer.Mux`. This would
necessitate exporting `newDRPCServer` and `drpcServer.mux` as
`NewDRPCServer` and `drpcServer.Mux` respectively. But even then, this
won't resolve the situation; it will introduce a cyclic dependency. Moving
these tests to `kvserver_test` is not a viable option, as they rely on
private methods within `kvserver`. Even if we managed to untangle this
dependency mess, the whole setup just feels wrong and screams for a proper
refactor. We should simply align with how gRPC handles this.

Additionally, as part of this change, it's more appropriate for the DRPC
server implementation residing in `pkg/rpc` to be exposed as an interface
to other packages, rather than as a concrete object.

Epic: CRDB-48925
Release note: None

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

Author: craig[bot]

pkg/rpc/BUILD.bazel

@@ -80,7 +80,9 @@ 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/drpc.go

@@ -9,14 +9,19 @@ import (
 	"context"
 	"crypto/tls"
 	"math"
+	"net"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/rpc/rpcbase"
+	"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"
 )
@@ -93,3 +98,79 @@ func (c *closeEntirePoolConn) Close() error {
 }
 
 type DRPCConnection = Connection[drpc.Conn]
+
+// 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")
+
+// DRPCServer defines the interface for a DRPC server, which can serve
+// connections and provides a drpc.Mux for registering handlers.
+type DRPCServer interface {
+	// Serve starts serving DRPC requests on the provided listener.
+	Serve(ctx context.Context, lis net.Listener) error
+	drpc.Mux
+}
+
+// drpcServer implements the DRPCServer interface.
+type drpcServer struct {
+	*drpcserver.Server
+	drpc.Mux
+}
+
+// NewDRPCServer creates a new DRPCServer with the provided rpc context.
+func NewDRPCServer(_ context.Context, rpcCtx *Context) (DRPCServer, error) {
+	d := &drpcServer{}
+	mux := drpcmux.New()
+	d.Server = 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}},
+	})
+	d.Mux = mux
+
+	// 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 d, nil
+}
+
+// NewDummyDRPCServer returns a DRPCServer that is disabled and always returns
+// ErrDRPCDisabled.
+func NewDummyDRPCServer() DRPCServer {
+	return &drpcOffServer{}
+}
+
+// drpcOffServer is a disabled DRPC server implementation. It immediately closes
+// accepted connections and returns ErrDRPCDisabled for all Serve calls.
+// Register is a no-op.
+type drpcOffServer struct{}
+
+// Serve implements the DRPCServer interface for drpcOffServer. It closes any
+// accepted connection and returns ErrDRPCDisabled.
+func (srv *drpcOffServer) Serve(_ context.Context, lis net.Listener) error {
+	conn, err := lis.Accept()
+	if err != nil {
+		return err
+	}
+	_ = conn.Close()
+	return ErrDRPCDisabled
+}
+
+// Register implements the drpc.Mux interface for drpcOffServer. It is a no-op
+// when DRPC is disabled.
+func (srv *drpcOffServer) Register(interface{}, drpc.Description) error {
+	return nil
+}

pkg/server/BUILD.bazel

@@ -375,13 +375,8 @@ 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/drpc_server.go

@@ -8,116 +8,58 @@ package server
 import (
 	"context"
 	"crypto/tls"
-	"math"
-	"net"
 
 	"github.com/cockroachdb/cockroach/pkg/rpc"
 	"github.com/cockroachdb/cockroach/pkg/server/srverrors"
-	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
 	"google.golang.org/grpc/codes"
-	"storj.io/drpc"
 	"storj.io/drpc/drpcerr"
-	"storj.io/drpc/drpcmanager"
-	"storj.io/drpc/drpcmux"
-	"storj.io/drpc/drpcserver"
-	"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")
-
-type drpcServerI interface {
-	Serve(ctx context.Context, lis net.Listener) error
-}
-
-type drpcMuxI interface {
-	Register(srv interface{}, desc drpc.Description) error
-}
-
+// drpcServer wraps a DRPCServer and manages its enabled state and TLS
+// configuration. It also tracks the current serving mode for health checks.
 type drpcServer struct {
+	// Embeds logic for managing server mode (initializing, draining, operational).
 	serveModeHandler
-	srv     drpcServerI
-	mux     drpcMuxI
-	tlsCfg  *tls.Config
+	// Underlying DRPC server implementation.
+	rpc.DRPCServer
+	// Indicates if DRPC is enabled for this server.
 	enabled bool
+	// TLS configuration for secure connections.
+	tlsCfg *tls.Config
 }
 
-var _ drpcServerI = (*drpcserver.Server)(nil)
-var _ drpcServerI = (*drpcOffServer)(nil)
-
+// newDRPCServer creates and configures a new drpcServer instance. It enables
+// DRPC if the experimental setting is on, otherwise returns a dummy server.
+//
 // TODO: Register DRPC Heartbeat service
-func newDRPCServer(_ context.Context, rpcCtx *rpc.Context) (*drpcServer, error) {
-	var dmux drpcMuxI = &drpcOffServer{}
-	var dsrv drpcServerI = &drpcOffServer{}
-	var tlsCfg *tls.Config
-	enabled := false
-
+func newDRPCServer(ctx context.Context, rpcCtx *rpc.Context) (*drpcServer, error) {
+	drpcServer := &drpcServer{}
 	if rpc.ExperimentalDRPCEnabled.Get(&rpcCtx.Settings.SV) {
-		enabled = true
-		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()
+		d, err := rpc.NewDRPCServer(ctx, rpcCtx)
 		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
-	}
-
-	d := &drpcServer{
-		srv:     dsrv,
-		mux:     dmux,
-		tlsCfg:  tlsCfg,
-		enabled: enabled,
+		drpcServer.DRPCServer = d
+		drpcServer.enabled = true
+	} else {
+		drpcServer.DRPCServer = rpc.NewDummyDRPCServer()
+		drpcServer.enabled = false
 	}
 
-	d.setMode(modeInitializing)
-
-	return d, nil
-}
-
-// 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()
+	tlsCfg, err := rpcCtx.GetServerTLSConfig()
 	if err != nil {
-		return err
+		return nil, err
 	}
-	_ = conn.Close()
-	return ErrDRPCDisabled
-}
 
-func (srv *drpcOffServer) Register(interface{}, drpc.Description) error {
-	return nil
+	drpcServer.tlsCfg = tlsCfg
+	drpcServer.setMode(modeInitializing)
+
+	return drpcServer, nil
 }
 
+// health returns an error if the server is not operational, encoding the error
+// with a DRPC code. Returns nil if the server is healthy and operational.
 func (s *drpcServer) health(ctx context.Context) error {
 	sm := s.mode.get()
 	switch sm {

pkg/server/server.go

@@ -975,7 +975,7 @@ func NewServer(cfg Config, stopper *stop.Stopper) (serverctl.ServerStartupInterf
 		cfg.LicenseEnforcer,
 	)
 	kvpb.RegisterInternalServer(grpcServer.Server, node)
-	if err := kvpb.DRPCRegisterKVBatch(drpcServer.mux, node.AsDRPCKVBatchServer()); err != nil {
+	if err := kvpb.DRPCRegisterKVBatch(drpcServer, node.AsDRPCKVBatchServer()); err != nil {
 		return nil, err
 	}
 	kvserver.RegisterPerReplicaServer(grpcServer.Server, node.perReplicaServer)

pkg/server/start_listen.go

@@ -207,9 +207,9 @@ func startListenRPCAndSQL(
 		_ = stopper.RunAsyncTask(drpcCtx, "serve-drpc", func(ctx context.Context) {
 			if cfg := drpc.tlsCfg; cfg != nil {
 				drpcTLSL := tls.NewListener(drpcL, cfg)
-				netutil.FatalIfUnexpected(drpc.srv.Serve(ctx, drpcTLSL))
+				netutil.FatalIfUnexpected(drpc.Serve(ctx, drpcTLSL))
 			} else {
-				netutil.FatalIfUnexpected(drpc.srv.Serve(ctx, drpcL))
+				netutil.FatalIfUnexpected(drpc.Serve(ctx, drpcL))
 			}
 		})
 		_ = stopper.RunAsyncTask(workersCtx, "serve-loopback-grpc", func(context.Context) {