From 852667e8b89a2389e3528e863b08a6cf4461029f Mon Sep 17 00:00:00 2001
From: Joe Lanford <joe.lanford@gmail.com>
Date: Wed, 24 Aug 2022 09:18:14 -0400
Subject: [PATCH] do not fatally exit from serve commands

When we fatally exit directly in the serve functions, we do not give a
chance for the registered cleanup functions to run, resulting in
temporary files being left behind.

This commit ensures that we always cleanly return an error from the
serve function, thus ensuring that defered cleanups always happen.

Signed-off-by: Joe Lanford <joe.lanford@gmail.com>
---
 cmd/opm/registry/serve.go |  2 +-
 cmd/opm/serve/serve.go    | 45 +++++++++++++++++++++++++++------------
 2 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/cmd/opm/registry/serve.go b/cmd/opm/registry/serve.go
index b6defb75e..78aba5e52 100644
--- a/cmd/opm/registry/serve.go
+++ b/cmd/opm/registry/serve.go
@@ -115,7 +115,7 @@ func serveFunc(cmd *cobra.Command, _ []string) error {
 
 	lis, err := net.Listen("tcp", ":"+port)
 	if err != nil {
-		logger.Fatalf("failed to listen: %s", err)
+		return fmt.Errorf("failed to listen: %s", err)
 	}
 
 	timeout, err := cmd.Flags().GetString("timeout-seconds")
diff --git a/cmd/opm/serve/serve.go b/cmd/opm/serve/serve.go
index 0733cafd9..633ad9ef3 100644
--- a/cmd/opm/serve/serve.go
+++ b/cmd/opm/serve/serve.go
@@ -82,7 +82,9 @@ will not be reflected in the served content.
 
 func (s *serve) run(ctx context.Context) error {
 	p := newProfilerInterface(s.pprofAddr, s.logger)
-	p.startEndpoint()
+	if err := p.startEndpoint(); err != nil {
+		return fmt.Errorf("could not start pprof endpoint: %v", err)
+	}
 	if err := p.startCpuProfileCache(); err != nil {
 		return fmt.Errorf("could not start CPU profile: %v", err)
 	}
@@ -111,7 +113,7 @@ func (s *serve) run(ctx context.Context) error {
 
 	lis, err := net.Listen("tcp", ":"+s.port)
 	if err != nil {
-		s.logger.Fatalf("failed to listen: %s", err)
+		return fmt.Errorf("failed to listen: %s", err)
 	}
 
 	grpcServer := grpc.NewServer()
@@ -125,7 +127,9 @@ func (s *serve) run(ctx context.Context) error {
 		return grpcServer.Serve(lis)
 	}, func() {
 		grpcServer.GracefulStop()
-		p.stopEndpoint(p.logger.Context)
+		if err := p.stopEndpoint(ctx); err != nil {
+			s.logger.Warnf("error shutting down pprof server: %v", err)
+		}
 	})
 
 }
@@ -143,7 +147,8 @@ type profilerInterface struct {
 	cacheReady bool
 	cacheLock  sync.RWMutex
 
-	logger *logrus.Entry
+	logger   *logrus.Entry
+	closeErr chan error
 }
 
 func newProfilerInterface(a string, log *logrus.Entry) *profilerInterface {
@@ -158,10 +163,10 @@ func (p *profilerInterface) isEnabled() bool {
 	return p.addr != ""
 }
 
-func (p *profilerInterface) startEndpoint() {
+func (p *profilerInterface) startEndpoint() error {
 	// short-circuit if not enabled
 	if !p.isEnabled() {
-		return
+		return nil
 	}
 
 	mux := http.NewServeMux()
@@ -177,14 +182,22 @@ func (p *profilerInterface) startEndpoint() {
 		Handler: mux,
 	}
 
-	// goroutine exits with main
-	go func() {
+	lis, err := net.Listen("tcp", p.addr)
+	if err != nil {
+		return err
+	}
 
-		p.logger.Info("starting pprof endpoint")
-		if err := p.server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
-			p.logger.Fatal(err)
-		}
+	p.closeErr = make(chan error)
+	go func() {
+		p.closeErr <- func() error {
+			p.logger.Info("starting pprof endpoint")
+			if err := p.server.Serve(lis); err != nil && !errors.Is(err, http.ErrServerClosed) {
+				return err
+			}
+			return nil
+		}()
 	}()
+	return nil
 }
 
 func (p *profilerInterface) startCpuProfileCache() error {
@@ -218,10 +231,14 @@ func (p *profilerInterface) httpHandler(w http.ResponseWriter, r *http.Request)
 	w.Write(p.cache.Bytes())
 }
 
-func (p *profilerInterface) stopEndpoint(ctx context.Context) {
+func (p *profilerInterface) stopEndpoint(ctx context.Context) error {
+	if !p.isEnabled() {
+		return nil
+	}
 	if err := p.server.Shutdown(ctx); err != nil {
-		p.logger.Fatal(err)
+		return err
 	}
+	return <-p.closeErr
 }
 
 func (p *profilerInterface) isCacheReady() bool {