feat: add admin server with pprof (#236)

This is a port of GoogleCloudPlatform/cloud-sql-proxy#1564.

README.md

@@ -277,6 +277,21 @@ when you want to proxy this traffic. Otherwise, it is optional. See
 [`http.ProxyFromEnvironment`](https://pkg.go.dev/net/http@go1.17.3#ProxyFromEnvironment)
 for possible values.
 
+## 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.
+
+
 ## Support policy
 
 ### Major version lifecycle

cmd/root.go

@@ -22,6 +22,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"net/http/pprof"
 	"net/url"
 	"os"
 	"os/signal"
@@ -254,6 +255,18 @@ Configuration using environment variables
         ALLOYDB_PROXY_INSTANCE_URI_1=projects/PROJECT/locations/REGION/clusters/CLUSTER/instances/INSTANCE2 \
             ./alloydb-auth-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.
 `
 
 const envPrefix = "ALLOYDB_PROXY"
@@ -387,6 +400,10 @@ the maximum time has passed. Defaults to 0s.`)
 		"Address for Prometheus and health check server")
 	pflags.StringVar(&c.conf.HTTPPort, "http-port", "9090",
 		"Port for the Prometheus server to use")
+	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 HTTP endpoints /startup, /liveness, and /readiness
 that report on the proxy's health. Endpoints are available on localhost
@@ -682,6 +699,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

@@ -78,6 +78,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
 	}
@@ -296,6 +299,22 @@ func TestNewCommandArguments(t *testing.T) {
 				ImpersonationChain: "sv1@developer.gserviceaccount.com",
 			}),
 		},
+		{
+			desc: "using the debug flag",
+			args: []string{"--debug",
+				"projects/proj/locations/region/clusters/clust/instances/inst"},
+			want: withDefaults(&proxy.Config{
+				Debug: true,
+			}),
+		},
+		{
+			desc: "using the admin port flag",
+			args: []string{"--admin-port", "7777",
+				"projects/proj/locations/region/clusters/clust/instances/inst"},
+			want: withDefaults(&proxy.Config{
+				AdminPort: "7777",
+			}),
+		},
 	}
 
 	for _, tc := range tcs {
@@ -495,6 +514,22 @@ func TestNewCommandWithEnvironmentConfig(t *testing.T) {
 				HTTPPort: "5555",
 			}),
 		},
+		{
+			desc:     "using the debug envvar",
+			envName:  "ALLOYDB_PROXY_DEBUG",
+			envValue: "true",
+			want: withDefaults(&proxy.Config{
+				Debug: true,
+			}),
+		},
+		{
+			desc:     "using the admin port envvar",
+			envName:  "ALLOYDB_PROXY_ADMIN_PORT",
+			envValue: "7777",
+			want: withDefaults(&proxy.Config{
+				AdminPort: "7777",
+			}),
+		},
 	}
 	for _, tc := range tcs {
 		t.Run(tc.desc, func(t *testing.T) {
@@ -843,6 +878,26 @@ func TestCommandWithCustomDialer(t *testing.T) {
 	}, 10)
 }
 
+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
@@ -857,25 +912,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)
@@ -884,3 +920,22 @@ 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",
+		"projects/proj/locations/region/clusters/clust/instances/inst"})
+	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

@@ -147,6 +147,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.