Consul hook: add service port tag for proxy ports (#129)

hook/consul/hook.go

@@ -21,6 +21,9 @@ const (
 	consulNameLabelKey = "consul"
 	consulTagValue     = "tag"
 	serviceHost        = "127.0.0.1"
+	servicePortName    = "service"
+	proxyPortName      = "proxyingress"
+	servicePortTag     = "service-port:%d"
 )
 
 // instance represents a service in consul
@@ -92,6 +95,7 @@ func (h *Hook) RegisterIntoConsul(taskInfo mesosutils.TaskInfo) error {
 	}
 
 	ports := taskInfo.GetPorts()
+	servicePort := getServicePort(ports)
 	globalTags := append(taskInfo.GetLabelKeysByValue(consulTagValue), h.config.ConsulGlobalTag)
 
 	var instancesToRegister []instance
@@ -107,6 +111,10 @@ func (h *Hook) RegisterIntoConsul(taskInfo mesosutils.TaskInfo) error {
 		consulServiceID := fmt.Sprintf("%s_%s_%d", taskID, portServiceName, port.GetNumber())
 		portTags := mesosutils.GetLabelKeysByValue(port.GetLabels().GetLabels(), consulTagValue)
 		portTags = append(portTags, globalTags...)
+		if port.Name != nil && *port.Name == proxyPortName && servicePort != nil {
+			portTags = append(portTags, fmt.Sprintf(servicePortTag, servicePort.Number))
+		}
+
 		log.Infof("Adding service ID %q to deregister before termination", consulServiceID)
 		instancesToRegister = append(instancesToRegister, instance{
 			consulServiceName: portServiceName,
@@ -180,6 +188,15 @@ func getServiceLabel(port mesos.Port) (string, error) {
 	return label.GetValue(), nil
 }
 
+func getServicePort(ports []mesos.Port) *mesos.Port {
+	for _, port := range ports {
+		if port.Name != nil && *port.Name == servicePortName {
+			return &port
+		}
+	}
+	return nil
+}
+
 func marathonAppNameToServiceName(name mesosutils.TaskID) string {
 	var sanitizer = strings.NewReplacer("_", ".", "/", "-")
 	// Remove all spaces and initial slashes, replace above characters

hook/consul/hook_test.go

@@ -374,6 +374,61 @@ func TestIfNewHookCreatesNoopHookWhenHookDisabled(t *testing.T) {
 	require.IsType(t, hook.NoopHook{}, h)
 }
 
+func TestIfServicePortTagAdded(t *testing.T) {
+	servicePortName := "service"
+	proxyPortName := "proxyingress"
+	otherPortName := "otherport"
+	tagName := "tag"
+
+	consulMainName := "serviceName"
+	consulAdditionalName := "additionalServiceName"
+
+	taskInfo := prepareTaskInfo("taskId", "taskName", "taskName", []string{"global-tag"}, []mesos.Port{
+		{Number: 667},
+		{Number: 668, Name: &servicePortName},
+		{Number: 669, Name: &otherPortName, Labels: &mesos.Labels{Labels: []mesos.Label{
+			{Key: "consul", Value: &consulAdditionalName},
+			{Key: "other-port-tag", Value: &tagName},
+		}}},
+		{Number: 670, Name: &proxyPortName, Labels: &mesos.Labels{Labels: []mesos.Label{
+			{Key: "consul", Value: &consulMainName},
+			{Key: "main-port-tag", Value: &tagName},
+		}}},
+	})
+
+	expectedMainInstance := instance{
+		consulServiceName: "serviceName",
+		port:              670,
+		tags:              []string{"marathon", "main-port-tag", "global-tag", "service-port:668"},
+	}
+
+	expectedAdditionalInstance := instance{
+		consulServiceName: "additionalServiceName",
+		port:              669,
+		tags:              []string{"marathon", "other-port-tag", "global-tag"},
+	}
+
+	// Create a test Consul server
+	config, server := createTestConsulServer(t)
+	client, _ := api.NewClient(config) // #nosec
+	defer stopConsul(server)
+
+	h := &Hook{config: Config{ConsulGlobalTag: "marathon"}, client: client}
+
+	err := h.RegisterIntoConsul(taskInfo)
+
+	require.NoError(t, err)
+	require.Len(t, h.serviceInstances, 2)
+
+	opts := api.QueryOptions{}
+	services, _, err := client.Catalog().Services(&opts)
+
+	require.Contains(t, services, consulMainName)
+	requireEqualElements(t, expectedMainInstance.tags, services[consulMainName])
+	require.Contains(t, services, consulAdditionalName)
+	requireEqualElements(t, expectedAdditionalInstance.tags, services[consulAdditionalName])
+}
+
 func stopConsul(server *testutil.TestServer) {
 	_ = server.Stop()
 }