diff --git a/cmd/single/start.go b/cmd/single/start.go index 90d86684bd..8a499dbe92 100644 --- 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 index a1b5baf5ea..e4623c159d 100644 --- 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 {