state: Refactor Task IP sourcing

This commit refactors the Task IP sourcing code to gracefully account
for the new `NetworkInfo` source with the same structure and be easily
changed to distinguish between IPv4 and IPv6.

records/config.go

@@ -103,7 +103,7 @@ func NewConfig() Config {
 		HTTPOn:             true,
 		ExternalOn:         true,
 		RecurseOn:          true,
-		IPSources:          []string{"mesos", "host"},
+		IPSources:          []string{"netinfo", "mesos", "host"},
 	}
 }
 

records/state/state.go

@@ -2,6 +2,7 @@ package state
 
 import (
 	"bytes"
+	"net"
 	"strconv"
 	"strings"
 
@@ -78,66 +79,103 @@ type Task struct {
 	SlaveIP string `json:"-"`
 }
 
-var ipLabels = map[string]string{
-	"docker": "Docker.NetworkSettings.IPAddress",
-	"mesos":  "MesosContainerizer.NetworkSettings.IPAddress",
+// HasDiscoveryInfo return whether the DiscoveryInfo was provided in the state.json
+func (t *Task) HasDiscoveryInfo() bool {
+	return t.DiscoveryInfo.Name != ""
 }
 
-// IP extracts the IP from a task given the prioritized list of IP sources
+// IP returns the first Task IP found in the given sources.
 func (t *Task) IP(srcs ...string) string {
-	for _, src := range srcs {
-		switch src {
-		case "host":
-			return t.SlaveIP
-		case "docker", "mesos":
-			if ip := t.containerIP(src); ip != "" {
-				return ip
+	if ips := t.IPs(srcs...); len(ips) > 0 {
+		return ips[0].String()
+	}
+	return ""
+}
+
+// IPs returns a slice of IPs sourced from the given sources with ascending
+// priority.
+func (t *Task) IPs(srcs ...string) (ips []net.IP) {
+	if t == nil {
+		return nil
+	}
+	for i := range srcs {
+		if src, ok := sources[srcs[i]]; ok {
+			for _, srcIP := range src(t) {
+				if ip := net.ParseIP(srcIP); len(ip) > 0 {
+					ips = append(ips, ip)
+				}
 			}
 		}
 	}
-	return ""
+	return ips
 }
 
-// HasDiscoveryInfo return whether the DiscoveryInfo was provided in the state.json
-func (t *Task) HasDiscoveryInfo() bool {
-	return t.DiscoveryInfo.Name != ""
+// sources maps the string representation of IP sources to their functions.
+var sources = map[string]func(*Task) []string{
+	"host":    hostIPs,
+	"mesos":   mesosIPs,
+	"docker":  dockerIPs,
+	"netinfo": networkInfoIPs,
 }
 
-// containerIP extracts a container ip from a Mesos state.json task. If no
-// container ip is provided, an empty string is returned.
-func (t *Task) containerIP(src string) string {
-	ipLabel := ipLabels[src]
+// hostIPs is an IPSource which returns the IP addresses of the slave a Task
+// runs on.
+func hostIPs(t *Task) []string { return []string{t.SlaveIP} }
 
-	// find TASK_RUNNING statuses
-	var latestContainerIP string
-	var latestTimestamp float64
-	for _, status := range t.Statuses {
-		if status.State != "TASK_RUNNING" || status.Timestamp <= latestTimestamp {
-			continue
+// networkInfoIPs returns IP addresses from a given Task's
+// []Status.ContainerStatus.[]NetworkInfos.IPAddress
+func networkInfoIPs(t *Task) []string {
+	return statusIPs(t.Statuses, func(s *Status) []string {
+		ips := make([]string, len(s.ContainerStatus.NetworkInfos))
+		for i, netinfo := range s.ContainerStatus.NetworkInfos {
+			ips[i] = netinfo.IPAddress
 		}
+		return ips
+	})
+}
+
+const (
+	// DockerIPLabel is the key of the Label which holds the Docker containerizer IP value.
+	DockerIPLabel = "Docker.NetworkSettings.IPAddress"
+	// MesosIPLabel is the key of the label which holds the Mesos containerizer IP value.
+	MesosIPLabel = "MesosContainerizer.NetworkSettings.IPAddress"
+)
+
+// dockerIPs returns IP addresses from the values of all
+// Task.[]Status.[]Labels whose keys are equal to "Docker.NetworkSettings.IPAddress".
+func dockerIPs(t *Task) []string {
+	return statusIPs(t.Statuses, labels(DockerIPLabel))
+}
 
-		// first try to extract the address from the NetworkInfo structure
-		if len(status.ContainerStatus.NetworkInfos) > 0 {
-			// TODO(CD): handle multiple NetworkInfo objects
-			// TODO(CD): only create A records if the address is IPv4
-			// TODO(CD): create AAAA records if the address is IPv6
-			latestContainerIP = status.ContainerStatus.NetworkInfos[0].IPAddress
-			latestTimestamp = status.Timestamp
-			continue // Skip label inspection step below for this element
+// mesosIPs returns IP addresses from the values of all
+// Task.[]Status.[]Labels whose keys are equal to
+// "MesosContainerizer.NetworkSettings.IPAddress".
+func mesosIPs(t *Task) []string {
+	return statusIPs(t.Statuses, labels(MesosIPLabel))
+}
+
+// statusIPs returns the latest running status IPs extracted with the given src
+func statusIPs(st []Status, src func(*Status) []string) []string {
+	for i := len(st) - 1; i >= 0; i-- {
+		if st[i].State == "TASK_RUNNING" {
+			return src(&st[i])
 		}
+	}
+	return nil
+}
 
-		// next, fall back to the docker-inspect label
-		// TODO(CD): deprecate and then remove this address discovery method
-		for _, label := range status.Labels {
-			if label.Key == ipLabel {
-				latestContainerIP = label.Value
-				latestTimestamp = status.Timestamp
-				break
+// labels returns all given Status.[]Labels' values whose keys are equal
+// to the given key
+func labels(key string) func(*Status) []string {
+	return func(s *Status) []string {
+		vs := make([]string, 0, len(s.Labels))
+		for _, l := range s.Labels {
+			if l.Key == key {
+				vs = append(vs, l.Value)
 			}
 		}
+		return vs
 	}
-
-	return latestContainerIP
 }
 
 // Framework holds a framework as defined in the /state.json Mesos HTTP endpoint.

records/state/state_test.go

@@ -1,11 +1,13 @@
-package state
+package state_test
 
 import (
 	"encoding/json"
+	"net"
 	"reflect"
 	"testing"
 
 	"github.com/mesos/mesos-go/upid"
+	. "github.com/mesosphere/mesos-dns/records/state"
 )
 
 func TestResources_Ports(t *testing.T) {
@@ -39,62 +41,134 @@ func TestPID_UnmarshalJSON(t *testing.T) {
 	}
 }
 
-func TestTask_containerIP(t *testing.T) {
-	makeTask := func(networkInfoAddress string, label Label) Task {
-		labels := make([]Label, 0, 1)
-		if label.Key != "" {
-			labels = append(labels, label)
+func TestTask_IPs(t *testing.T) {
+	for i, tt := range []struct {
+		*Task
+		srcs []string
+		want []net.IP
+	}{
+		{nil, nil, nil},
+		{nil, []string{}, nil},
+		{nil, []string{"host"}, nil},
+		{ // no IPs for the given sources
+			Task: task(statuses(status(state("TASK_RUNNING"), netinfo("1.2.3.4")))),
+			srcs: []string{"host", "mesos"},
+			want: nil,
+		},
+		{ // unknown IP sources are ignored
+			Task: task(statuses(status(state("TASK_RUNNING"), netinfo("1.2.3.4")))),
+			srcs: []string{"foo", "netinfo", "bar"},
+			want: ips("1.2.3.4"),
+		},
+		{ // source order
+			Task: task(
+				slaveIP("2.3.4.5"),
+				statuses(status(state("TASK_RUNNING"), netinfo("1.2.3.4"))),
+			),
+			srcs: []string{"host", "netinfo"},
+			want: ips("2.3.4.5", "1.2.3.4"),
+		},
+		{ // statuses state
+			Task: task(
+				statuses(
+					status(state("TASK_RUNNING"), netinfo("1.2.3.4")),
+					status(state("TASK_STOPPED"), netinfo("2.3.4.5")),
+				),
+			),
+			srcs: []string{"netinfo"},
+			want: ips("1.2.3.4"),
+		},
+		{ // statuses ordering
+			Task: task(
+				statuses(
+					status(state("TASK_RUNNING"), netinfo("1.2.3.4")),
+					status(state("TASK_RUNNING"), labels(DockerIPLabel, "2.3.4.5")),
+				),
+			),
+			srcs: []string{"docker", "netinfo"},
+			want: ips("2.3.4.5"),
+		},
+		{ // label ordering
+			Task: task(
+				statuses(
+					status(
+						state("TASK_RUNNING"),
+						labels(DockerIPLabel, "1.2.3.4", DockerIPLabel, "2.3.4.5"),
+					),
+				),
+			),
+			srcs: []string{"docker"},
+			want: ips("1.2.3.4", "2.3.4.5"),
+		},
+	} {
+		if got := tt.IPs(tt.srcs...); !reflect.DeepEqual(got, tt.want) {
+			t.Logf("%+v", tt.Task)
+			t.Errorf("test #%d: got %+v, want %+v", i, got, tt.want)
 		}
+	}
+}
 
-		var containerStatus ContainerStatus
-		if networkInfoAddress != "" {
-			containerStatus = ContainerStatus{
-				NetworkInfos: []NetworkInfo{
-					NetworkInfo{
-						IPAddress: networkInfoAddress,
-					},
-				},
-			}
-		}
+// test helpers
 
-		return Task{
-			State: "TASK_RUNNING",
-			Statuses: []Status{
-				Status{
-					Timestamp:       1.0,
-					State:           "TASK_RUNNING",
-					Labels:          labels,
-					ContainerStatus: containerStatus,
-				},
-			},
-		}
+type (
+	taskOpt   func(*Task)
+	statusOpt func(*Status)
+)
+
+func ips(ss ...string) []net.IP {
+	addrs := make([]net.IP, len(ss))
+	for i := range ss {
+		addrs[i] = net.ParseIP(ss[i])
 	}
+	return addrs
+}
 
-	// Verify IP extraction from NetworkInfo
-	task := makeTask("1.2.3.4", Label{})
-	if task.containerIP("mesos") != "1.2.3.4" {
-		t.Errorf("Failed to extract IP from NetworkInfo")
+func task(opts ...taskOpt) *Task {
+	var t Task
+	for _, opt := range opts {
+		opt(&t)
 	}
+	return &t
+}
 
-	// Verify IP extraction from NetworkInfo takes precedence over
-	// labels
-	task = makeTask("1.2.3.4",
-		Label{Key: ipLabels["mesos"], Value: "2.4.6.8"})
-	if task.containerIP("mesos") != "1.2.3.4" {
-		t.Errorf("Failed to extract IP from NetworkInfo when label also supplied")
+func statuses(st ...Status) taskOpt {
+	return func(t *Task) {
+		t.Statuses = append(t.Statuses, st...)
 	}
+}
 
-	// Verify IP extraction from the Mesos label without NetworkInfo
-	task = makeTask("",
-		Label{Key: ipLabels["mesos"], Value: "1.2.3.4"})
-	if task.containerIP("mesos") != "1.2.3.4" {
-		t.Errorf("Failed to extract IP from Mesos label")
+func slaveIP(ip string) taskOpt {
+	return func(t *Task) { t.SlaveIP = ip }
+}
+
+func status(opts ...statusOpt) Status {
+	var s Status
+	for _, opt := range opts {
+		opt(&s)
 	}
+	return s
+}
 
-	// Verify IP extraction from the Docker label without NetworkInfo
-	task = makeTask("",
-		Label{Key: ipLabels["docker"], Value: "1.2.3.4"})
-	if task.containerIP("docker") != "1.2.3.4" {
-		t.Errorf("Failed to extract IP from Docker label")
+func labels(kvs ...string) statusOpt {
+	if len(kvs)%2 != 0 {
+		panic("odd number")
+	}
+	return func(s *Status) {
+		for i := 0; i < len(kvs); i += 2 {
+			s.Labels = append(s.Labels, Label{Key: kvs[i], Value: kvs[i+1]})
+		}
+	}
+}
+
+func state(st string) statusOpt {
+	return func(s *Status) { s.State = st }
+}
+
+func netinfo(ips ...string) statusOpt {
+	return func(s *Status) {
+		netinfos := &s.ContainerStatus.NetworkInfos
+		for _, ip := range ips {
+			*netinfos = append(*netinfos, NetworkInfo{IPAddress: ip})
+		}
 	}
 }