feat: add support for health-check flag (GoogleCloudPlatform#85)

This is an adaptation of GoogleCloudPlatform/cloud-sql-proxy#1271.

Author: enocom

README.md

@@ -271,7 +271,7 @@ and from a AlloyDB instance. The `ALL_PROXY` environment variable supports
 `socks5h` protocol.
 
 The `HTTPS_PROXY` (or `HTTP_PROXY`) specifies the proxy for all HTTP(S) traffic
-to the SQL Admin API. Specifying `HTTPS_PROXY` or `HTTP_PROXY` is only necessary
+to the AlloyDB Admin API. Specifying `HTTPS_PROXY` or `HTTP_PROXY` is only necessary
 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.

cmd/root.go

@@ -32,6 +32,7 @@ import (
 	"contrib.go.opencensus.io/exporter/prometheus"
 	"contrib.go.opencensus.io/exporter/stackdriver"
 	"github.com/GoogleCloudPlatform/alloydb-auth-proxy/alloydb"
+	"github.com/GoogleCloudPlatform/alloydb-auth-proxy/internal/healthcheck"
 	"github.com/GoogleCloudPlatform/alloydb-auth-proxy/internal/log"
 	"github.com/GoogleCloudPlatform/alloydb-auth-proxy/internal/proxy"
 	"github.com/spf13/cobra"
@@ -88,6 +89,7 @@ type Command struct {
 	telemetryProject           string
 	telemetryPrefix            string
 	prometheusNamespace        string
+	healthCheck                bool
 	httpPort                   string
 }
 
@@ -186,6 +188,10 @@ the maximum time has passed. Defaults to 0s.`)
 		"Enable Prometheus for metric collection using the provided namespace")
 	cmd.PersistentFlags().StringVar(&c.httpPort, "http-port", "9090",
 		"Port for the Prometheus server to use")
+	cmd.PersistentFlags().BoolVar(&c.healthCheck, "health-check", false,
+		`Enables HTTP endpoints /startup, /liveness, and /readiness
+that report on the proxy's health. Endpoints are available on localhost
+only. Uses the port specified by the http-port flag.`)
 
 	// Global and per instance flags
 	cmd.PersistentFlags().StringVarP(&c.conf.Addr, "address", "a", "127.0.0.1",
@@ -241,18 +247,18 @@ func parseConfig(cmd *Command, conf *proxy.Config, args []string) error {
 		cmd.logger.Infof("Using API Endpoint %v", conf.APIEndpointURL)
 	}
 
-	if userHasSet("http-port") && !userHasSet("prometheus-namespace") {
-		return newBadCommandError("cannot specify --http-port without --prometheus-namespace")
+	if userHasSet("http-port") && !userHasSet("prometheus-namespace") && !userHasSet("health-check") {
+		cmd.logger.Infof("Ignoring --http-port because --prometheus-namespace or --health-check was not set")
 	}
 
 	if !userHasSet("telemetry-project") && userHasSet("telemetry-prefix") {
-		cmd.logger.Infof("Ignoring telementry-prefix as telemetry-project was not set")
+		cmd.logger.Infof("Ignoring --telementry-prefix as --telemetry-project was not set")
 	}
 	if !userHasSet("telemetry-project") && userHasSet("disable-metrics") {
-		cmd.logger.Infof("Ignoring disable-metrics as telemetry-project was not set")
+		cmd.logger.Infof("Ignoring --disable-metrics as --telemetry-project was not set")
 	}
 	if !userHasSet("telemetry-project") && userHasSet("disable-traces") {
-		cmd.logger.Infof("Ignoring disable-traces as telemetry-project was not set")
+		cmd.logger.Infof("Ignoring --disable-traces as --telemetry-project was not set")
 	}
 
 	var ics []proxy.InstanceConnConfig
@@ -328,9 +334,8 @@ 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.
+	// Configure collectors before the proxy has started to ensure we are
+	// collecting metrics before *ANY* AlloyDB Admin API calls are made.
 	enableMetrics := !cmd.disableMetrics
 	enableTraces := !cmd.disableTraces
 	if cmd.telemetryProject != "" && (enableMetrics || enableTraces) {
@@ -358,40 +363,22 @@ func runSignalWrapper(cmd *Command) error {
 		}()
 	}
 
-	shutdownCh := make(chan error)
-
+	var (
+		needsHTTPServer bool
+		mux             = http.NewServeMux()
+	)
 	if cmd.prometheusNamespace != "" {
+		needsHTTPServer = true
 		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.logger.Errorf("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)
-			}
-		}()
 	}
 
+	shutdownCh := make(chan error)
 	// watch for sigterm / sigint signals
 	signals := make(chan os.Signal, 1)
 	signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT)
@@ -429,18 +416,55 @@ func runSignalWrapper(cmd *Command) error {
 		cmd.logger.Errorf("The proxy has encountered a terminal error: %v", err)
 		return err
 	case p = <-startCh:
+		cmd.logger.Infof("The proxy has started successfully and is ready for new connections!")
 	}
-	cmd.logger.Infof("The proxy has started successfully and is ready for new connections!")
-	defer p.Close()
 	defer func() {
 		if cErr := p.Close(); cErr != nil {
 			cmd.logger.Errorf("error during shutdown: %v", cErr)
 		}
 	}()
 
-	go func() {
-		shutdownCh <- p.Serve(ctx)
-	}()
+	notify := func() {}
+	if cmd.healthCheck {
+		needsHTTPServer = true
+		hc := healthcheck.NewCheck(p, cmd.logger)
+		mux.HandleFunc("/startup", hc.HandleStartup)
+		mux.HandleFunc("/readiness", hc.HandleReadiness)
+		mux.HandleFunc("/liveness", hc.HandleLiveness)
+		notify = hc.NotifyStarted
+	}
+
+	// Start the HTTP server if anything requiring HTTP is specified.
+	if needsHTTPServer {
+		server := &http.Server{
+			Addr:    fmt.Sprintf("localhost:%s", cmd.httpPort),
+			Handler: mux,
+		}
+		// Start the HTTP server.
+		go func() {
+			err := server.ListenAndServe()
+			if err == http.ErrServerClosed {
+				return
+			}
+			if err != nil {
+				shutdownCh <- fmt.Errorf("failed to start HTTP server: %v", err)
+			}
+		}()
+		// Handle shutdown of the HTTP server gracefully.
+		go func() {
+			select {
+			case <-ctx.Done():
+				// Give the HTTP server a second to shutdown cleanly.
+				ctx2, cancel := context.WithTimeout(context.Background(), time.Second)
+				defer cancel()
+				if err := server.Shutdown(ctx2); err != nil {
+					cmd.logger.Errorf("failed to shutdown Prometheus HTTP server: %v\n", err)
+				}
+			}
+		}()
+	}
+
+	go func() { shutdownCh <- p.Serve(ctx, notify) }()
 
 	err := <-shutdownCh
 	switch {

cmd/root_test.go

@@ -350,10 +350,6 @@ func TestNewCommandWithErrors(t *testing.T) {
 			desc: "using the unix socket and port query params",
 			args: []string{"projects/proj/locations/region/clusters/clust/instances/inst?unix-socket=/path&port=5000"},
 		},
-		{
-			desc: "enabling a Prometheus port without a namespace",
-			args: []string{"--http-port", "1111", "proj:region:inst"},
-		},
 		{
 			desc: "using an invalid url for host flag",
 			args: []string{"--host", "https://invalid:url[/]", "proj:region:inst"},

internal/healthcheck/healthcheck.go

@@ -0,0 +1,110 @@
+// Copyright 2022 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package healthcheck tests and communicates the health of the AlloyDB Auth
+// proxy.
+package healthcheck
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"sync"
+
+	"github.com/GoogleCloudPlatform/alloydb-auth-proxy/alloydb"
+	"github.com/GoogleCloudPlatform/alloydb-auth-proxy/internal/proxy"
+)
+
+// Check provides HTTP handlers for use as healthchecks typically in a
+// Kubernetes context.
+type Check struct {
+	once    *sync.Once
+	started chan struct{}
+	proxy   *proxy.Client
+	logger  alloydb.Logger
+}
+
+// NewCheck is the initializer for Check.
+func NewCheck(p *proxy.Client, l alloydb.Logger) *Check {
+	return &Check{
+		once:    &sync.Once{},
+		started: make(chan struct{}),
+		proxy:   p,
+		logger:  l,
+	}
+}
+
+// NotifyStarted notifies the check that the proxy has started up successfully.
+func (c *Check) NotifyStarted() {
+	c.once.Do(func() { close(c.started) })
+}
+
+// HandleStartup reports whether the Check has been notified of startup.
+func (c *Check) HandleStartup(w http.ResponseWriter, _ *http.Request) {
+	select {
+	case <-c.started:
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte("ok"))
+	default:
+		w.WriteHeader(http.StatusServiceUnavailable)
+		w.Write([]byte("error"))
+	}
+}
+
+var errNotStarted = errors.New("proxy is not started")
+
+// HandleReadiness ensures the Check has been notified of successful startup,
+// that the proxy has not reached maximum connections, and that all connections
+// are healthy.
+func (c *Check) HandleReadiness(w http.ResponseWriter, _ *http.Request) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	select {
+	case <-c.started:
+	default:
+		c.logger.Errorf("[Health Check] Readiness failed: %v", errNotStarted)
+		w.WriteHeader(http.StatusServiceUnavailable)
+		w.Write([]byte(errNotStarted.Error()))
+		return
+	}
+
+	if open, max := c.proxy.ConnCount(); max > 0 && open == max {
+		err := fmt.Errorf("max connections reached (open = %v, max = %v)", open, max)
+		c.logger.Errorf("[Health Check] Readiness failed: %v", err)
+		w.WriteHeader(http.StatusServiceUnavailable)
+		w.Write([]byte(err.Error()))
+		return
+	}
+
+	err := c.proxy.CheckConnections(ctx)
+	if err != nil {
+		c.logger.Errorf("[Health Check] Readiness failed: %v", err)
+		w.WriteHeader(http.StatusServiceUnavailable)
+		w.Write([]byte(err.Error()))
+		return
+	}
+
+	w.WriteHeader(http.StatusOK)
+	w.Write([]byte("ok"))
+}
+
+// HandleLiveness indicates the process is up and responding to HTTP requests.
+// If this check fails (because it's not reachable), the process is in a bad
+// state and should be restarted.
+func (c *Check) HandleLiveness(w http.ResponseWriter, _ *http.Request) {
+	w.WriteHeader(http.StatusOK)
+	w.Write([]byte("ok"))
+}