feat: add admin server with pprof (#1564)

README.md

@@ -434,6 +434,20 @@ To enable Prometheus, use the `--prometheus` flag. This will start an HTTP
 server on localhost with a `/metrics` endpoint. The Prometheus namespace may
 optionally be set with `--prometheus-namespace`.
 
+## Localhost Admin Server
+
+The Proxy includes support for an admin server on localhost. By default, the
+admin server is not enabled. To enable the server, pass the `--debug` flag.
+This will start the server on localhost at port 9091. To change the port, use
+the `--admin-port` flag.
+
+The admin server includes Go's pprof tool and is available at `/debug/pprof/`.
+
+See the [documentation on pprof][pprof] for details on how to use the
+profiler.
+
+[pprof]: https://pkg.go.dev/net/http/pprof.
+
 ## Frequently Asked Questions
 
 ### Why would I use the proxy?

cmd/root.go

@@ -22,6 +22,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"net/http/pprof"
 	"net/url"
 	"os"
 	"os/signal"
@@ -238,6 +239,19 @@ Configuration using environment variables
         CSQL_PROXY_INSTANCE_CONNECTION_NAME_1=my-other-project:us-central1:my-other-server \
             ./cloud-sql-proxy
 
+Localhost Admin Server
+
+    The Proxy includes support for an admin server on localhost. By default,
+    the admin server is not enabled. To enable the server, pass the --debug
+    flag. This will start the server on localhost at port 9091. To change the
+    port, use the --admin-port flag.
+
+    The admin server includes Go's pprof tool and is available at
+    /debug/pprof/.
+
+    See the documentation on pprof for details on how to use the
+    profiler at https://pkg.go.dev/net/http/pprof.
+
 (*) indicates a flag that may be used as a query parameter
 
 `
@@ -360,6 +374,10 @@ func NewCommand(opts ...Option) *Command {
 		"Address for Prometheus and health check server")
 	pflags.StringVar(&c.conf.HTTPPort, "http-port", "9090",
 		"Port for Prometheus and health check server")
+	pflags.BoolVar(&c.conf.Debug, "debug", false,
+		"Enable the admin server on localhost")
+	pflags.StringVar(&c.conf.AdminPort, "admin-port", "9091",
+		"Port for localhost-only admin server")
 	pflags.BoolVar(&c.conf.HealthCheck, "health-check", false,
 		"Enables health check endpoints /startup, /liveness, and /readiness on localhost.")
 	pflags.StringVar(&c.conf.APIEndpointURL, "sqladmin-api-endpoint", "",
@@ -712,6 +730,22 @@ func runSignalWrapper(cmd *Command) error {
 		notify = hc.NotifyStarted
 	}
 
+	go func() {
+		if !cmd.conf.Debug {
+			return
+		}
+		m := http.NewServeMux()
+		m.HandleFunc("/debug/pprof/", pprof.Index)
+		m.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+		m.HandleFunc("/debug/pprof/profile", pprof.Profile)
+		m.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+		m.HandleFunc("/debug/pprof/trace", pprof.Trace)
+		addr := net.JoinHostPort("localhost", cmd.conf.AdminPort)
+		cmd.logger.Infof("Starting admin server on %v", addr)
+		if lErr := http.ListenAndServe(addr, m); lErr != nil {
+			cmd.logger.Errorf("Failed to start admin HTTP server: %v", lErr)
+		}
+	}()
 	// Start the HTTP server if anything requiring HTTP is specified.
 	if needsHTTPServer {
 		server := &http.Server{

cmd/root_test.go

@@ -56,6 +56,9 @@ func withDefaults(c *proxy.Config) *proxy.Config {
 	if c.HTTPPort == "" {
 		c.HTTPPort = "9090"
 	}
+	if c.AdminPort == "" {
+		c.AdminPort = "9091"
+	}
 	if c.TelemetryTracingSampleRate == 0 {
 		c.TelemetryTracingSampleRate = 10_000
 	}
@@ -359,6 +362,20 @@ func TestNewCommandArguments(t *testing.T) {
 				ImpersonationChain: "[email protected]",
 			}),
 		},
+		{
+			desc: "using the debug flag",
+			args: []string{"--debug", "proj:region:inst"},
+			want: withDefaults(&proxy.Config{
+				Debug: true,
+			}),
+		},
+		{
+			desc: "using the admin port flag",
+			args: []string{"--admin-port", "7777", "proj:region:inst"},
+			want: withDefaults(&proxy.Config{
+				AdminPort: "7777",
+			}),
+		},
 	}
 
 	for _, tc := range tcs {
@@ -698,6 +715,22 @@ func TestNewCommandWithEnvironmentConfig(t *testing.T) {
 				HTTPPort: "5555",
 			}),
 		},
+		{
+			desc:     "using the debug envvar",
+			envName:  "CSQL_PROXY_DEBUG",
+			envValue: "true",
+			want: withDefaults(&proxy.Config{
+				Debug: true,
+			}),
+		},
+		{
+			desc:     "using the admin port envvar",
+			envName:  "CSQL_PROXY_ADMIN_PORT",
+			envValue: "7777",
+			want: withDefaults(&proxy.Config{
+				AdminPort: "7777",
+			}),
+		},
 	}
 	for _, tc := range tcs {
 		t.Run(tc.desc, func(t *testing.T) {
@@ -1010,6 +1043,26 @@ func TestCommandWithCustomDialer(t *testing.T) {
 	}
 }
 
+func tryDial(addr string) (*http.Response, error) {
+	var (
+		resp     *http.Response
+		attempts int
+		err      error
+	)
+	for {
+		if attempts > 10 {
+			return resp, err
+		}
+		resp, err = http.Get(addr)
+		if err != nil {
+			attempts++
+			time.Sleep(time.Second)
+			continue
+		}
+		return resp, err
+	}
+}
+
 func TestPrometheusMetricsEndpoint(t *testing.T) {
 	c := NewCommand(WithDialer(&spyDialer{}))
 	// Keep the test output quiet
@@ -1024,25 +1077,6 @@ func TestPrometheusMetricsEndpoint(t *testing.T) {
 
 	// try to dial metrics server for a max of ~10s to give the proxy time to
 	// start up.
-	tryDial := func(addr string) (*http.Response, error) {
-		var (
-			resp     *http.Response
-			attempts int
-			err      error
-		)
-		for {
-			if attempts > 10 {
-				return resp, err
-			}
-			resp, err = http.Get(addr)
-			if err != nil {
-				attempts++
-				time.Sleep(time.Second)
-				continue
-			}
-			return resp, err
-		}
-	}
 	resp, err := tryDial("http://localhost:9090/metrics") // default port set by http-port flag
 	if err != nil {
 		t.Fatalf("failed to dial metrics endpoint: %v", err)
@@ -1051,3 +1085,21 @@ func TestPrometheusMetricsEndpoint(t *testing.T) {
 		t.Fatalf("expected a 200 status, got = %v", resp.StatusCode)
 	}
 }
+
+func TestPProfServer(t *testing.T) {
+	c := NewCommand(WithDialer(&spyDialer{}))
+	c.SilenceUsage = true
+	c.SilenceErrors = true
+	c.SetArgs([]string{"--debug", "--admin-port", "9191", "my-project:my-region:my-instance"})
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go c.ExecuteContext(ctx)
+	resp, err := tryDial("http://localhost:9191/debug/pprof/")
+	if err != nil {
+		t.Fatalf("failed to dial endpoint: %v", err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("expected a 200 status, got = %v", resp.StatusCode)
+	}
+}

internal/proxy/proxy.go

@@ -202,6 +202,11 @@ type Config struct {
 	HTTPAddress string
 	// HTTPPort sets the port for the health check and prometheus server.
 	HTTPPort string
+	// AdminPort configures the port for the localhost-only admin server.
+	AdminPort string
+
+	// Debug enables a debug handler on localhost.
+	Debug bool
 
 	// OtherUserAgents is a list of space separate user agents that will be
 	// appended to the default user agent.