Use tags and selectors to allow registering one port under couple of …

…names (#140)
allegro · Jul 6, 2020 · e6c3fd2 · e6c3fd2
1 parent 8fa54b2
commit e6c3fd2
Show file tree

Hide file tree

Showing 2 changed files with 113 additions and 19 deletions.
diff --git a/hook/consul/hook.go b/hook/consul/hook.go
@@ -102,26 +102,29 @@ func (h *Hook) RegisterIntoConsul(taskInfo mesosutils.TaskInfo) error {
 
 	var instancesToRegister []instance
 	for _, port := range ports {
-		portServiceName, err := getServiceLabel(port)
+		portServiceNames, err := getServiceLabels(port)
 		if err != nil {
 			log.Debugf("Pre-registration check for port failed: %s", err.Error())
 			continue
 		}
-		// consulServiceID is generated the same way as it is in marathon-consul - because
-		// it registers the service
-		// See: https://github.com/allegro/marathon-consul/blob/v1.1.0/consul/consul.go#L299-L301
-		consulServiceID := fmt.Sprintf("%s_%s_%d", taskID, portServiceName, port.GetNumber())
-		marathonTaskTag := fmt.Sprintf("marathon-task:%s", taskID)
-		portTags := mesosutils.GetLabelKeysByValue(port.GetLabels().GetLabels(), consulTagValue)
-		portTags = append(portTags, globalTags...)
-		portTags = append(portTags, marathonTaskTag)
-		log.Infof("Adding service ID %q to deregister before termination", consulServiceID)
-		instancesToRegister = append(instancesToRegister, instance{
-			consulServiceName: portServiceName,
-			consulServiceID:   consulServiceID,
-			port:              port.GetNumber(),
-			tags:              portTags,
-		})
+
+		for _, portServiceName := range portServiceNames {
+			// consulServiceID is generated the same way as it is in marathon-consul - because
+			// it registers the service
+			// See: https://github.com/allegro/marathon-consul/blob/v1.1.0/consul/consul.go#L299-L301
+			consulServiceID := fmt.Sprintf("%s_%s_%d", taskID, portServiceName, port.GetNumber())
+			marathonTaskTag := fmt.Sprintf("marathon-task:%s", taskID)
+			portTags := getPortTags(port, portServiceName)
+			portTags = append(portTags, globalTags...)
+			portTags = append(portTags, marathonTaskTag)
+			log.Infof("Adding service ID %q to deregister before termination", consulServiceID)
+			instancesToRegister = append(instancesToRegister, instance{
+				consulServiceName: portServiceName,
+				consulServiceID:   consulServiceID,
+				port:              port.GetNumber(),
+				tags:              portTags,
+			})
+		}
 	}
 
 	if len(instancesToRegister) == 0 {
@@ -161,6 +164,28 @@ func (h *Hook) RegisterIntoConsul(taskInfo mesosutils.TaskInfo) error {
 	return nil
 }
 
+func getPortTags(port mesos.Port, serviceName string) []string {
+	var keys []string
+	labels := port.GetLabels().GetLabels()
+
+	for _, label := range labels {
+		value := label.GetValue()
+		valueAndSelector := strings.Split(value, ":")
+		if len(valueAndSelector) > 1 {
+			value := valueAndSelector[0]
+			serviceSelector := valueAndSelector[1]
+
+			if value == consulTagValue && serviceSelector == serviceName {
+				keys = append(keys, label.GetKey())
+			}
+		} else if value == consulTagValue {
+			keys = append(keys, label.GetKey())
+		}
+	}
+
+	return keys
+}
+
 // DeregisterFromConsul will deregister service IDs from Consul that were created
 // during AfterTaskStartEvent hook event.
 func (h *Hook) DeregisterFromConsul(taskInfo mesosutils.TaskInfo) error {
@@ -220,12 +245,15 @@ func resolvePlaceholders(values []string, placeholders map[string]string) []stri
 	return resolved
 }
 
-func getServiceLabel(port mesos.Port) (string, error) {
+func getServiceLabels(port mesos.Port) ([]string, error) {
 	label := mesosutils.FindLabel(port.GetLabels().GetLabels(), consulNameLabelKey)
 	if label == nil {
-		return "", fmt.Errorf("port %d has no label %q", port.GetNumber(), consulNameLabelKey)
+		return nil, fmt.Errorf("port %d has no label %q", port.GetNumber(), consulNameLabelKey)
 	}
-	return label.GetValue(), nil
+
+	labels := strings.Split(label.GetValue(), ",")
+
+	return labels, nil
 }
 
 func marathonAppNameToServiceName(name mesosutils.TaskID) string {

diff --git a/hook/consul/hook_test.go b/hook/consul/hook_test.go
@@ -155,6 +155,72 @@ func TestIfUsesLabelledPortsForServiceIDGenAndRegisterMultiplePorts(t *testing.T
 	require.Contains(t, services, consulNameSecond)
 }
 
+func TestIfUsesCompoundLabelledPortsForServiceIDGenAndRegisterSinglePortWithMultipleNames(t *testing.T) {
+	// given
+	consulName := "consulName,consulName-secured"
+	taskID := "taskID"
+	securedTagValue := "tag:consulName-secured"
+	insecureTagValue := "tag:consulName"
+	commonTagValue := "tag"
+	taskInfo := prepareTaskInfo(taskID, consulName, consulName, []string{"metrics"}, []mesos.Port{
+		{
+			Number: 998,
+			Labels: &mesos.Labels{
+				Labels: []mesos.Label{
+					{
+						Key:   "consul",
+						Value: &consulName,
+					},
+					{
+						Key:   "secure",
+						Value: &securedTagValue,
+					},
+					{
+						Key:   "insecure",
+						Value: &insecureTagValue,
+					},
+					{
+						Key:   "common",
+						Value: &commonTagValue,
+					},
+				},
+			},
+		},
+	})
+	expectedService := instance{
+		consulServiceName: "consulName",
+		consulServiceID:   createServiceID(taskID, "consulName", 998),
+		port:              998,
+		tags:              []string{"insecure", "common", "metrics", "marathon", "marathon-task:taskID"},
+	}
+	expectedService2 := instance{
+		consulServiceName: "consulName-secured",
+		consulServiceID:   createServiceID(taskID, "consulName-secured", 998),
+		port:              998,
+		tags:              []string{"secure", "common", "metrics", "marathon", "marathon-task:taskID"},
+	}
+
+	// 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}
+
+	// when
+	err := h.RegisterIntoConsul(taskInfo)
+	opts := api.QueryOptions{}
+	services, _, err := client.Catalog().Services(&opts)
+
+	// then
+	require.NoError(t, err)
+	require.Len(t, h.serviceInstances, 2)
+	require.Contains(t, services, "consulName")
+	requireEqualElements(t, expectedService.tags, services["consulName"])
+	require.Contains(t, services, "consulName-secured")
+	requireEqualElements(t, expectedService2.tags, services["consulName-secured"])
+}
+
 func TestIfUsesPortLabelsForRegistration(t *testing.T) {
 	consulName := "consulName"
 	consulNameSecond := "consulName-secured"