feat: add support for Cloud Monitoring and Cloud Trace (#1212)

Users may now enable metrics and traces from the Go connector by
supplying a "telemetry-project" and project ID. Optionally, users may
also set a metrics prefix.

Separately, this commit separates metrics and tracing, allowing each to
be enabled on their own. In addition, it adds support for configuring the default
probabilistic sampling rate and otherwise uses the default where
1/10,000 calls will be traced.

Author: enocom

cmd/root.go

@@ -31,10 +31,12 @@ import (
 
 	"cloud.google.com/go/cloudsqlconn"
 	"contrib.go.opencensus.io/exporter/prometheus"
+	"contrib.go.opencensus.io/exporter/stackdriver"
 	"github.com/GoogleCloudPlatform/cloudsql-proxy/v2/cloudsql"
 	"github.com/GoogleCloudPlatform/cloudsql-proxy/v2/internal/gcloud"
 	"github.com/GoogleCloudPlatform/cloudsql-proxy/v2/internal/proxy"
 	"github.com/spf13/cobra"
+	"go.opencensus.io/trace"
 	"golang.org/x/oauth2"
 )
 
@@ -67,8 +69,13 @@ type Command struct {
 	*cobra.Command
 	conf *proxy.Config
 
-	prometheusNamespace string
-	httpPort            string
+	disableTraces              bool
+	telemetryTracingSampleRate int
+	disableMetrics             bool
+	telemetryProject           string
+	telemetryPrefix            string
+	prometheusNamespace        string
+	httpPort                   string
 }
 
 // Option is a function that configures a Command.
@@ -120,6 +127,16 @@ 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.telemetryProject, "telemetry-project", "",
+		"Enable Cloud Monitoring and Cloud Trace integration with the provided project ID.")
+	cmd.PersistentFlags().BoolVar(&c.disableTraces, "disable-traces", false,
+		"Disable Cloud Trace integration (used with telemetry-project)")
+	cmd.PersistentFlags().IntVar(&c.telemetryTracingSampleRate, "telemetry-sample-rate", 10_000,
+		"Configure the denominator of the probabilistic sample rate of traces sent to Cloud Trace\n(e.g., 10,000 traces 1/10,000 calls).")
+	cmd.PersistentFlags().BoolVar(&c.disableMetrics, "disable-metrics", false,
+		"Disable Cloud Monitoring integration (used with telemetry-project)")
+	cmd.PersistentFlags().StringVar(&c.telemetryPrefix, "telemetry-prefix", "",
+		"Prefix to use for Cloud Monitoring metrics.")
 	cmd.PersistentFlags().StringVar(&c.prometheusNamespace, "prometheus-namespace", "",
 		"Enable Prometheus for metric collection using the provided namespace")
 	cmd.PersistentFlags().StringVar(&c.httpPort, "http-port", "9090",
@@ -196,6 +213,16 @@ func parseConfig(cmd *cobra.Command, conf *proxy.Config, args []string) error {
 		return newBadCommandError("cannot specify --http-port without --prometheus-namespace")
 	}
 
+	if !userHasSet("telemetry-project") && userHasSet("telemetry-prefix") {
+		cmd.Println("Ignoring telementry-prefix as telemetry-project was not set")
+	}
+	if !userHasSet("telemetry-project") && userHasSet("disable-metrics") {
+		cmd.Println("Ignoring disable-metrics as telemetry-project was not set")
+	}
+	if !userHasSet("telemetry-project") && userHasSet("disable-traces") {
+		cmd.Println("Ignoring disable-traces as telemetry-project was not set")
+	}
+
 	var ics []proxy.InstanceConnConfig
 	for _, a := range args {
 		// Assume no query params initially
@@ -268,6 +295,36 @@ func runSignalWrapper(cmd *Command) error {
 	ctx, cancel := context.WithCancel(cmd.Context())
 	defer cancel()
 
+	// Configure Cloud Trace and/or Cloud Monitoring based on command
+	// invocation. If a project has not been enabled, no traces or metrics are
+	// enabled.
+	enableMetrics := !cmd.disableMetrics
+	enableTraces := !cmd.disableTraces
+	if cmd.telemetryProject != "" && (enableMetrics || enableTraces) {
+		sd, err := stackdriver.NewExporter(stackdriver.Options{
+			ProjectID:    cmd.telemetryProject,
+			MetricPrefix: cmd.telemetryPrefix,
+		})
+		if err != nil {
+			return err
+		}
+		if enableMetrics {
+			err = sd.StartMetricsExporter()
+			if err != nil {
+				return err
+			}
+		}
+		if enableTraces {
+			s := trace.ProbabilitySampler(1 / float64(cmd.telemetryTracingSampleRate))
+			trace.ApplyConfig(trace.Config{DefaultSampler: s})
+			trace.RegisterExporter(sd)
+		}
+		defer func() {
+			sd.Flush()
+			sd.StopMetricsExporter()
+		}()
+	}
+
 	shutdownCh := make(chan error)
 
 	if cmd.prometheusNamespace != "" {

go.mod

@@ -6,6 +6,7 @@ require (
 	cloud.google.com/go/cloudsqlconn v0.2.1-0.20220401153611-87e713b37755
 	cloud.google.com/go/compute v1.7.0
 	contrib.go.opencensus.io/exporter/prometheus v0.4.1
+	contrib.go.opencensus.io/exporter/stackdriver v0.13.13
 	github.com/GoogleCloudPlatform/cloudsql-proxy v1.29.0
 	github.com/coreos/go-systemd/v22 v22.3.2
 	github.com/denisenkom/go-mssqldb v0.12.2
@@ -15,6 +16,7 @@ require (
 	github.com/jackc/pgx/v4 v4.16.1
 	github.com/lib/pq v1.10.6
 	github.com/spf13/cobra v1.2.1
+	go.opencensus.io v0.23.0
 	go.uber.org/zap v1.21.0
 	golang.org/x/crypto v0.0.0-20220507011949-2cf3adece122 // indirect
 	golang.org/x/net v0.0.0-20220624214902-1bab6f366d9e