From ac2045cd527bdafcddac6e95436db5785d763e02 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Senart?= <tomas@mesosphere.io>
Date: Sat, 1 Aug 2015 08:39:46 +0200
Subject: [PATCH] Generate SRV records for each slave

Closes #200
---
 factories/fake.json       |  2 +-
 records/generator.go      | 21 +++++++++++++++++++--
 records/generator_test.go | 23 +++++++++++++++++++++++
 3 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/factories/fake.json b/factories/fake.json
index 27f21321..585a544c 100644
--- a/factories/fake.json
+++ b/factories/fake.json
@@ -380,7 +380,7 @@
             },
             "hostname": "1.2.3.12",
             "id": "20140827-000744-3041283216-5050-2116-1",
-            "pid": "slave(1)@1.2.3:12",
+            "pid": "slave(1)@1.2.3.12:5051",
             "registered_time": 1414913537.19731,
             "reregistered_time": 1414913537.19735,
             "resources": {
diff --git a/records/generator.go b/records/generator.go
index 9bc9a751..e48e1f46 100644
--- a/records/generator.go
+++ b/records/generator.go
@@ -3,6 +3,7 @@
 package records
 
 import (
+	"bytes"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -15,6 +16,8 @@ import (
 
 	"github.com/mesosphere/mesos-dns/logging"
 	"github.com/mesosphere/mesos-dns/records/labels"
+
+	"github.com/mesos/mesos-go/upid"
 )
 
 // Map host/service name to DNS answer
@@ -69,6 +72,16 @@ type Framework struct {
 type Slave struct {
 	ID       string `json:"id"`
 	Hostname string `json:"hostname"`
+	PID      PID    `json:"pid"`
+}
+
+// PID holds a Mesos PID and implements the json.Unmarshaler interface.
+type PID struct{ *upid.UPID }
+
+// UnmarshalJSON implements the json.Unmarshaler interface for PIDs.
+func (p *PID) UnmarshalJSON(data []byte) (err error) {
+	p.UPID, err = upid.Parse(string(bytes.Trim(data, `" `)))
+	return err
 }
 
 // StateJSON holds the state defined in the /state.json Mesos HTTP endpoint.
@@ -264,9 +277,13 @@ func (rg *RecordGenerator) InsertState(sj StateJSON, domain string, ns string,
 	rg.As = make(rrs)
 
 	for _, slave := range sj.Slaves {
-		address, ok := hostToIP4(slave.Hostname)
+		address, ok := hostToIP4(slave.PID.Host)
 		if ok {
-			rg.insertRR("slave."+domain+".", address, "A")
+			a := "slave." + domain + "."
+			rg.insertRR(a, address, "A")
+			srv := net.JoinHostPort(a, slave.PID.Port)
+			rg.insertRR("_slave._tcp."+domain+".", srv, "SRV")
+			rg.insertRR("_slave._udp."+domain+".", srv, "SRV")
 		} else {
 			logging.VeryVerbose.Printf("string '%q' for slave with id %q is not a valid IP address", address, slave.ID)
 			address = labels.AsDomainFrag(address, spec)
diff --git a/records/generator_test.go b/records/generator_test.go
index f5c00cdd..4834cd8c 100644
--- a/records/generator_test.go
+++ b/records/generator_test.go
@@ -7,6 +7,7 @@ import (
 	"testing"
 	"testing/quick"
 
+	"github.com/mesos/mesos-go/upid"
 	"github.com/mesosphere/mesos-dns/logging"
 	"github.com/mesosphere/mesos-dns/records/labels"
 )
@@ -249,6 +250,8 @@ func TestInsertState(t *testing.T) {
 			"liquor-store-7581-1.marathon.mesos.:31737",
 		}},
 		{rg.SRVs, "SRV", "_liquor-store.marathon.mesos.", nil},
+		{rg.SRVs, "SRV", "_slave._tcp.mesos.", []string{"slave.mesos.:5051"}},
+		{rg.SRVs, "SRV", "_slave._udp.mesos.", []string{"slave.mesos.:5051"}},
 	} {
 		if got := tt.rrs[tt.name]; !reflect.DeepEqual(got, tt.want) {
 			t.Errorf("test #%d: %s record for %q: got: %q, want: %q", i, tt.kind, tt.name, got, tt.want)
@@ -256,6 +259,26 @@ func TestInsertState(t *testing.T) {
 	}
 }
 
+func TestPID_UnmarshalJSON(t *testing.T) {
+	for i, tt := range []struct {
+		data string
+		want PID
+		err  error
+	}{
+		{`"slave(1)@127.0.0.1:5051"`, PID{&upid.UPID{"slave(1)", "127.0.0.1", "5051"}}, nil},
+		{`  "slave(1)@127.0.0.1:5051"  `, PID{&upid.UPID{"slave(1)", "127.0.0.1", "5051"}}, nil},
+		{`"  slave(1)@127.0.0.1:5051  "`, PID{&upid.UPID{"slave(1)", "127.0.0.1", "5051"}}, nil},
+	} {
+		var pid PID
+		if err := json.Unmarshal([]byte(tt.data), &pid); !reflect.DeepEqual(err, tt.err) {
+			t.Errorf("test #%d: got err: %v, want: %v", i, err, tt.want)
+		}
+		if got := pid; !reflect.DeepEqual(got, tt.want) {
+			t.Errorf("test #%d: got: %v, want: %v", i, got, tt.want)
+		}
+	}
+}
+
 // ensure we only generate one A record for each host
 func TestNTasks(t *testing.T) {
 	rg := &RecordGenerator{}