Merge branch 'main' into add-atomic-cache-ops

management/internals/server/boot.go

@@ -182,6 +182,10 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
 			grpc.ChainStreamInterceptor(realip.StreamServerInterceptorOpts(realipOpts...), streamInterceptor, proxyStream),
 		}
 
+		// Append interceptors contributed by registered gRPC extensions. These
+		// run after the built-in chain (ChainUnaryInterceptor is additive).
+		gRPCOpts = appendExtensionInterceptors(gRPCOpts, s.grpcExtensions)
+
 		if s.Config.HttpConfig.LetsEncryptDomain != "" {
 			certManager, err := encryption.CreateCertManager(s.Config.Datadir, s.Config.HttpConfig.LetsEncryptDomain)
 			if err != nil {
@@ -213,6 +217,9 @@ func (s *BaseServer) GRPCServer() *grpc.Server {
 		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
 	})
 }

management/internals/server/grpc_extension.go

@@ -0,0 +1,74 @@
+package server
+
+import (
+	"context"
+
+	"google.golang.org/grpc"
+)
+
+// GRPCExtension bundles an external module's contribution to the management
+// gRPC server: the registration of one or more services onto the shared
+// grpc.Server, any server-wide interceptors those services require, and an
+// optional shutdown hook. It is a generic extension point with no knowledge of
+// any specific service.
+type GRPCExtension struct {
+	// Register is invoked with the shared grpc.Server (as a ServiceRegistrar)
+	// after the built-in services are registered. It may register any number of
+	// services. May be nil.
+	Register func(grpc.ServiceRegistrar)
+	// UnaryInterceptors are appended to the server's unary interceptor chain,
+	// running after the built-in interceptors. May be empty.
+	UnaryInterceptors []grpc.UnaryServerInterceptor
+	// StreamInterceptors are appended to the server's stream interceptor chain,
+	// running after the built-in interceptors. May be empty.
+	StreamInterceptors []grpc.StreamServerInterceptor
+	// Shutdown, if non-nil, is called once during Stop() with the context
+	// governing server shutdown, which carries a deadline. The hook MUST
+	// return promptly and MUST abandon its work once that context is
+	// cancelled or expires: it runs before the rest of Stop()'s cleanup
+	// (store, event store, embedded IdP) and before Stop() itself checks the
+	// context's deadline, so a hook that ignores the context will delay all
+	// of that cleanup and prevent Stop() from returning on time. May be nil.
+	Shutdown func(ctx context.Context)
+}
+
+// RegisterGRPCExtension registers a gRPC extension. Call before the gRPC server
+// is first built (i.e. before Start); registrations after that have no effect.
+func (s *BaseServer) RegisterGRPCExtension(ext GRPCExtension) {
+	s.grpcExtensions = append(s.grpcExtensions, ext)
+}
+
+// appendExtensionInterceptors appends each extension's interceptors to the gRPC
+// server options as additional chained interceptors. grpc.ChainUnaryInterceptor
+// and grpc.ChainStreamInterceptor are additive, so the returned options run the
+// extension interceptors after any interceptors already present in opts.
+func appendExtensionInterceptors(opts []grpc.ServerOption, exts []GRPCExtension) []grpc.ServerOption {
+	for _, ext := range exts {
+		if len(ext.UnaryInterceptors) > 0 {
+			opts = append(opts, grpc.ChainUnaryInterceptor(ext.UnaryInterceptors...))
+		}
+		if len(ext.StreamInterceptors) > 0 {
+			opts = append(opts, grpc.ChainStreamInterceptor(ext.StreamInterceptors...))
+		}
+	}
+	return opts
+}
+
+// registerExtensions registers each extension's services onto reg.
+func registerExtensions(reg grpc.ServiceRegistrar, exts []GRPCExtension) {
+	for _, ext := range exts {
+		if ext.Register != nil {
+			ext.Register(reg)
+		}
+	}
+}
+
+// runExtensionShutdownHooks calls each extension's shutdown hook, if set,
+// passing ctx through so hooks can honor its deadline/cancellation.
+func runExtensionShutdownHooks(ctx context.Context, exts []GRPCExtension) {
+	for _, ext := range exts {
+		if ext.Shutdown != nil {
+			ext.Shutdown(ctx)
+		}
+	}
+}

management/internals/server/grpc_extension_test.go

