Merge #148584

148584: server: register status, TS, Login/Logout, Admin  services with DRPC server r=cthumuluru-crdb,aa-joshi a=shubhamdhama

Enable the `Status`, `Login`, `Logout`, `Timeseries`, `Admin` services on the DRPC server in addition to gRPC. This is
controlled by `rpc.experimental_drpc.enabled` (off by default).

This change is part of a series and is similar to: #146926

Note: This only registers the service; the client is not updated to use the
DRPC client, so this service will not have any functional effect.

Fixes: #148721
Fixes: #148722
Fixes: #148724
Fixes: #148723

Epic: CRDB-48925
Release note: None


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

Author: craig[bot]

pkg/kv/kvserver/loqrecovery/server.go

@@ -119,7 +119,7 @@ func NewServer(
 func (s Server) ServeLocalReplicas(
 	ctx context.Context,
 	_ *serverpb.RecoveryCollectLocalReplicaInfoRequest,
-	stream serverpb.Admin_RecoveryCollectLocalReplicaInfoServer,
+	stream serverpb.RPCAdmin_RecoveryCollectLocalReplicaInfoStream,
 ) error {
 	v := s.settings.Version.ActiveVersion(ctx)
 	var stores []*kvserver.Store
@@ -149,7 +149,7 @@ func (s Server) ServeLocalReplicas(
 func (s Server) ServeClusterReplicas(
 	ctx context.Context,
 	req *serverpb.RecoveryCollectReplicaInfoRequest,
-	outStream serverpb.Admin_RecoveryCollectReplicaInfoServer,
+	outStream serverpb.RPCAdmin_RecoveryCollectReplicaInfoStream,
 	kvDB *kv.DB,
 ) (err error) {
 	var descriptors, nodes, replicas atomic.Int64

pkg/server/admin.go

@@ -82,6 +82,8 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	grpcstatus "google.golang.org/grpc/status"
+	"storj.io/drpc"
+	"storj.io/drpc/drpcerr"
 )
 
 // Number of empty ranges for table descriptors that aren't actually tables. These
@@ -253,6 +255,26 @@ func (s *adminServer) RegisterService(g *grpc.Server) {
 	serverpb.RegisterAdminServer(g, s)
 }
 
+type drpcSystemAdminServer struct {
+	*systemAdminServer
+}
+
+// RegisterDRPCService registers the Admin service with the DRPC server running
+// in the system tenant.
+func (s *systemAdminServer) RegisterDRPCService(d drpc.Mux) error {
+	return serverpb.DRPCRegisterAdmin(d, &drpcSystemAdminServer{systemAdminServer: s})
+}
+
+type drpcAdminServer struct {
+	*adminServer
+}
+
+// RegisterDRPCService registers the Admin service with the DRPC server running
+// in a secondary tenant.
+func (s *adminServer) RegisterDRPCService(d drpc.Mux) error {
+	return serverpb.DRPCRegisterAdmin(d, &drpcAdminServer{adminServer: s})
+}
+
 // RegisterGateway starts the gateway (i.e. reverse proxy) that proxies HTTP requests
 // to the appropriate gRPC endpoints.
 func (s *adminServer) RegisterGateway(
@@ -3312,9 +3334,23 @@ func (s *systemAdminServer) SendKVBatch(
 	return br, nil
 }
 
+func (s *drpcSystemAdminServer) RecoveryCollectReplicaInfo(
+	request *serverpb.RecoveryCollectReplicaInfoRequest,
+	stream serverpb.DRPCAdmin_RecoveryCollectReplicaInfoStream,
+) error {
+	return s.recoveryCollectReplicaInfo(request, stream)
+}
+
 func (s *systemAdminServer) RecoveryCollectReplicaInfo(
 	request *serverpb.RecoveryCollectReplicaInfoRequest,
 	stream serverpb.Admin_RecoveryCollectReplicaInfoServer,
+) error {
+	return s.recoveryCollectReplicaInfo(request, stream)
+}
+
+func (s *systemAdminServer) recoveryCollectReplicaInfo(
+	request *serverpb.RecoveryCollectReplicaInfoRequest,
+	stream serverpb.RPCAdmin_RecoveryCollectReplicaInfoStream,
 ) error {
 	ctx := stream.Context()
 	ctx = s.server.AnnotateCtx(ctx)
@@ -3327,9 +3363,23 @@ func (s *systemAdminServer) RecoveryCollectReplicaInfo(
 	return s.server.recoveryServer.ServeClusterReplicas(ctx, request, stream, s.server.db)
 }
 
+func (s *drpcSystemAdminServer) RecoveryCollectLocalReplicaInfo(
+	request *serverpb.RecoveryCollectLocalReplicaInfoRequest,
+	stream serverpb.DRPCAdmin_RecoveryCollectLocalReplicaInfoStream,
+) error {
+	return s.recoveryCollectLocalReplicaInfo(request, stream)
+}
+
 func (s *systemAdminServer) RecoveryCollectLocalReplicaInfo(
 	request *serverpb.RecoveryCollectLocalReplicaInfoRequest,
 	stream serverpb.Admin_RecoveryCollectLocalReplicaInfoServer,
+) error {
+	return s.recoveryCollectLocalReplicaInfo(request, stream)
+}
+
+func (s *systemAdminServer) recoveryCollectLocalReplicaInfo(
+	request *serverpb.RecoveryCollectLocalReplicaInfoRequest,
+	stream serverpb.RPCAdmin_RecoveryCollectLocalReplicaInfoStream,
 ) error {
 	ctx := stream.Context()
 	ctx = s.server.AnnotateCtx(ctx)
@@ -4043,3 +4093,21 @@ func (s *systemAdminServer) ReadFromTenantInfo(
 
 	return &serverpb.ReadFromTenantInfoResponse{ReadFrom: dstID, ReadAt: progress.GetStreamIngest().ReplicatedTime}, nil
 }
+
+// RecoveryCollectReplicaInfo is unimplemented here because adminServer also
+// have it delegated from embedded serverpb.UnimplementedAdminServer.
+func (s *drpcAdminServer) RecoveryCollectReplicaInfo(
+	request *serverpb.RecoveryCollectReplicaInfoRequest,
+	stream serverpb.DRPCAdmin_RecoveryCollectReplicaInfoStream,
+) error {
+	return drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented)
+}
+
+// RecoveryCollectLocalReplicaInfo is unimplemented here because adminServer
+// also have it delegated from embedded serverpb.UnimplementedAdminServer.
+func (s *drpcAdminServer) RecoveryCollectLocalReplicaInfo(
+	request *serverpb.RecoveryCollectLocalReplicaInfoRequest,
+	stream serverpb.DRPCAdmin_RecoveryCollectLocalReplicaInfoStream,
+) error {
+	return drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented)
+}

pkg/server/authserver/BUILD.bazel

@@ -44,6 +44,7 @@ go_library(
         "@com_github_cockroachdb_logtags//:logtags",
         "@com_github_cockroachdb_redact//:redact",
         "@com_github_grpc_ecosystem_grpc_gateway//runtime:go_default_library",
+        "@io_storj_drpc//:drpc",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//metadata",

pkg/server/authserver/api.go

@@ -18,12 +18,15 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire"
 	"github.com/grpc-ecosystem/grpc-gateway/runtime"
 	"google.golang.org/grpc"
+	"storj.io/drpc"
 )
 
 type Server interface {
 	RegisterService(*grpc.Server)
 	RegisterGateway(ctx context.Context, mux *runtime.ServeMux, conn *grpc.ClientConn) error
 
+	RegisterDRPCService(drpc.Mux) error
+
 	// UserLogin verifies an incoming request by a user to create an web
 	// authentication session. It checks the provided credentials against
 	// the system.users table, and if successful creates a new

pkg/server/authserver/authentication.go

@@ -47,6 +47,7 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
+	"storj.io/drpc"
 )
 
 const (
@@ -132,6 +133,15 @@ type authenticationServer struct {
 func (s *authenticationServer) RegisterService(g *grpc.Server) {
 	serverpb.RegisterLogInServer(g, s)
 	serverpb.RegisterLogOutServer(g, s)
+
+}
+
+// RegisterService registers the LogIn and LogOut services with DRPC.
+func (s *authenticationServer) RegisterDRPCService(d drpc.Mux) error {
+	if err := serverpb.DRPCRegisterLogIn(d, s); err != nil {
+		return err
+	}
+	return serverpb.DRPCRegisterLogOut(d, s)
 }
 
 // RegisterGateway starts the gateway (i.e. reverse proxy) that proxies HTTP requests

pkg/server/drain.go

@@ -88,6 +88,12 @@ var (
 // instructs the process to terminate.
 // This method is part of the serverpb.AdminClient interface.
 func (s *adminServer) Drain(req *serverpb.DrainRequest, stream serverpb.Admin_DrainServer) error {
+	return s.drain(req, stream)
+}
+
+func (s *adminServer) drain(
+	req *serverpb.DrainRequest, stream serverpb.RPCAdmin_DrainStream,
+) error {
 	ctx := stream.Context()
 	ctx = s.AnnotateCtx(ctx)
 
@@ -114,6 +120,22 @@ func (s *adminServer) Drain(req *serverpb.DrainRequest, stream serverpb.Admin_Dr
 	return s.drainServer.handleDrain(ctx, req, stream)
 }
 
+// Drain puts the node into the specified drain mode(s) and optionally
+// instructs the process to terminate.
+func (s *drpcAdminServer) Drain(
+	req *serverpb.DrainRequest, stream serverpb.DRPCAdmin_DrainStream,
+) error {
+	return s.adminServer.drain(req, stream)
+}
+
+// Drain puts the node into the specified drain mode(s) and optionally
+// instructs the process to terminate.
+func (s *drpcSystemAdminServer) Drain(
+	req *serverpb.DrainRequest, stream serverpb.DRPCAdmin_DrainStream,
+) error {
+	return s.adminServer.drain(req, stream)
+}
+
 type drainServer struct {
 	stopper *stop.Stopper
 	// stopTrigger is used to request that the server is shut down.
@@ -162,7 +184,7 @@ func (s *drainServer) setNode(node *Node, nodeLiveness *liveness.NodeLiveness) {
 }
 
 func (s *drainServer) handleDrain(
-	ctx context.Context, req *serverpb.DrainRequest, stream serverpb.Admin_DrainServer,
+	ctx context.Context, req *serverpb.DrainRequest, stream serverpb.RPCAdmin_DrainStream,
 ) error {
 	log.Ops.Infof(ctx, "drain request received with doDrain = %v, shutdown = %v", req.DoDrain, req.Shutdown)
 
@@ -240,7 +262,7 @@ func delegateDrain(
 	ctx context.Context,
 	req *serverpb.DrainRequest,
 	client serverpb.RPCAdminClient,
-	stream serverpb.Admin_DrainServer,
+	stream serverpb.RPCAdmin_DrainStream,
 ) error {
 	// Retrieve the stream interface to the target node.
 	drainClient, err := client.Drain(ctx, req)

pkg/server/drpc_server.go

@@ -13,6 +13,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/server/srverrors"
 	"github.com/cockroachdb/errors"
 	"google.golang.org/grpc/codes"
+	"storj.io/drpc"
 	"storj.io/drpc/drpcerr"
 )
 
@@ -75,3 +76,9 @@ func (s *drpcServer) health(ctx context.Context) error {
 		return srverrors.ServerError(ctx, errors.Newf("unknown mode: %v", sm))
 	}
 }
+
+// drpcServiceRegistrar is implemented by servers that create a DRPC server and
+// registers it with drpc.Mux
+type drpcServiceRegistrar interface {
+	RegisterDRPCService(drpc.Mux) error
+}

pkg/server/server.go

@@ -1253,6 +1253,12 @@ func NewServer(cfg Config, stopper *stop.Stopper) (serverctl.ServerStartupInterf
 		gw.RegisterService(grpcServer.Server)
 	}
 
+	for _, s := range []drpcServiceRegistrar{sAdmin, sStatus, sAuth, &sTS} {
+		if err := s.RegisterDRPCService(drpcServer); err != nil {
+			return nil, err
+		}
+	}
+
 	// Tell the node event logger (join, restart) how to populate SQL entries
 	// into system.eventlog.
 	node.InitLogger(sqlServer.execCfg)

pkg/server/status.go

@@ -95,6 +95,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"storj.io/drpc"
 )
 
 const (
@@ -695,6 +696,11 @@ func (s *statusServer) RegisterService(g *grpc.Server) {
 	serverpb.RegisterStatusServer(g, s)
 }
 
+// RegisterService registers the DRPC service.
+func (s *statusServer) RegisterDRPCService(d drpc.Mux) error {
+	return serverpb.DRPCRegisterStatus(d, s)
+}
+
 // RegisterGateway starts the gateway (i.e. reverse
 // proxy) that proxies HTTP requests to the appropriate gRPC endpoints.
 func (s *statusServer) RegisterGateway(
@@ -709,6 +715,11 @@ func (s *systemStatusServer) RegisterService(g *grpc.Server) {
 	serverpb.RegisterStatusServer(g, s)
 }
 
+// RegisterService registers the DRPC service.
+func (s *systemStatusServer) RegisterDRPCService(d drpc.Mux) error {
+	return serverpb.DRPCRegisterStatus(d, s)
+}
+
 func (s *statusServer) parseNodeID(nodeIDParam string) (roachpb.NodeID, bool, error) {
 	id, local, err := s.serverIterator.parseServerID(nodeIDParam)
 	return roachpb.NodeID(id), local, err

pkg/server/tenant.go

@@ -480,6 +480,12 @@ func newTenantServer(
 		gw.RegisterService(args.grpc.Server)
 	}
 
+	for _, s := range []drpcServiceRegistrar{sAdmin, sStatus, sAuth, args.tenantTimeSeriesServer} {
+		if err := s.RegisterDRPCService(args.drpc); err != nil {
+			return nil, err
+		}
+	}
+
 	// Tell the status/admin servers how to access SQL structures.
 	sStatus.setStmtDiagnosticsRequester(sqlServer.execCfg.StmtDiagnosticsRecorder)
 	serverIterator.sqlServer = sqlServer