feat: add support for Prometheus metrics

cmd/root.go

@@ -20,6 +20,7 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"net/http"
 	"net/url"
 	"os"
 	"os/signal"
@@ -28,6 +29,7 @@ import (
 	"syscall"
 
 	"cloud.google.com/go/cloudsqlconn"
+	"contrib.go.opencensus.io/exporter/prometheus"
 	"github.com/GoogleCloudPlatform/cloudsql-proxy/v2/cloudsql"
 	"github.com/GoogleCloudPlatform/cloudsql-proxy/v2/internal/gcloud"
 	"github.com/GoogleCloudPlatform/cloudsql-proxy/v2/internal/proxy"
@@ -63,6 +65,9 @@ func Execute() {
 type Command struct {
 	*cobra.Command
 	conf *proxy.Config
+
+	prometheusNamespace string
+	prometheusPort      string
 }
 
 // Option is a function that configures a Command.
@@ -114,6 +119,10 @@ any client SSL certificates.`,
 		"Path to a service account key to use for authentication.")
 	cmd.PersistentFlags().BoolVarP(&c.conf.GcloudAuth, "gcloud-auth", "g", false,
 		"Use gcloud's user configuration to retrieve a token for authentication.")
+	cmd.PersistentFlags().StringVarP(&c.prometheusNamespace, "prometheus-namespace", "n", "",
+		"Enable Prometheus for metric collection using the provided namespace")
+	cmd.PersistentFlags().StringVarP(&c.prometheusPort, "prometheus-port", "s", "9090",
+		"Port for the Prometheus server to use")
 
 	// Global and per instance flags
 	cmd.PersistentFlags().StringVarP(&c.conf.Addr, "address", "a", "127.0.0.1",
@@ -182,6 +191,10 @@ func parseConfig(cmd *cobra.Command, conf *proxy.Config, args []string) error {
 	}
 	conf.DialerOpts = opts
 
+	if userHasSet("prometheus-port") && !userHasSet("prometheus-namespace") {
+		return newBadCommandError("cannot specify --prometheus-port without --prometheus-namespace")
+	}
+
 	var ics []proxy.InstanceConnConfig
 	for _, a := range args {
 		// Assume no query params initially
@@ -256,6 +269,32 @@ func runSignalWrapper(cmd *Command) error {
 
 	shutdownCh := make(chan error)
 
+	if cmd.prometheusNamespace != "" {
+		e, err := prometheus.NewExporter(prometheus.Options{
+			Namespace: cmd.prometheusNamespace,
+		})
+		if err != nil {
+			return err
+		}
+		mux := http.NewServeMux()
+		mux.Handle("/metrics", e)
+		addr := fmt.Sprintf("localhost:%s", cmd.prometheusPort)
+		server := &http.Server{Addr: addr, Handler: mux}
+		go func() {
+			select {
+			case <-ctx.Done():
+				if err := server.Shutdown(context.Background()); err != nil {
+					shutdownCh <- fmt.Errorf("failed to stop prometheus HTTP server: %v", err)
+				}
+			}
+		}()
+		go func() {
+			if err := server.ListenAndServe(); err != nil {
+				shutdownCh <- fmt.Errorf("failed to start prometheus HTTP server: %v", err)
+			}
+		}()
+	}
+
 	// watch for sigterm / sigint signals
 	signals := make(chan os.Signal, 1)
 	signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT)

cmd/root_test.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"errors"
 	"net"
+	"net/http"
 	"sync"
 	"testing"
 	"time"
@@ -280,6 +281,10 @@ func TestNewCommandWithErrors(t *testing.T) {
 			desc: "using the unix socket and port query params",
 			args: []string{"proj:region:inst?unix-socket=/path&port=5000"},
 		},
+		{
+			desc: "enabling a Prometheus port without a namespace",
+			args: []string{"--prometheus-port", "1111", "proj:region:inst"},
+		},
 	}
 
 	for _, tc := range tcs {
@@ -349,3 +354,48 @@ func TestCommandWithCustomDialer(t *testing.T) {
 		t.Fatalf("want = %v, got = %v", want, got)
 	}
 }
+
+func TestPrometheusMetricsEndpoint(t *testing.T) {
+	c := NewCommand(WithDialer(&spyDialer{}))
+	// Keep the test output quiet
+	c.SilenceUsage = true
+	c.SilenceErrors = true
+	c.SetArgs([]string{
+		"--prometheus-namespace", "prometheus",
+		"--prometheus-port", "9999",
+		"my-project:my-region:my-instance"})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	go c.ExecuteContext(ctx)
+
+	// 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:9999/metrics")
+	if err != nil {
+		t.Fatalf("failed to dial metrics endpoint: %v", err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("expected a 200 status, got = %v", resp.StatusCode)
+	}
+}

go.mod

@@ -5,6 +5,7 @@ go 1.16
 require (
 	cloud.google.com/go/cloudsqlconn v0.2.1-0.20220401153611-87e713b37755
 	cloud.google.com/go/compute v1.6.1
+	contrib.go.opencensus.io/exporter/prometheus v0.4.1 // indirect
 	github.com/GoogleCloudPlatform/cloudsql-proxy v1.29.0
 	github.com/coreos/go-systemd/v22 v22.3.2
 	github.com/denisenkom/go-mssqldb v0.12.0