Merge pull request #154224 from cockroachdb/blathers/backport-release-25.4-154123

Author: Nukitt

pkg/server/drpc_server.go

@@ -31,7 +31,17 @@ type drpcServer struct {
 // newDRPCServer creates and configures a new drpcServer instance. It enables
 // DRPC if the experimental setting is on, otherwise returns a dummy server.
 func newDRPCServer(ctx context.Context, rpcCtx *rpc.Context) (*drpcServer, error) {
-	d, err := rpc.NewDRPCServer(ctx, rpcCtx)
+	drpcServer := &drpcServer{}
+	drpcServer.setMode(modeInitializing)
+
+	d, err := rpc.NewDRPCServer(
+		ctx,
+		rpcCtx,
+		rpc.WithInterceptor(
+			func(path string) error {
+				return drpcServer.intercept(path)
+			}),
+	)
 	if err != nil {
 		return nil, err
 	}
@@ -41,12 +51,8 @@ func newDRPCServer(ctx context.Context, rpcCtx *rpc.Context) (*drpcServer, error
 		return nil, err
 	}
 
-	drpcServer := &drpcServer{
-		DRPCServer: d,
-		tlsCfg:     tlsCfg,
-	}
-
-	drpcServer.setMode(modeInitializing)
+	drpcServer.DRPCServer = d
+	drpcServer.tlsCfg = tlsCfg
 
 	if err := rpc.DRPCRegisterHeartbeat(drpcServer, rpcCtx.NewHeartbeatService()); err != nil {
 		return nil, err

pkg/server/grpc_server.go

@@ -68,24 +68,6 @@ func (s *grpcServer) health(ctx context.Context) error {
 	}
 }
 
-var rpcsAllowedWhileBootstrapping = map[string]struct{}{
-	"/cockroach.rpc.Heartbeat/Ping":             {},
-	"/cockroach.gossip.Gossip/Gossip":           {},
-	"/cockroach.server.serverpb.Init/Bootstrap": {},
-	"/cockroach.server.serverpb.Admin/Health":   {},
-}
-
-// intercept implements filtering rules for each server state.
-func (s *grpcServer) intercept(fullName string) error {
-	if s.operational() {
-		return nil
-	}
-	if _, allowed := rpcsAllowedWhileBootstrapping[fullName]; !allowed {
-		return NewWaitingForInitError(fullName)
-	}
-	return nil
-}
-
 // NewWaitingForInitError creates an error indicating that the server cannot run
 // the specified method until the node has been initialized.
 func NewWaitingForInitError(methodName string) error {

pkg/server/serve_mode.go

@@ -36,6 +36,24 @@ const (
 	modeDraining
 )
 
+var rpcsAllowedWhileBootstrapping = map[string]struct{}{
+	"/cockroach.rpc.Heartbeat/Ping":             {},
+	"/cockroach.gossip.Gossip/Gossip":           {},
+	"/cockroach.server.serverpb.Init/Bootstrap": {},
+	"/cockroach.server.serverpb.Admin/Health":   {},
+}
+
+// intercept implements filtering rules for each server state.
+func (s *serveModeHandler) intercept(fullName string) error {
+	if s.operational() {
+		return nil
+	}
+	if _, allowed := rpcsAllowedWhileBootstrapping[fullName]; !allowed {
+		return NewWaitingForInitError(fullName)
+	}
+	return nil
+}
+
 func (s *serveMode) set(mode serveMode) {
 	atomic.StoreInt32((*int32)(s), int32(mode))
 }