[management] batch disconnect and lock free reconnect (#8373)

management/internals/server/boot.go

@@ -269,6 +269,9 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
 			grpc.KeepaliveParams(kasp),
 			grpc.ChainUnaryInterceptor(realip.UnaryServerInterceptorOpts(realipOpts...), unaryInterceptor, proxyUnary),
 			grpc.ChainStreamInterceptor(realip.StreamServerInterceptorOpts(realipOpts...), streamInterceptor, proxyStream),
+			// Stop waits for the sync handlers, so their teardown has queued every
+			// disconnect before FlushDisconnects writes them.
+			grpc.WaitForHandlers(true),
 		}
 
 		// Append interceptors contributed by registered gRPC extensions. These
@@ -292,6 +295,21 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
 		}
 
 		gRPCAPIHandler := grpc.NewServer(gRPCOpts...)
+		mgmtProto.RegisterManagementServiceServer(gRPCAPIHandler, s.ManagementGRPCServer())
+
+		mgmtProto.RegisterProxyServiceServer(gRPCAPIHandler, s.ReverseProxyGRPCServer())
+		log.Info("ProxyService registered on gRPC server")
+
+		// Register services contributed by external modules via the extension seam.
+		registerExtensions(gRPCAPIHandler, s.grpcExtensions)
+
+		return gRPCAPIHandler
+	})
+}
+
+// ManagementGRPCServer returns the management gRPC service, created once and served by GRPCServer.
+func (s *BaseServer) ManagementGRPCServer() *nbgrpc.Server {
+	return Create(s, func() *nbgrpc.Server {
 		srv, err := nbgrpc.NewServer(s.Config, s.AccountManager(), s.SettingsManager(), s.JobManager(), s.SecretsManager(), s.Metrics(), s.AuthManager(), s.IntegratedValidator(), s.NetworkMapController(), s.OAuthConfigProvider(), s.SessionStore())
 		if err != nil {
 			log.Fatalf("failed to create management server: %v", err)
@@ -301,15 +319,7 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
 		if serviceMgr != nil {
 			serviceMgr.StartExposeReaper(context.Background())
 		}
-		mgmtProto.RegisterManagementServiceServer(gRPCAPIHandler, srv)
-
-		mgmtProto.RegisterProxyServiceServer(gRPCAPIHandler, s.ReverseProxyGRPCServer())
-		log.Info("ProxyService registered on gRPC server")
-
-		// Register services contributed by external modules via the extension seam.
-		registerExtensions(gRPCAPIHandler, s.grpcExtensions)
-
-		return gRPCAPIHandler
+		return srv
 	})
 }
 

management/internals/server/container.go

@@ -1,6 +1,9 @@
 package server
 
-import "fmt"
+import (
+	"fmt"
+	"reflect"
+)
 
 // Create a dependency and add it to the BaseServer's container. A string key identifier will be based on its type definition.
 func Create[T any](s Server, createFunc func() T) T {
@@ -33,15 +36,41 @@ func InjectNamed[T any](c Server, name string, thing T) {
 }
 
 func maybeCreate[T any](s Server, createFunc func() T) (result T, isNew bool) {
-	key := fmt.Sprintf("%T", (*T)(nil))[1:]
-	return maybeCreateKeyed(s, key, createFunc)
+	return maybeCreateKeyed(s, typeKey(reflect.TypeFor[T]()), createFunc)
 }
 
 func maybeCreateNamed[T any](s Server, name string, createFunc func() T) (result T, isNew bool) {
-	key := fmt.Sprintf("%T:%s", (*T)(nil), name)[1:]
+	key := fmt.Sprintf("%s:%s", typeKey(reflect.TypeFor[T]()), name)
 	return maybeCreateKeyed(s, key, createFunc)
 }
 
+// typeKey names t by the full package path of every named type it is built from, so
+// same-named types of different packages, such as google.golang.org/grpc.Server and a
+// local grpc.Server, never share a container key.
+func typeKey(t reflect.Type) string {
+	if t.Name() != "" {
+		if t.PkgPath() == "" {
+			return t.Name()
+		}
+		return t.PkgPath() + "." + t.Name()
+	}
+
+	switch t.Kind() {
+	case reflect.Pointer:
+		return "*" + typeKey(t.Elem())
+	case reflect.Slice:
+		return "[]" + typeKey(t.Elem())
+	case reflect.Array:
+		return fmt.Sprintf("[%d]%s", t.Len(), typeKey(t.Elem()))
+	case reflect.Map:
+		return "map[" + typeKey(t.Key()) + "]" + typeKey(t.Elem())
+	case reflect.Chan:
+		return t.ChanDir().String() + " " + typeKey(t.Elem())
+	default:
+		return t.String()
+	}
+}
+
 func maybeCreateKeyed[T any](s Server, key string, createFunc func() T) (result T, isNew bool) {
 	if t, ok := s.GetContainer(key); ok {
 		if t == nil {

management/internals/server/server.go

@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"crypto/tls"
+	"errors"
 	"fmt"
 	"net"
 	"net/http"
@@ -40,6 +41,15 @@ const (
 	ContainerKeyBaseServer = "baseServer"
 
 	applicationName = "management"
+
+	// shutdownTimeout stays below the 10s that systemd (TimeoutStopSec) and Docker grant
+	// before killing the process.
+	shutdownTimeout = 9 * time.Second
+	grpcStopTimeout = 3 * time.Second
+	// shutdownCloseReserve is kept from the flush for closing the stores.
+	shutdownCloseReserve = 2 * time.Second
+	// shutdownIdPReserve is kept from the extension shutdown hooks for stopping the embedded IdP.
+	shutdownIdPReserve = 1 * time.Second
 )
 
 type Server interface {
@@ -297,9 +307,13 @@ func (s *BaseServer) setupTLS(ctx context.Context) (bool, error) {
 	}
 }
 
-// Stop attempts a graceful shutdown, waiting up to 5 seconds for active connections to finish
+// Stop attempts a graceful shutdown within shutdownTimeout. Writing the disconnects of the
+// drained streams may use the budget left after stopping the services and gRPC streams,
+// except for shutdownCloseReserve, which stays for closing the stores; the extension hooks
+// leave shutdownIdPReserve for stopping the embedded IdP. It returns an error when the
+// disconnects could not all be written or the shutdown ran out of time.
 func (s *BaseServer) Stop() error {
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
 	defer cancel()
 	defer s.RunStopHandlers()
 	if s.domainCleanupStop != nil {
@@ -318,12 +332,20 @@ func (s *BaseServer) Stop() error {
 	if s.certListener != nil {
 		_ = s.certListener.Close()
 	}
-	s.GRPCServer().Stop()
+	s.ManagementGRPCServer().StartDraining()
+	s.stopGRPCServer()
+	deadline, _ := ctx.Deadline()
+	flushCtx, cancelFlush := context.WithDeadline(ctx, deadline.Add(-shutdownCloseReserve))
+	flushErr := s.ManagementGRPCServer().FlushDisconnects(flushCtx)
+	cancelFlush()
+
 	if s.proxyAuthClose != nil {
 		s.proxyAuthClose()
 		s.proxyAuthClose = nil
 	}
-	runExtensionShutdownHooks(ctx, s.grpcExtensions)
+	hooksCtx, cancelHooks := context.WithDeadline(ctx, deadline.Add(-shutdownIdPReserve))
+	runExtensionShutdownHooks(hooksCtx, s.grpcExtensions)
+	cancelHooks()
 	_ = s.Store().Close(ctx)
 	_ = s.EventStore().Close(ctx)
 	if s.update != nil {
@@ -338,12 +360,29 @@ func (s *BaseServer) Stop() error {
 	case err := <-s.Errors():
 		if err != nil {
 			log.WithContext(ctx).Infof("stopped Management Service: %v", err)
-			return nil
+			return flushErr
 		}
 		log.WithContext(ctx).Infof("stopped Management Service")
-		return nil
+		return flushErr
 	case <-ctx.Done():
-		return ctx.Err()
+		return errors.Join(flushErr, ctx.Err())
+	}
+}
+
+// stopGRPCServer stops the gRPC server, which waits for the stream handlers to finish
+// their teardown, but gives up waiting after grpcStopTimeout so a stuck handler cannot
+// block the shutdown.
+func (s *BaseServer) stopGRPCServer() {
+	stopped := make(chan struct{})
+	go func() {
+		s.GRPCServer().Stop()
+		close(stopped)
+	}()
+
+	select {
+	case <-stopped:
+	case <-time.After(grpcStopTimeout):
+		log.Warnf("gRPC server did not stop within %s, continuing shutdown", grpcStopTimeout)
 	}
 }