Start the profiling server properly (#4271)

Signed-off-by: Eduardo Apolinario <[email protected]> Co-authored-by: Eduardo Apolinario <[email protected]>
flyteorg · Oct 20, 2023 · e2b1b0a · e2b1b0a
1 parent 8749fce
commit e2b1b0a
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 9 deletions.
diff --git a/cmd/single/start.go b/cmd/single/start.go
@@ -168,11 +168,7 @@ func startPropeller(ctx context.Context, cfg Propeller) error {
 		}
 
 		g.Go(func() error {
-			err := profutils.StartProfilingServerWithDefaultHandlers(childCtx, propellerCfg.ProfilerPort.Port, handlers)
-			if err != nil {
-				logger.Fatalf(childCtx, "Failed to Start profiling and metrics server. Error: %v", err)
-			}
-			return err
+			return profutils.StartProfilingServerWithDefaultHandlers(childCtx, propellerCfg.ProfilerPort.Port, handlers)
 		})
 
 		g.Go(func() error {

diff --git a/flytestdlib/profutils/server.go b/flytestdlib/profutils/server.go
@@ -101,13 +101,13 @@ func StartProfilingServer(ctx context.Context, pprofPort int) error {
 	}
 
 	go func() {
-		e := srv.ListenAndServe()
-		if e != nil {
-			logger.Errorf(ctx, "Failed to start profiling server. Error: %v", e)
+		<-ctx.Done()
+		if err := srv.Shutdown(context.Background()); err != nil {
+			logger.Errorf(ctx, "Failed to gracefully shutdown profiling server. Error: %v", err)
 		}
 	}()
 
-	return srv.Shutdown(ctx)
+	return srv.ListenAndServe()
 }
 
 func configureGlobalHTTPHandler(handlers map[string]http.Handler) error {