Add /health endpoint for ndt-server (#376)

* Add /health endpoint for NDT

* Add comments

* Refactor

* Add tests

* Fix wording

ndt-server.go

@@ -38,6 +38,7 @@ var (
 	ndt5Addr          = flag.String("ndt5_addr", ":3001", "The address and port to use for the unencrypted ndt5 test")
 	ndt5WsAddr        = flag.String("ndt5_ws_addr", "127.0.0.1:3002", "The address and port to use for the ndt5 WS test")
 	ndt5WssAddr       = flag.String("ndt5_wss_addr", ":3010", "The address and port to use for the ndt5 WSS test")
+	healthAddr        = flag.String("health_addr", ":8000", "The address and port to use for health checks")
 	certFile          = flag.String("cert", "", "The file with server certificates in PEM format.")
 	keyFile           = flag.String("key", "", "The file with server key in PEM format.")
 	tlsVersion        = flag.String("tls.version", "", "Minimum TLS version. Valid values: 1.2 or 1.3")
@@ -47,6 +48,7 @@ var (
 	tokenVerifyKey    = flagx.FileBytesArray{}
 	tokenRequired5    bool
 	tokenRequired7    bool
+	isLameDuck        bool
 	tokenMachine      string
 
 	// A metric to use to signal that the server is in lame duck mode.
@@ -69,7 +71,7 @@ func init() {
 
 func catchSigterm() {
 	// Disable lame duck status.
-	lameDuck.Set(0)
+	setLameDuck(0)
 
 	// Register channel to receive SIGTERM events.
 	c := make(chan os.Signal, 1)
@@ -84,7 +86,7 @@ func catchSigterm() {
 		fmt.Println("Canceled")
 	}
 	// Set lame duck status. This will remain set until exit.
-	lameDuck.Set(1)
+	setLameDuck(1)
 	// When we receive a second SIGTERM, cancel the context and shut everything
 	// down. This should cause main() to exit cleanly.
 	select {
@@ -150,6 +152,22 @@ func parseDeploymentLabels() []metadata.NameValue {
 	return serverMetadata
 }
 
+// Set internal lame duck status and metric.
+func setLameDuck(status float64) {
+	isLameDuck = status != 0
+	lameDuck.Set(status)
+}
+
+// Handle requests to the /health endpoint.
+// Writes out a 200 status code only if the server is not in lame duck mode.
+func handleHealth(rw http.ResponseWriter, req *http.Request) {
+	if isLameDuck {
+		rw.WriteHeader(http.StatusInternalServerError)
+		return
+	}
+	rw.WriteHeader(http.StatusOK)
+}
+
 func main() {
 	flag.Parse()
 	rtx.Must(flagx.ArgsFromEnv(flag.CommandLine), "Could not parse env args")
@@ -259,6 +277,16 @@ func main() {
 		log.Printf("Cert=%q and Key=%q means no TLS services will be started.\n", *certFile, *keyFile)
 	}
 
+	// Set up handler for /health endpoint.
+	healthMux := http.NewServeMux()
+	healthMux.Handle("/health", http.HandlerFunc(handleHealth))
+	healthServer := httpServer(
+		*healthAddr,
+		healthMux,
+	)
+	rtx.Must(listener.ListenAndServeAsync(healthServer), "Could not start health server")
+	defer healthServer.Close()
+
 	// Serve until the context is canceled.
 	<-ctx.Done()
 }

ndt-server_test.go

@@ -382,3 +382,33 @@ func sortNameValueSlice(nv []metadata.NameValue) {
 		return nv[i].Name < nv[j].Name
 	})
 }
+
+func Test_handleHealth(t *testing.T) {
+	tests := []struct {
+		name     string
+		lameDuck bool
+		want     int
+	}{
+		{
+			name:     "not-lame-duck",
+			lameDuck: false,
+			want:     200,
+		},
+		{
+			name:     "lame-duck",
+			lameDuck: true,
+			want:     500,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			isLameDuck = tt.lameDuck
+			writer := httptest.NewRecorder()
+			handleHealth(writer, nil)
+
+			if writer.Code != tt.want {
+				t.Errorf("ndt-server.handleHealth() got = %d, want %d", writer.Code, tt.want)
+			}
+		})
+	}
+}