diff --git a/gapis/server/grpc.go b/gapis/server/grpc.go index e735446d22..b7e984f2ea 100644 --- a/gapis/server/grpc.go +++ b/gapis/server/grpc.go @@ -74,6 +74,8 @@ func NewWithListener(ctx context.Context, l net.Listener, cfg Config, srvChan ch } if cfg.IdleTimeout != 0 { crash.Go(func() { s.stopIfIdle(ctx, server, cfg.IdleTimeout, stop) }) + } else { + crash.Go(func() { s.stopOnInterrupt(ctx, server, stop) }) } return nil }, grpc.UnaryInterceptor(auth.ServerInterceptor(cfg.AuthToken))) @@ -117,7 +119,7 @@ func (s *grpcServer) inRPC() func() { } // stopIfIdle calls GracefulStop on server if there are no writes the the -// keepAlive chan within idleTimeout. +// keepAlive chan within idleTimeout or if the current process is interrupted. // This function blocks until there's an idle timeout, or ctx is cancelled. func (s *grpcServer) stopIfIdle(ctx context.Context, server *grpc.Server, idleTimeout time.Duration, stop func()) { // Split the idleTimeout into N smaller chunks, and check that there was @@ -158,6 +160,21 @@ func (s *grpcServer) stopIfIdle(ctx context.Context, server *grpc.Server, idleTi } } +// stopOnInterrupt calls GracefulStop on server if the current process is interrupted. +func (s *grpcServer) stopOnInterrupt(ctx context.Context, server *grpc.Server, stop func()) { + stoppedSignal, stopped := task.NewSignal() + defer func() { + stop() + server.GracefulStop() + stopped(ctx) + }() + + // Wait for the server to stop before terminating the app. + app.AddCleanupSignal(stoppedSignal) + + <-task.ShouldStop(ctx) +} + func (s *grpcServer) addInterrupter(f func()) (remove func()) { li := s.lastInterrupter s.lastInterrupter++