@@ -0,0 +1,160 @@
+package server
+
+import (
+	"context"
+	"net"
+	"sync/atomic"
+	"testing"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/health"
+	healthgrpc "google.golang.org/grpc/health/grpc_health_v1"
+	"google.golang.org/grpc/test/bufconn"
+)
+
+// Test that an extension's interceptors and service registration are actually
+// wired onto a real in-process gRPC server via the helpers, and that shutdown
+// hooks run. This validates the load-bearing assumption that
+// grpc.ChainUnaryInterceptor is additive (extension interceptors run in
+// addition to any base chain).
+func TestGRPCExtensionAppliedToServer(t *testing.T) {
+	var unaryCalls atomic.Int32
+	var streamShutdownCalled atomic.Bool
+
+	ext := GRPCExtension{
+		Register: func(reg grpc.ServiceRegistrar) {
+			healthgrpc.RegisterHealthServer(reg, health.NewServer())
+		},
+		UnaryInterceptors: []grpc.UnaryServerInterceptor{
+			func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
+				unaryCalls.Add(1)
+				return handler(ctx, req)
+			},
+		},
+		Shutdown: func(ctx context.Context) { streamShutdownCalled.Store(true) },
+	}
+	exts := []GRPCExtension{ext}
+
+	// Base options mimic GRPCServer(): a pre-existing chain the extension appends to.
+	var baseUnaryCalls atomic.Int32
+	opts := []grpc.ServerOption{
+		grpc.ChainUnaryInterceptor(func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
+			baseUnaryCalls.Add(1)
+			return handler(ctx, req)
+		}),
+	}
+	opts = appendExtensionInterceptors(opts, exts)
+
+	srv := grpc.NewServer(opts...)
+	registerExtensions(srv, exts)
+
+	lis := bufconn.Listen(1024 * 1024)
+	go func() { _ = srv.Serve(lis) }()
+	t.Cleanup(srv.Stop)
+
+	conn, err := grpc.NewClient("passthrough:///bufnet",
+		grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) { return lis.DialContext(ctx) }),
+		grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = conn.Close() })
+
+	_, err = healthgrpc.NewHealthClient(conn).Check(context.Background(), &healthgrpc.HealthCheckRequest{})
+	if err != nil {
+		t.Fatalf("health check via extension-registered service failed: %v", err)
+	}
+	if baseUnaryCalls.Load() != 1 {
+		t.Errorf("base interceptor calls = %d, want 1 (base chain must be preserved)", baseUnaryCalls.Load())
+	}
+	if unaryCalls.Load() != 1 {
+		t.Errorf("extension interceptor calls = %d, want 1", unaryCalls.Load())
+	}
+
+	runExtensionShutdownHooks(context.Background(), exts)
+	if !streamShutdownCalled.Load() {
+		t.Error("extension shutdown hook was not called")
+	}
+}
+
+// TestGRPCExtensionShutdownHookReceivesCallerContext asserts that each hook receives
+// a non-nil context and that it is the very same context the caller passed
+// in, so hooks can rely on values/deadlines placed on it by Stop().
+func TestGRPCExtensionShutdownHookReceivesCallerContext(t *testing.T) {
+	type sentinelKey struct{}
+	want := "shutdown-ctx-sentinel"
+	ctx := context.WithValue(context.Background(), sentinelKey{}, want)
+
+	var called bool
+	ext := GRPCExtension{
+		Shutdown: func(hookCtx context.Context) {
+			called = true
+			if hookCtx == nil {
+				t.Fatal("hook received a nil context")
+			}
+			got, _ := hookCtx.Value(sentinelKey{}).(string)
+			if got != want {
+				t.Errorf("hook context sentinel = %q, want %q (not the caller's context)", got, want)
+			}
+		},
+	}
+
+	runExtensionShutdownHooks(ctx, []GRPCExtension{ext})
+	if !called {
+		t.Fatal("shutdown hook was not called")
+	}
+}
+
+// TestGRPCExtensionShutdownHookObservesCancellation documents, by test, that
+// hooks can honor cancellation/deadlines: a hook given an already-cancelled
+// context must see ctx.Err() != nil and a closed Done() channel.
+func TestGRPCExtensionShutdownHookObservesCancellation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	var called bool
+	ext := GRPCExtension{
+		Shutdown: func(hookCtx context.Context) {
+			called = true
+			if hookCtx.Err() == nil {
+				t.Error("hook context Err() = nil, want non-nil for a cancelled context")
+			}
+			select {
+			case <-hookCtx.Done():
+			default:
+				t.Error("hook context Done() channel is not closed for a cancelled context")
+			}
+		},
+	}
+
+	runExtensionShutdownHooks(ctx, []GRPCExtension{ext})
+	if !called {
+		t.Fatal("shutdown hook was not called")
+	}
+}
+
+// TestGRPCExtensionShutdownHookNilSkipped asserts that an extension
+// with a nil Shutdown hook is skipped without panicking, and that hooks for
+// other extensions still run.
+func TestGRPCExtensionShutdownHookNilSkipped(t *testing.T) {
+	var called atomic.Bool
+	exts := []GRPCExtension{
+		{Shutdown: nil},
+		{Shutdown: func(context.Context) { called.Store(true) }},
+	}
+
+	runExtensionShutdownHooks(context.Background(), exts)
+	if !called.Load() {
+		t.Error("shutdown hook for non-nil extension was not called")
+	}
+}
+
+func TestRegisterGRPCExtensionAccumulates(t *testing.T) {
+	s := &BaseServer{}
+	s.RegisterGRPCExtension(GRPCExtension{})
+	s.RegisterGRPCExtension(GRPCExtension{})
+	if len(s.grpcExtensions) != 2 {
+		t.Fatalf("grpcExtensions len = %d, want 2", len(s.grpcExtensions))
+	}
+}

management/internals/server/server.go

@@ -68,6 +68,11 @@ type BaseServer struct {
 
 	proxyAuthClose func()
 
+	// grpcExtensions holds additional gRPC services, interceptors, and shutdown
+	// hooks registered by external modules via RegisterGRPCExtension. Populated
+	// during boot (single-threaded), consumed by GRPCServer() and Stop().
+	grpcExtensions []GRPCExtension
+
 	listener    net.Listener
 	certManager *autocert.Manager
 	update      *version.Update
@@ -257,6 +262,7 @@ func (s *BaseServer) Stop() error {
 		s.proxyAuthClose()
 		s.proxyAuthClose = nil
 	}
+	runExtensionShutdownHooks(ctx, s.grpcExtensions)
 	_ = s.Store().Close(ctx)
 	_ = s.EventStore().Close(ctx)
 	if s.update != nil {