[v0.51] [AN] Backport graceful shutdown fix (#8662, #8664)

module/grpcserver/interceptor_shutdown.go

@@ -0,0 +1,59 @@
+package grpcserver
+
+import (
+	"context"
+
+	"go.uber.org/atomic"
+	"google.golang.org/grpc"
+
+	"github.com/onflow/flow-go/module/irrecoverable"
+)
+
+// ShutdownStreamInterceptor cancels the per-stream context as soon as the node's
+// [irrecoverable.SignalerContext] is cancelled. This lets long-lived streaming RPCs
+// (e.g. block subscriptions) observe shutdown and return promptly, which in turn allows
+// [grpc.Server.GracefulStop] to finish without waiting for the client to disconnect.
+//
+// Without this interceptor, `stream.Context()` is only cancelled when the underlying
+// transport dies; [grpc.Server.GracefulStop] does not cancel it. That is why active
+// subscriptions block shutdown until either the client disconnects or the server is
+// force-stopped via [grpc.Server.Stop].
+//
+// If `signalerCtx` has not been populated yet (the server is still initializing), the
+// stream is passed through unchanged.
+func ShutdownStreamInterceptor(signalerCtx *atomic.Pointer[irrecoverable.SignalerContext]) grpc.StreamServerInterceptor {
+	return func(srv any, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+		sigCtx := signalerCtx.Load()
+		if sigCtx == nil {
+			return handler(srv, ss)
+		}
+
+		ctx, cancel := context.WithCancel(ss.Context())
+		defer cancel()
+
+		// Fan the SignalerContext's Done into the stream's context. The watcher exits either
+		// when shutdown begins (SignalerContext cancelled) or when the handler returns
+		// (defer cancel() above), so no goroutine is leaked.
+		go func() {
+			select {
+			case <-(*sigCtx).Done():
+				cancel()
+			case <-ctx.Done():
+			}
+		}()
+
+		return handler(srv, &shutdownAwareStream{ServerStream: ss, ctx: ctx})
+	}
+}
+
+// shutdownAwareStream wraps a [grpc.ServerStream] so that its Context reflects both the
+// original transport-level cancellation and node shutdown.
+type shutdownAwareStream struct {
+	grpc.ServerStream
+	ctx context.Context
+}
+
+// Context returns the shutdown-aware context for the stream.
+func (s *shutdownAwareStream) Context() context.Context {
+	return s.ctx
+}

module/grpcserver/server.go

@@ -3,6 +3,7 @@ package grpcserver
 import (
 	"net"
 	"sync"
+	"time"
 
 	"go.uber.org/atomic"
 
@@ -18,6 +19,11 @@ import (
 	"github.com/onflow/flow-go/module/irrecoverable"
 )
 
+// DefaultGracefulStopTimeout is the time GracefulStop is allowed to wait for active streaming
+// RPCs to finish before the server is force-stopped via Stop. Long-lived streaming subscriptions
+// would otherwise block shutdown indefinitely.
+const DefaultGracefulStopTimeout = 5 * time.Second
+
 // GrpcServer wraps `grpc.Server` and allows to manage it using `component.Component` interface. It can be injected
 // into different engines making it possible to use single grpc server for multiple services which live in different modules.
 type GrpcServer struct {
@@ -30,25 +36,32 @@ type GrpcServer struct {
 	// within handler code.
 	grpcSignalerCtx *atomic.Pointer[irrecoverable.SignalerContext]
 
-	grpcListenAddr string // the GRPC server address as ip:port
+	grpcListenAddr      string        // the GRPC server address as ip:port
+	gracefulStopTimeout time.Duration // how long to wait for active RPCs to finish before force-stopping
 
 	addrLock    sync.RWMutex
 	grpcAddress net.Addr
 }
 
 var _ component.Component = (*GrpcServer)(nil)
 
-// NewGrpcServer returns a new grpc server.
+// NewGrpcServer returns a new grpc server. If gracefulStopTimeout is zero,
+// DefaultGracefulStopTimeout is used.
 func NewGrpcServer(log zerolog.Logger,
 	grpcListenAddr string,
 	grpcServer *grpc.Server,
 	grpcSignalerCtx *atomic.Pointer[irrecoverable.SignalerContext],
+	gracefulStopTimeout time.Duration,
 ) *GrpcServer {
+	if gracefulStopTimeout <= 0 {
+		gracefulStopTimeout = DefaultGracefulStopTimeout
+	}
 	server := &GrpcServer{
-		log:             log,
-		server:          grpcServer,
-		grpcListenAddr:  grpcListenAddr,
-		grpcSignalerCtx: grpcSignalerCtx,
+		log:                 log,
+		server:              grpcServer,
+		grpcListenAddr:      grpcListenAddr,
+		grpcSignalerCtx:     grpcSignalerCtx,
+		gracefulStopTimeout: gracefulStopTimeout,
 	}
 	server.Component = component.NewComponentManagerBuilder().
 		AddWorker(server.serveGRPCWorker).
@@ -104,8 +117,24 @@ func (g *GrpcServer) GRPCAddress() net.Addr {
 }
 
 // shutdownWorker is a worker routine which shuts down server when the context is cancelled.
+// It attempts a graceful stop first. If active streaming RPCs do not finish within
+// gracefulStopTimeout, the server is force-stopped to avoid blocking shutdown indefinitely.
 func (g *GrpcServer) shutdownWorker(ctx irrecoverable.SignalerContext, ready component.ReadyFunc) {
 	ready()
 	<-ctx.Done()
-	g.server.GracefulStop()
+	gracefulDone := make(chan struct{})
+	go func() {
+		defer close(gracefulDone)
+		g.server.GracefulStop()
+	}()
+	timer := time.NewTimer(g.gracefulStopTimeout)
+	defer timer.Stop()
+	select {
+	case <-gracefulDone:
+	case <-timer.C:
+		g.log.Warn().
+			Dur("timeout", g.gracefulStopTimeout).
+			Msg("graceful stop timed out; force-stopping gRPC server")
+		g.server.Stop()
+	}
 }

module/grpcserver/server_builder.go

@@ -1,6 +1,8 @@
 package grpcserver
 
 import (
+	"time"
+
 	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
 	"github.com/rs/zerolog"
 	"go.uber.org/atomic"
@@ -26,6 +28,14 @@ func WithStreamInterceptor() Option {
 	}
 }
 
+// WithGracefulStopTimeout sets how long the server waits for active streaming RPCs to finish
+// before force-stopping during shutdown. Defaults to DefaultGracefulStopTimeout.
+func WithGracefulStopTimeout(d time.Duration) Option {
+	return func(c *GrpcServerBuilder) {
+		c.gracefulStopTimeout = d
+	}
+}
+
 // GrpcServerBuilder created for separating the creation and starting GrpcServer,
 // cause services need to be registered before the server starts.
 type GrpcServerBuilder struct {
@@ -36,6 +46,7 @@ type GrpcServerBuilder struct {
 
 	transportCredentials         credentials.TransportCredentials // the GRPC credentials
 	stateStreamInterceptorEnable bool
+	gracefulStopTimeout          time.Duration
 }
 
 // NewGrpcServerBuilder creates a new builder for configuring and initializing a gRPC server.
@@ -88,6 +99,12 @@ func NewGrpcServerBuilder(
 	var streamInterceptors []grpc.StreamServerInterceptor
 
 	unaryInterceptors = append(unaryInterceptors, IrrecoverableCtxInjector(signalerCtx))
+
+	// ShutdownStreamInterceptor must be registered before any interceptor or handler that
+	// reads stream.Context(), so subsequent interceptors and the handler see the
+	// shutdown-aware context.
+	streamInterceptors = append(streamInterceptors, ShutdownStreamInterceptor(signalerCtx))
+
 	if rpcMetricsEnabled {
 		unaryInterceptors = append(unaryInterceptors, grpc_prometheus.UnaryServerInterceptor)
 
@@ -124,5 +141,9 @@ func NewGrpcServerBuilder(
 }
 
 func (b *GrpcServerBuilder) Build() *GrpcServer {
-	return NewGrpcServer(b.log, b.gRPCListenAddr, b.server, b.signalerCtx)
+	timeout := b.gracefulStopTimeout
+	if timeout == 0 {
+		timeout = DefaultGracefulStopTimeout
+	}
+	return NewGrpcServer(b.log, b.gRPCListenAddr, b.server, b.signalerCtx, timeout)
 }