feat: add support for Prometheus metrics

cmd/root.go

@@ -20,14 +20,17 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"net/http"
 	"net/url"
 	"os"
 	"os/signal"
 	"strconv"
 	"strings"
 	"syscall"
+	"time"
 
 	"cloud.google.com/go/alloydbconn"
+	"contrib.go.opencensus.io/exporter/prometheus"
 	"github.com/GoogleCloudPlatform/alloydb-auth-proxy/alloydb"
 	"github.com/GoogleCloudPlatform/alloydb-auth-proxy/internal/gcloud"
 	"github.com/GoogleCloudPlatform/alloydb-auth-proxy/internal/proxy"
@@ -75,6 +78,9 @@ func Execute() {
 type Command struct {
 	*cobra.Command
 	conf *proxy.Config
+
+	prometheusNamespace string
+	httpPort            string
 }
 
 // Option is a function that configures a Command.
@@ -126,6 +132,10 @@ without having to manage 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().StringVar(&c.prometheusNamespace, "prometheus-namespace", "",
+		"Enable Prometheus for metric collection using the provided namespace")
+	cmd.PersistentFlags().StringVar(&c.httpPort, "http-port", "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",
@@ -194,6 +204,10 @@ func parseConfig(cmd *cobra.Command, conf *proxy.Config, args []string) error {
 	}
 	conf.DialerOpts = opts
 
+	if userHasSet("http-port") && !userHasSet("prometheus-namespace") {
+		return newBadCommandError("cannot specify --http-port without --prometheus-namespace")
+	}
+
 	var ics []proxy.InstanceConnConfig
 	for _, a := range args {
 		// Assume no query params initially
@@ -268,14 +282,46 @@ 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.httpPort)
+		server := &http.Server{Addr: addr, Handler: mux}
+		go func() {
+			select {
+			case <-ctx.Done():
+				// Give the HTTP server a second to shutdown cleanly.
+				ctx2, _ := context.WithTimeout(context.Background(), time.Second)
+				if err := server.Shutdown(ctx2); err != nil {
+					cmd.Printf("failed to shutdown Prometheus HTTP server: %v\n", err)
+				}
+			}
+		}()
+		go func() {
+			err := server.ListenAndServe()
+			if err == http.ErrServerClosed {
+				return
+			}
+			if 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)
 	go func() {
 		var s os.Signal
 		select {
 		case s = <-signals:
-		case <-cmd.Context().Done():
+		case <-ctx.Done():
 			// this should only happen when the context supplied in tests in canceled
 			s = syscall.SIGINT
 		}

cmd/root_test.go

@@ -19,6 +19,7 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"net/http"
 	"sync"
 	"testing"
 	"time"
@@ -160,6 +161,10 @@ func TestNewCommandArguments(t *testing.T) {
 				}},
 			}),
 		},
+		{
+			desc: "enabling a Prometheus port without a namespace",
+			args: []string{"--htto-port", "1111", "proj:region:inst"},
+		},
 	}
 
 	for _, tc := range tcs {
@@ -411,3 +416,47 @@ func TestCommandWithCustomDialer(t *testing.T) {
 		return nil
 	}, 10)
 }
+
+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",
+		"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:9090/metrics") // default port set by http-port flag
+	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

@@ -4,7 +4,7 @@ go 1.16
 
 require (
 	cloud.google.com/go/alloydbconn v0.1.1
-	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
+	contrib.go.opencensus.io/exporter/prometheus v0.4.1
 	github.com/google/go-cmp v0.5.8
 	github.com/lib/pq v1.10.5 // indirect
 	github.com/spf13/cobra v1.5.0