server: register Admin service with DRPC server

Enable these service 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.

Epic: CRDB-48925
Release note: None

Author: shubhamdhama

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/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/server.go

@@ -1259,6 +1259,9 @@ func NewServer(cfg Config, stopper *stop.Stopper) (serverctl.ServerStartupInterf
 	if err := sTS.RegisterDRPCService(drpcServer); err != nil {
 		return nil, err
 	}
+	if err := sAdmin.RegisterDRPCService(drpcServer); err != nil {
+		return nil, err
+	}
 
 	// Tell the node event logger (join, restart) how to populate SQL entries
 	// into system.eventlog.

pkg/server/tenant.go

@@ -486,6 +486,9 @@ func newTenantServer(
 	if err := args.tenantTimeSeriesServer.RegisterDRPCService(args.drpc); err != nil {
 		return nil, err
 	}
+	if err := sAdmin.RegisterDRPCService(args.drpc); err != nil {
+		return nil, err
+	}
 
 	// Tell the status/admin servers how to access SQL structures.
 	sStatus.setStmtDiagnosticsRequester(sqlServer.execCfg.StmtDiagnosticsRecorder)