chore: optionally enables pprof on the controller (#1406)

**Description**

This also enables pprof server optionally on the controller in addition
to extproc.

**Related Issues/PRs (if applicable)**

Follow up on #1395

---------

Signed-off-by: Takeshi Yoneda <[email protected]>

Author: mathetake

cmd/controller/main.go

@@ -33,6 +33,7 @@ import (
 	"github.com/envoyproxy/ai-gateway/internal/controller"
 	"github.com/envoyproxy/ai-gateway/internal/extensionserver"
 	"github.com/envoyproxy/ai-gateway/internal/internalapi"
+	"github.com/envoyproxy/ai-gateway/internal/pprof"
 )
 
 type flags struct {
@@ -306,6 +307,7 @@ func main() {
 	setupLog.Info("configuring kubernetes cache", "watch-namespaces", parsedFlags.watchNamespaces, "sync-timeout", parsedFlags.cacheSyncTimeout)
 
 	ctx := ctrl.SetupSignalHandler()
+	pprof.Run(ctx)
 	mgrOpts := ctrl.Options{
 		Cache:            setupCache(parsedFlags),
 		Controller:       config.Controller{CacheSyncTimeout: parsedFlags.cacheSyncTimeout},

cmd/extproc/main.go

@@ -7,43 +7,21 @@ package main
 
 import (
 	"context"
-	"errors"
 	"log"
-	"net/http"
-	"net/http/pprof"
 	"os"
 	"os/signal"
 	"syscall"
 	"time"
 
 	"github.com/envoyproxy/ai-gateway/cmd/extproc/mainlib"
+	"github.com/envoyproxy/ai-gateway/internal/pprof"
 )
 
 func main() {
-	// Run the pprof server if the DISABLE_PPROF environment variable is not set.
-	//
-	// Enabling the pprof server by default helps with debugging performance issues in production.
-	// The impact should be negligible when the actual pprof endpoints are not being accessed.
-	if _, ok := os.LookupEnv("DISABLE_PPROF"); !ok {
-		go func() {
-			const pprofPort = "6060" // The same default port as in th Go pprof documentation.
-			mux := http.NewServeMux()
-			mux.HandleFunc("/debug/pprof/", pprof.Index)
-			mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
-			mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
-			mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
-			mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
-			server := &http.Server{Addr: ":" + pprofPort, Handler: mux, ReadHeaderTimeout: 5 * time.Second}
-			log.Printf("starting pprof server on port %s", pprofPort)
-			if err := server.ListenAndServe(); err != nil && errors.Is(err, http.ErrServerClosed) {
-				log.Printf("pprof server stopped: %v", err)
-			}
-		}()
-	}
-
 	ctx, cancel := context.WithCancel(context.Background())
 	signalsChan := make(chan os.Signal, 1)
 	signal.Notify(signalsChan, syscall.SIGINT, syscall.SIGTERM)
+	pprof.Run(ctx)
 	go func() {
 		<-signalsChan
 		log.Printf("signal received, shutting down...")

internal/pprof/pprof.go

@@ -0,0 +1,57 @@
+// Copyright Envoy AI Gateway Authors
+// SPDX-License-Identifier: Apache-2.0
+// The full text of the Apache license is available in the LICENSE file at
+// the root of the repo.
+
+package pprof
+
+import (
+	"context"
+	"errors"
+	"log"
+	"net/http"
+	"net/http/pprof"
+	"os"
+	"time"
+)
+
+const (
+	pprofPort = "6060" // The same default port as in th Go pprof documentation.
+	// DisableEnvVarKey is the environment variable name to disable the pprof server.
+	// If this environment variable is set to any value, the pprof server will not be started.
+	DisableEnvVarKey = "DISABLE_PPROF"
+)
+
+// Run the pprof server if the DISABLE_PPROF environment variable is not set.
+// This is non-blocking and will run the pprof server in a separate goroutine until the provided context is cancelled.
+//
+// Enabling the pprof server by default helps with debugging performance issues in production.
+// The impact should be negligible when the actual pprof endpoints are not being accessed.
+func Run(ctx context.Context) {
+	if _, ok := os.LookupEnv(DisableEnvVarKey); !ok {
+		mux := http.NewServeMux()
+		mux.HandleFunc("/debug/pprof/", pprof.Index)
+		mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+		mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
+		mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+		mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
+		server := &http.Server{Addr: ":" + pprofPort, Handler: mux, ReadHeaderTimeout: 5 * time.Second}
+		go func() {
+			log.Printf("starting pprof server on port %s", pprofPort)
+			if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+				log.Printf("pprof server stopped: %v", err)
+			}
+		}()
+		go func() {
+			<-ctx.Done()
+			log.Printf("shutting down pprof server...")
+			shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer cancel()
+			if err := server.Shutdown(shutdownCtx); err != nil {
+				log.Printf("error shutting down pprof server: %v", err)
+			} else {
+				log.Print("pprof server shut down gracefully")
+			}
+		}()
+	}
+}

internal/pprof/pprof_test.go

@@ -0,0 +1,44 @@
+// Copyright Envoy AI Gateway Authors
+// SPDX-License-Identifier: Apache-2.0
+// The full text of the Apache license is available in the LICENSE file at
+// the root of the repo.
+
+package pprof
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestRun_disabled(t *testing.T) {
+	t.Setenv(DisableEnvVarKey, "anything")
+	ctx, cancel := context.WithCancel(context.Background())
+	Run(ctx)
+	// Try accessing the pprof server here if needed.
+	response, err := http.Get("http://localhost:6060/debug/pprof/") //nolint:bodyclose
+	require.Error(t, err)
+	require.Nil(t, response)
+	cancel()
+}
+
+func TestRun_enabled(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	Run(ctx)
+	// Try accessing the pprof server here if needed.
+	resp, err := http.Get("http://localhost:6060/debug/pprof/cmdline")
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, resp.Body.Close())
+	}()
+	require.NotNil(t, resp)
+	body, err := io.ReadAll(resp.Body)
+	require.NoError(t, err)
+	require.Contains(t, string(body),
+		// Test binary name should be present in the cmdline output.
+		"pprof.test")
+	cancel()
+}

manifests/charts/ai-gateway-helm/values.yaml

@@ -107,9 +107,10 @@ controller:
   ## @param controller.extraEnvVars Array with extra environment variables to add to controller containers
   ## e.g:
   ## extraEnvVars:
-  ##   - name: FOO
-  ##     value: "bar"
-  ##
+  ##   # This disables pprof endpoint at "localhost:6060/debug/pprof" that can be accessed via port-forwarding for profiling.
+  ##   # Enabled by default for troubleshooting purposes given the impact is negligible when not in use.
+  ##   - name: DISABLE_PPROF
+  ##     value: "true"
   extraEnvVars: []
 
   # Example of volumes

tests/internal/testenvironment/test_environment.go

@@ -24,6 +24,7 @@ import (
 	"golang.org/x/sync/errgroup"
 
 	"github.com/envoyproxy/ai-gateway/cmd/extproc/mainlib"
+	"github.com/envoyproxy/ai-gateway/internal/pprof"
 	internaltesting "github.com/envoyproxy/ai-gateway/internal/testing"
 	testsinternal "github.com/envoyproxy/ai-gateway/tests/internal"
 )
@@ -313,6 +314,8 @@ func requireExtProc(t testing.TB, out io.Writer, bin, config string, env []strin
 		"-mcpWriteTimeout", mcpWriteTimeout.String(),
 		"-logLevel", "info",
 	}
+	// Disable pprof for tests to avoid port conflicts.
+	env = append(env, fmt.Sprintf("%s=true", pprof.DisableEnvVarKey))
 	t.Logf("Starting ExtProc with args: %v", args)
 	if inProcess {
 		go func() {