diff --git a/module/grpcserver/interceptor_shutdown.go b/module/grpcserver/interceptor_shutdown.go new file mode 100644 index 00000000000..ad9bcc8eb19 --- /dev/null +++ b/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 +} diff --git a/module/grpcserver/server.go b/module/grpcserver/server.go index 4cd2ada4db9..14c756950f2 100644 --- a/module/grpcserver/server.go +++ b/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,7 +36,8 @@ 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 @@ -38,17 +45,23 @@ type GrpcServer struct { 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() + } } diff --git a/module/grpcserver/server_builder.go b/module/grpcserver/server_builder.go index f5fb47941f8..0e93c11e8a9 100644 --- a/module/grpcserver/server_builder.go +++ b/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) } diff --git a/module/grpcserver/server_test.go b/module/grpcserver/server_test.go new file mode 100644 index 00000000000..a2083a1edc6 --- /dev/null +++ b/module/grpcserver/server_test.go @@ -0,0 +1,202 @@ +package grpcserver_test + +import ( + "context" + "testing" + "time" + + "github.com/rs/zerolog" + "github.com/stretchr/testify/require" + "go.uber.org/atomic" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + + "github.com/onflow/flow-go/module/grpcserver" + "github.com/onflow/flow-go/module/irrecoverable" + "github.com/onflow/flow-go/utils/unittest" +) + +// blockingStreamService is the interface gRPC uses to type-check the registered handler. +type blockingStreamService interface { + Stream(grpc.ServerStream) error +} + +// blockingStreamServiceDesc is a gRPC service descriptor with a single server-streaming method. +// The handler blocks until the stream context is cancelled, simulating a long-lived subscription. +var blockingStreamServiceDesc = grpc.ServiceDesc{ + ServiceName: "test.BlockingStream", + HandlerType: (*blockingStreamService)(nil), + Methods: []grpc.MethodDesc{}, + Streams: []grpc.StreamDesc{ + { + StreamName: "Stream", + Handler: blockingStreamHandler, + ServerStreams: true, + }, + }, +} + +type blockingStreamServer struct { + // started is closed when the stream handler has been entered. + started chan struct{} + // blockDuration will cause `Stream` to block for the set duration + blockDuration time.Duration +} + +var _ blockingStreamService = (*blockingStreamServer)(nil) + +func (s *blockingStreamServer) Stream(stream grpc.ServerStream) error { + close(s.started) + if s.blockDuration > 0 { + // this is to simulate the case that after `grpcServer.Stop()` is called, + // `<-gracefulDone` channel is still blocking, so that we can verify + // the caller is not waiting for `<-gracefulDone` return before shutdown, + // otherwise, the waiting might be still blocking for longer or indefinitely. + time.Sleep(s.blockDuration) + } + <-stream.Context().Done() + return nil +} + +func blockingStreamHandler(srv any, stream grpc.ServerStream) error { + return srv.(blockingStreamService).Stream(stream) +} + +// TestGrpcServerShutdown_WithActiveStream verifies that GrpcServer shuts down within +// gracefulStopTimeout even when a long-lived streaming RPC is active and the client +// has not disconnected. Without the fix, GracefulStop() would block indefinitely. +func TestGrpcServerShutdown_WithActiveStream(t *testing.T) { + gracefulStopTimeout := 200 * time.Millisecond + + rawServer := grpc.NewServer() + handler := &blockingStreamServer{ + started: make(chan struct{}), + blockDuration: gracefulStopTimeout * 10, + } + rawServer.RegisterService(&blockingStreamServiceDesc, handler) + + signalerCtx := atomic.NewPointer[irrecoverable.SignalerContext](nil) + server := grpcserver.NewGrpcServer( + zerolog.Nop(), + "localhost:0", + rawServer, + signalerCtx, + gracefulStopTimeout, + ) + + ctx, cancel := irrecoverable.NewMockSignalerContextWithCancel(t, context.Background()) + server.Start(ctx) + unittest.RequireComponentsReadyBefore(t, 2*time.Second, server) + + conn, err := grpc.NewClient( + server.GRPCAddress().String(), + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + require.NoError(t, err) + defer conn.Close() + + // Open a stream; do not cancel clientCtx so the stream stays open indefinitely. + clientCtx := t.Context() + _, err = conn.NewStream(clientCtx, &grpc.StreamDesc{ServerStreams: true}, "/test.BlockingStream/Stream") + require.NoError(t, err) + + // Wait until the server-side handler is running. + unittest.RequireCloseBefore(t, handler.started, 2*time.Second, "stream handler did not start") + + // Trigger node shutdown. The stream is still open on the client side. + cancel() + + // The server must complete shutdown within gracefulStopTimeout plus a small buffer. + // Before the fix, this would hang indefinitely because GracefulStop() waits for all + // active streaming RPCs to finish, and the client never disconnects. + unittest.RequireComponentsDoneBefore(t, gracefulStopTimeout+500*time.Millisecond, server) +} + +// TestGrpcServerShutdown_ShutdownStreamInterceptor verifies that when the +// [grpcserver.ShutdownStreamInterceptor] is registered, an active streaming RPC's +// stream.Context() is cancelled as soon as the node's SignalerContext is cancelled. +// This allows GracefulStop to complete cleanly — well under gracefulStopTimeout — +// even when the client has not disconnected, so the force-stop fallback is not needed. +func TestGrpcServerShutdown_ShutdownStreamInterceptor(t *testing.T) { + // Give the graceful path a generous window so we can prove that the interceptor — + // not the force-stop fallback — is what unblocks shutdown. + gracefulStopTimeout := 10 * time.Second + + signalerCtx := atomic.NewPointer[irrecoverable.SignalerContext](nil) + rawServer := grpc.NewServer( + grpc.ChainStreamInterceptor(grpcserver.ShutdownStreamInterceptor(signalerCtx)), + ) + handler := &blockingStreamServer{ + started: make(chan struct{}), + blockDuration: time.Second, + } + rawServer.RegisterService(&blockingStreamServiceDesc, handler) + + server := grpcserver.NewGrpcServer( + zerolog.Nop(), + "localhost:0", + rawServer, + signalerCtx, + gracefulStopTimeout, + ) + + ctx, cancel := irrecoverable.NewMockSignalerContextWithCancel(t, context.Background()) + server.Start(ctx) + unittest.RequireComponentsReadyBefore(t, 2*time.Second, server) + + conn, err := grpc.NewClient( + server.GRPCAddress().String(), + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + require.NoError(t, err) + defer conn.Close() + + // Open a stream and never cancel the client-side context — mimicking a long-lived + // subscription that a well-behaved client is happy to keep open indefinitely. + clientCtx := t.Context() + _, err = conn.NewStream(clientCtx, &grpc.StreamDesc{ServerStreams: true}, "/test.BlockingStream/Stream") + require.NoError(t, err) + + unittest.RequireCloseBefore(t, handler.started, 2*time.Second, "stream handler did not start") + + // Trigger node shutdown. The interceptor should cancel the stream's context, the + // handler should return, and GracefulStop should complete immediately. + cancel() + + // Shutdown must complete well under gracefulStopTimeout; otherwise the force-stop + // fallback is what unblocked us, not the interceptor. + unittest.RequireComponentsDoneBefore(t, 2*time.Second, server) +} + +// TestGrpcServerShutdown_NoActiveStreams verifies that when no streaming RPCs are active, +// GrpcServer shuts down promptly via GracefulStop without waiting for the timeout. +func TestGrpcServerShutdown_NoActiveStreams(t *testing.T) { + gracefulStopTimeout := 5 * time.Second + + rawServer := grpc.NewServer() + rawServer.RegisterService( + &blockingStreamServiceDesc, + &blockingStreamServer{ + started: make(chan struct{}), + blockDuration: gracefulStopTimeout * 10, + }, + ) + + signalerCtx := atomic.NewPointer[irrecoverable.SignalerContext](nil) + server := grpcserver.NewGrpcServer( + zerolog.Nop(), + "localhost:0", + rawServer, + signalerCtx, + gracefulStopTimeout, + ) + + ctx, cancel := irrecoverable.NewMockSignalerContextWithCancel(t, context.Background()) + server.Start(ctx) + unittest.RequireComponentsReadyBefore(t, 2*time.Second, server) + + cancel() + + // With no active streams, GracefulStop() completes immediately — well under the 5s timeout. + unittest.RequireComponentsDoneBefore(t, 500*time.Millisecond, server) +}