Use IntOrString for TCP health check ports

Clean up code to be more testable. Add test cases for named and numeric ports in TCP health checks. Improve tests.
duda · Aug 11, 2014 · c67c1ed · c67c1ed
1 parent 7201227
commit c67c1ed
Show file tree

Hide file tree

Showing 4 changed files with 99 additions and 19 deletions.
diff --git a/pkg/api/types.go b/pkg/api/types.go
@@ -146,8 +146,8 @@ type HTTPGetProbe struct {
 
 // TCPSocketProbe describes a liveness probe based on opening a socket
 type TCPSocketProbe struct {
-	// Port is the port to connect to. Required.
-	Port int `yaml:"port,omitempty" json:"port,omitempty"`
+	// Required: Port to connect to.
+	Port util.IntOrString `yaml:"port,omitempty" json:"port,omitempty"`
 }
 
 // LivenessProbe describes a liveness probe to be examined to the container.

diff --git a/pkg/api/v1beta1/types.go b/pkg/api/v1beta1/types.go
@@ -149,8 +149,8 @@ type HTTPGetProbe struct {
 
 // TCPSocketProbe describes a liveness probe based on opening a socket
 type TCPSocketProbe struct {
-	// Port is the port to connect to. Required.
-	Port int `yaml:"port,omitempty" json:"port,omitempty"`
+	// Required: Port to connect to.
+	Port util.IntOrString `yaml:"port,omitempty" json:"port,omitempty"`
 }
 
 // LivenessProbe describes a liveness probe to be examined to the container.

diff --git a/pkg/health/health_check.go b/pkg/health/health_check.go
@@ -126,15 +126,41 @@ func (h *HTTPHealthChecker) HealthCheck(currentState api.PodState, container api
 
 type TCPHealthChecker struct{}
 
-func (t *TCPHealthChecker) HealthCheck(currentState api.PodState, container api.Container) (Status, error) {
+// Get the components of a TCP connection address.  For testability.
+func getTCPAddrParts(currentState api.PodState, container api.Container) (string, int, error) {
 	params := container.LivenessProbe.TCPSocket
 	if params == nil {
-		return Unknown, fmt.Errorf("error, no TCP parameters specified: %v", container)
+		return "", -1, fmt.Errorf("error, no TCP parameters specified: %v", container)
+	}
+	port := -1
+	switch params.Port.Kind {
+	case util.IntstrInt:
+		port = params.Port.IntVal
+	case util.IntstrString:
+		port = findPortByName(container, params.Port.StrVal)
+		if port == -1 {
+			// Last ditch effort - maybe it was an int stored as string?
+			var err error
+			if port, err = strconv.Atoi(params.Port.StrVal); err != nil {
+				return "", -1, err
+			}
+		}
+	}
+	if port == -1 {
+		return "", -1, fmt.Errorf("unknown port: %v", params.Port)
 	}
 	if len(currentState.PodIP) == 0 {
-		return Unknown, fmt.Errorf("no host specified.")
+		return "", -1, fmt.Errorf("no host specified.")
 	}
-	conn, err := net.Dial("tcp", net.JoinHostPort(currentState.PodIP, strconv.Itoa(params.Port)))
+
+	return currentState.PodIP, port, nil
+}
+
+// DoTCPCheck checks that a TCP socket to the address can be opened.
+// If the socket can be opened, it returns Healthy.
+// If the socket fails to open, it returns Unhealthy.
+func DoTCPCheck(addr string) (Status, error) {
+	conn, err := net.Dial("tcp", addr)
 	if err != nil {
 		return Unhealthy, nil
 	}
@@ -144,3 +170,11 @@ func (t *TCPHealthChecker) HealthCheck(currentState api.PodState, container api.
 	}
 	return Healthy, nil
 }
+
+func (t *TCPHealthChecker) HealthCheck(currentState api.PodState, container api.Container) (Status, error) {
+	host, port, err := getTCPAddrParts(currentState, container)
+	if err != nil {
+		return Unknown, err
+	}
+	return DoTCPCheck(net.JoinHostPort(host, strconv.Itoa(port)))
+}
diff --git a/pkg/health/health_check_test.go b/pkg/health/health_check_test.go
@@ -21,7 +21,6 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/url"
-	"strconv"
 	"testing"
 
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
@@ -209,11 +208,56 @@ func TestHTTPHealthChecker(t *testing.T) {
 	}
 }
 
+func TestGetTCPAddrParts(t *testing.T) {
+	testCases := []struct {
+		probe *api.TCPSocketProbe
+		ok    bool
+		host  string
+		port  int
+	}{
+		{&api.TCPSocketProbe{Port: util.MakeIntOrStringFromInt(-1)}, false, "", -1},
+		{&api.TCPSocketProbe{Port: util.MakeIntOrStringFromString("")}, false, "", -1},
+		{&api.TCPSocketProbe{Port: util.MakeIntOrStringFromString("-1")}, false, "", -1},
+		{&api.TCPSocketProbe{Port: util.MakeIntOrStringFromString("not-found")}, false, "", -1},
+		{&api.TCPSocketProbe{Port: util.MakeIntOrStringFromString("found")}, true, "1.2.3.4", 93},
+		{&api.TCPSocketProbe{Port: util.MakeIntOrStringFromInt(76)}, true, "1.2.3.4", 76},
+		{&api.TCPSocketProbe{Port: util.MakeIntOrStringFromString("118")}, true, "1.2.3.4", 118},
+	}
+
+	for _, test := range testCases {
+		state := api.PodState{PodIP: "1.2.3.4"}
+		container := api.Container{
+			Ports: []api.Port{{Name: "found", HostPort: 93}},
+			LivenessProbe: &api.LivenessProbe{
+				TCPSocket: test.probe,
+				Type:      "tcp",
+			},
+		}
+		host, port, err := getTCPAddrParts(state, container)
+		if !test.ok && err == nil {
+			t.Errorf("Expected error for %+v, got %s:%d", test, host, port)
+		}
+		if test.ok && err != nil {
+			t.Errorf("Unexpected error: %v", err)
+		}
+		if test.ok {
+			if host != test.host || port != test.port {
+				t.Errorf("Expected %s:%d, got %s:%d", test.host, test.port, host, port)
+			}
+		}
+	}
+}
+
 func TestTcpHealthChecker(t *testing.T) {
-	type tcpHealthTest struct {
-		probe          *api.LivenessProbe
+	tests := []struct {
+		probe          *api.TCPSocketProbe
 		expectedStatus Status
 		expectError    bool
+	}{
+		// The probe will be filled in below.  This is primarily testing that a connection is made.
+		{&api.TCPSocketProbe{}, Healthy, false},
+		{&api.TCPSocketProbe{}, Unhealthy, false},
+		{nil, Unknown, true},
 	}
 
 	checker := &TCPHealthChecker{}
@@ -225,17 +269,19 @@ func TestTcpHealthChecker(t *testing.T) {
 		t.Errorf("Unexpected error: %v", err)
 	}
 	host, port, err := net.SplitHostPort(u.Host)
-	portNum, _ := strconv.Atoi(port)
-
-	tests := []tcpHealthTest{
-		{&api.LivenessProbe{TCPSocket: &api.TCPSocketProbe{Port: portNum}}, Healthy, false},
-		{&api.LivenessProbe{TCPSocket: &api.TCPSocketProbe{Port: 100000}}, Unhealthy, false},
-		{&api.LivenessProbe{}, Unknown, true},
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
 	}
 	for _, test := range tests {
-		probe := test.probe
 		container := api.Container{
-			LivenessProbe: probe,
+			LivenessProbe: &api.LivenessProbe{
+				TCPSocket: test.probe,
+				Type:      "tcp",
+			},
+		}
+		params := container.LivenessProbe.TCPSocket
+		if params != nil && test.expectedStatus == Healthy {
+			params.Port = util.MakeIntOrStringFromString(port)
 		}
 		status, err := checker.HealthCheck(api.PodState{PodIP: host}, container)
 		if status != test.expectedStatus {