Merge pull request #23463 from asubiotto/haproxy-cfg

cli: update generated haproxy config with readiness endpoint

pkg/cli/haproxy.go

@@ -16,11 +16,19 @@ package cli
 
 import (
 	"context"
+	"flag"
 	"html/template"
 	"io"
 	"os"
+	"strings"
 
+	"io/ioutil"
+
+	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/cli/cliflags"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
+	"github.com/cockroachdb/cockroach/pkg/server/status"
 	"github.com/spf13/cobra"
 )
 
@@ -40,6 +48,41 @@ running haproxy in the same network.
 	RunE: MaybeDecorateGRPCError(runGenHAProxyCmd),
 }
 
+type haProxyNodeInfo struct {
+	NodeID   roachpb.NodeID
+	NodeAddr string
+	// The port on which health checks are performed.
+	CheckPort string
+}
+
+func nodeStatusesToNodeInfos(statuses []status.NodeStatus) []haProxyNodeInfo {
+	fs := flag.NewFlagSet("haproxy", flag.ContinueOnError)
+	checkPort := fs.String(cliflags.ServerHTTPPort.Name, base.DefaultHTTPPort, "" /* usage */)
+
+	// Discard parsing output.
+	fs.SetOutput(ioutil.Discard)
+
+	nodeInfos := make([]haProxyNodeInfo, len(statuses))
+	for i, status := range statuses {
+		nodeInfos[i].NodeID = status.Desc.NodeID
+		nodeInfos[i].NodeAddr = status.Desc.Address.AddressField
+
+		*checkPort = base.DefaultHTTPPort
+		// Iterate over the arguments until the ServerHTTPPort flag is found and
+		// parse the remainder of the arguments. This is done because Parse returns
+		// when it encounters an undefined flag and we do not want to define all
+		// possible flags.
+		for i, arg := range status.Args {
+			if strings.Contains(arg, cliflags.ServerHTTPPort.Name) {
+				_ = fs.Parse(status.Args[i:])
+			}
+		}
+
+		nodeInfos[i].CheckPort = *checkPort
+	}
+	return nodeInfos
+}
+
 func runGenHAProxyCmd(cmd *cobra.Command, args []string) error {
 	if len(args) > 0 {
 		return usageAndError(cmd)
@@ -75,7 +118,7 @@ func runGenHAProxyCmd(cmd *cobra.Command, args []string) error {
 		w = f
 	}
 
-	err = configTemplate.Execute(w, nodeStatuses.Nodes)
+	err = configTemplate.Execute(w, nodeStatusesToNodeInfos(nodeStatuses.Nodes))
 	if err != nil {
 		// Return earliest error, but still close the file.
 		_ = f.Close()
@@ -107,6 +150,7 @@ listen psql
     bind :26257
     mode tcp
     balance roundrobin
-{{range .}}    server cockroach{{.Desc.NodeID}} {{.Desc.Address.AddressField}} check
+    option httpchk GET /health?ready=1
+{{range .}}    server cockroach{{.NodeID}} {{.NodeAddr}} check port {{.CheckPort}}
 {{end}}
 `

pkg/cli/haproxy_test.go

@@ -0,0 +1,108 @@
+// Copyright 2017 The Cockroach Authors.
+//
+// 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 cli
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/server/status"
+	"github.com/cockroachdb/cockroach/pkg/util"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+)
+
+func TestNodeStatusToNodeInfoConversion(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	testCases := []struct {
+		input    []status.NodeStatus
+		expected []haProxyNodeInfo
+	}{
+		{
+			[]status.NodeStatus{
+				{
+					Desc: roachpb.NodeDescriptor{
+						NodeID: 1,
+						Address: util.UnresolvedAddr{
+							AddressField: "addr",
+						},
+					},
+					// Flags but no http port.
+					Args: []string{"--unwanted", "-unwanted"},
+				},
+			},
+			[]haProxyNodeInfo{
+				{
+					NodeID:    1,
+					NodeAddr:  "addr",
+					CheckPort: base.DefaultHTTPPort,
+				},
+			},
+		},
+		{
+			[]status.NodeStatus{
+				{
+					Args: []string{"--unwanted", "-http-port=1234"},
+				},
+				{
+					Args: nil,
+				},
+			},
+			[]haProxyNodeInfo{
+				{
+					CheckPort: "1234",
+				},
+				{
+					CheckPort: base.DefaultHTTPPort,
+				},
+			},
+		},
+		{
+			[]status.NodeStatus{
+				{
+					Args: []string{"--http-port", "5678", "--unwanted"},
+				},
+			},
+			[]haProxyNodeInfo{
+				{
+					CheckPort: "5678",
+				},
+			},
+		},
+		{
+			[]status.NodeStatus{
+				{
+					// We shouldn't see this, because the flag needs an argument on startup,
+					// but check that we fall back to the default port.
+					Args: []string{"-http-port"},
+				},
+			},
+			[]haProxyNodeInfo{
+				{
+					CheckPort: base.DefaultHTTPPort,
+				},
+			},
+		},
+	}
+
+	for _, testCase := range testCases {
+		output := nodeStatusesToNodeInfos(testCase.input)
+		if !reflect.DeepEqual(output, testCase.expected) {
+			t.Fatalf("unexpected output %v, expected %v", output, testCase.expected)
+		}
+	}
+}