server: refactor the grpc server

This patch creates the grpcServer object, to encapsulate the grpc.Server
object, the serveMode object and the interceptor that our server.Server
uses to filter out RPCs.
The idea is to work towards decoupling the gprc.Server from our
server.Server object. I'd like to have the grpc server be created before
Server.Start(): I'd like the cluster id and node id to be available by
the time Server.Start() is called so that we can get rid of the
nodeIDCountainer that everybody and their dog uses. But for getting
these ids a grpc server needs to be running early to handle the "init"
rpc which bootstraps a cluster.
This patch doesn't accomplish too much - it doesn't do anything about
actually starting to serve any requests without a Server (i.e. create
the needed listeners), but I think the patch stands on its own too as
good refactoring.

Release note: None

pkg/server/servemode.go → pkg/server/grpc_server.go

@@ -18,10 +18,31 @@ import (
 	"strings"
 	"sync/atomic"
 
+	"github.com/cockroachdb/cockroach/pkg/rpc"
+	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	grpcstatus "google.golang.org/grpc/status"
 )
 
+// grpcServer is a wrapper on top of a grpc.Server that includes an interceptor
+// and a mode of operation that can instruct the interceptor to refuse certain
+// RPCs.
+type grpcServer struct {
+	*grpc.Server
+	mode serveMode
+}
+
+func newGRPCServer(rpcCtx *rpc.Context) *grpcServer {
+	s := &grpcServer{}
+	s.mode.set(modeInitializing)
+	s.Server = rpc.NewServerWithInterceptor(rpcCtx, func(path string) error {
+		return s.intercept(path)
+	})
+	return s
+}
+
+type serveMode int32
+
 // A list of the server states for bootstrap process.
 const (
 	// modeInitializing is intended for server initialization process.
@@ -36,25 +57,31 @@ const (
 	modeDraining
 )
 
-type serveMode int32
+func (s *grpcServer) setMode(mode serveMode) {
+	s.mode.set(mode)
+}
 
-// Intercept implements filtering rules for each server state.
-func (s *Server) Intercept() func(string) error {
-	interceptors := map[string]struct{}{
-		"/cockroach.rpc.Heartbeat/Ping":             {},
-		"/cockroach.gossip.Gossip/Gossip":           {},
-		"/cockroach.server.serverpb.Init/Bootstrap": {},
-		"/cockroach.server.serverpb.Status/Details": {},
-	}
-	return func(fullName string) error {
-		if s.serveMode.operational() {
-			return nil
-		}
-		if _, allowed := interceptors[fullName]; !allowed {
-			return WaitingForInitError(fullName)
-		}
+func (s *grpcServer) operational() bool {
+	sMode := s.mode.get()
+	return sMode == modeOperational || sMode == modeDraining
+}
+
+var rpcsAllowedWhileBootstrapping = map[string]struct{}{
+	"/cockroach.rpc.Heartbeat/Ping":             {},
+	"/cockroach.gossip.Gossip/Gossip":           {},
+	"/cockroach.server.serverpb.Init/Bootstrap": {},
+	"/cockroach.server.serverpb.Status/Details": {},
+}
+
+// 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 s.waitingForInitError(fullName)
+	}
+	return nil
 }
 
 func (s *serveMode) set(mode serveMode) {
@@ -65,14 +92,9 @@ func (s *serveMode) get() serveMode {
 	return serveMode(atomic.LoadInt32((*int32)(s)))
 }
 
-func (s *serveMode) operational() bool {
-	sMode := s.get()
-	return sMode == modeOperational || sMode == modeDraining
-}
-
-// WaitingForInitError indicates that the server cannot run the specified
-// method until the node has been initialized.
-func WaitingForInitError(methodName string) error {
+// waitingForInitError creates an error indicating that the server cannot run
+// the specified method until the node has been initialized.
+func (s *grpcServer) waitingForInitError(methodName string) error {
 	return grpcstatus.Errorf(codes.Unavailable, "node waiting for init; %s not available", methodName)
 }
 

pkg/server/servemode_test.go

@@ -25,7 +25,8 @@ import (
 
 func TestWaitingForInitError(t *testing.T) {
 	defer leaktest.AfterTest(t)()
-	if err := WaitingForInitError("foo"); !IsWaitingForInit(err) {
+	s := &grpcServer{}
+	if err := s.waitingForInitError("foo"); !IsWaitingForInit(err) {
 		t.Errorf("WaitingForInitError() not recognized by IsWaitingForInit(): %v", err)
 	}
 	if err := grpcstatus.Errorf(codes.Unavailable, "foo"); IsWaitingForInit(err) {

pkg/server/server.go

@@ -149,12 +149,13 @@ func (mux *safeServeMux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 type Server struct {
 	nodeIDContainer base.NodeIDContainer
 
-	cfg              Config
-	st               *cluster.Settings
-	mux              safeServeMux
-	clock            *hlc.Clock
-	rpcContext       *rpc.Context
-	grpc             *grpc.Server
+	cfg        Config
+	st         *cluster.Settings
+	mux        safeServeMux
+	clock      *hlc.Clock
+	rpcContext *rpc.Context
+	// The gRPC server on which the different RPC handlers will be registered.
+	grpc             *grpcServer
 	gossip           *gossip.Gossip
 	nodeDialer       *nodedialer.Dialer
 	nodeLiveness     *storage.NodeLiveness
@@ -189,7 +190,6 @@ type Server struct {
 	adminMemMetrics    sql.MemoryMetrics
 	// sqlMemMetrics are used to track memory usage of sql sessions.
 	sqlMemMetrics sql.MemoryMetrics
-	serveMode
 }
 
 // NewServer creates a Server from a server.Config.
@@ -212,7 +212,6 @@ func NewServer(cfg Config, stopper *stop.Stopper) (*Server, error) {
 		cfg:      cfg,
 		registry: metric.NewRegistry(),
 	}
-	s.serveMode.set(modeInitializing)
 
 	// If the tracer has a Close function, call it after the server stops.
 	if tr, ok := cfg.AmbientCtx.Tracer.(stop.Closer); ok {
@@ -251,14 +250,14 @@ func NewServer(cfg Config, stopper *stop.Stopper) (*Server, error) {
 		}
 	}
 
-	s.grpc = rpc.NewServerWithInterceptor(s.rpcContext, s.Intercept())
+	s.grpc = newGRPCServer(s.rpcContext)
 
 	s.gossip = gossip.New(
 		s.cfg.AmbientCtx,
 		&s.rpcContext.ClusterID,
 		&s.nodeIDContainer,
 		s.rpcContext,
-		s.grpc,
+		s.grpc.Server,
 		s.stopper,
 		s.registry,
 		s.cfg.Locality,
@@ -340,7 +339,7 @@ func NewServer(cfg Config, stopper *stop.Stopper) (*Server, error) {
 	)
 
 	s.raftTransport = storage.NewRaftTransport(
-		s.cfg.AmbientCtx, st, s.nodeDialer, s.grpc, s.stopper,
+		s.cfg.AmbientCtx, st, s.nodeDialer, s.grpc.Server, s.stopper,
 	)
 
 	// Set up internal memory metrics for use by internal SQL executors.
@@ -481,9 +480,9 @@ func NewServer(cfg Config, stopper *stop.Stopper) (*Server, error) {
 	s.node = NewNode(
 		storeCfg, s.recorder, s.registry, s.stopper,
 		txnMetrics, nil /* execCfg */, &s.rpcContext.ClusterID)
-	roachpb.RegisterInternalServer(s.grpc, s.node)
-	storage.RegisterPerReplicaServer(s.grpc, s.node.perReplicaServer)
-	s.node.storeCfg.ClosedTimestamp.RegisterClosedTimestampServer(s.grpc)
+	roachpb.RegisterInternalServer(s.grpc.Server, s.node)
+	storage.RegisterPerReplicaServer(s.grpc.Server, s.node.perReplicaServer)
+	s.node.storeCfg.ClosedTimestamp.RegisterClosedTimestampServer(s.grpc.Server)
 
 	s.sessionRegistry = sql.NewSessionRegistry()
 	s.jobRegistry = jobs.MakeRegistry(
@@ -535,7 +534,7 @@ func NewServer(cfg Config, stopper *stop.Stopper) (*Server, error) {
 	}
 
 	s.distSQLServer = distsqlrun.NewServer(ctx, distSQLCfg)
-	distsqlpb.RegisterDistSQLServer(s.grpc, s.distSQLServer)
+	distsqlpb.RegisterDistSQLServer(s.grpc.Server, s.distSQLServer)
 
 	s.admin = newAdminServer(s)
 	s.status = newStatusServer(
@@ -555,7 +554,7 @@ func NewServer(cfg Config, stopper *stop.Stopper) (*Server, error) {
 	)
 	s.authentication = newAuthenticationServer(s)
 	for _, gw := range []grpcGatewayServer{s.admin, s.status, s.authentication, &s.tsServer} {
-		gw.RegisterService(s.grpc)
+		gw.RegisterService(s.grpc.Server)
 	}
 
 	// TODO(andrei): We're creating an initServer even through the inspection of
@@ -565,7 +564,7 @@ func NewServer(cfg Config, stopper *stop.Stopper) (*Server, error) {
 	// figure out early if our engines are bootstrapped and, if they are, create a
 	// dummy implementation of the InitServer that rejects all RPCs.
 	s.initServer = newInitServer(s.gossip.Connected, s.stopper.ShouldStop())
-	serverpb.RegisterInitServer(s.grpc, s.initServer)
+	serverpb.RegisterInitServer(s.grpc.Server, s.initServer)
 
 	nodeInfo := sql.NodeInfo{
 		AdminURL:  cfg.AdminURL,
@@ -1538,7 +1537,7 @@ func (s *Server) Start(ctx context.Context) error {
 	s.distSQLServer.Start()
 	s.pgServer.Start(ctx, s.stopper)
 
-	s.serveMode.set(modeOperational)
+	s.grpc.setMode(modeOperational)
 
 	log.Infof(ctx, "starting %s server at %s (use: %s)",
 		s.cfg.HTTPRequestScheme(), s.cfg.HTTPAddr, s.cfg.HTTPAdvertiseAddr)
@@ -1758,15 +1757,15 @@ func (s *Server) doDrain(
 		switch mode {
 		case serverpb.DrainMode_CLIENT:
 			if setTo {
-				s.serveMode.set(modeDraining)
+				s.grpc.setMode(modeDraining)
 				// Wait for drainUnreadyWait. This will fail load balancer checks and
 				// delay draining so that client traffic can move off this node.
 				time.Sleep(drainWait.Get(&s.st.SV))
 			}
 			if err := func() error {
 				if !setTo {
 					// Execute this last.
-					defer func() { s.serveMode.set(modeOperational) }()
+					defer func() { s.grpc.setMode(modeOperational) }()
 				}
 				// Since enabling the lease manager's draining mode will prevent
 				// the acquisition of new leases, the switch must be made after

pkg/server/status.go

@@ -590,7 +590,7 @@ func (s *statusServer) Details(
 		return resp, nil
 	}
 
-	serveMode := s.admin.server.serveMode.get()
+	serveMode := s.admin.server.grpc.mode.get()
 	if serveMode != modeOperational {
 		return nil, grpcstatus.Error(codes.Unavailable, "node is not ready")
 	}