Merge pull request #68 from wendigo/master

Multiple checks, fix for url encoding, TCP/Command checks
allegro · Mar 29, 2016 · c967bab · c967bab
2 parents 12ae8a1 + 65c47ca
commit c967bab
Show file tree

Hide file tree

Showing 3 changed files with 101 additions and 54 deletions.
diff --git a/apps/app.go b/apps/app.go
@@ -16,6 +16,9 @@ type HealthCheck struct {
 	IntervalSeconds        int    `json:"intervalSeconds"`
 	TimeoutSeconds         int    `json:"timeoutSeconds"`
 	MaxConsecutiveFailures int    `json:"maxConsecutiveFailures"`
+	Command                struct {
+		Value string `json:"value`
+	}
 }
 
 type AppWrapper struct {

diff --git a/consul/consul.go b/consul/consul.go
@@ -3,6 +3,7 @@ package consul
 import (
 	"fmt"
 	"net/url"
+	"strings"
 
 	log "github.com/Sirupsen/logrus"
 	"github.com/allegro/marathon-consul/apps"
@@ -203,26 +204,49 @@ func (c *Consul) marathonTaskToConsulService(task *apps.Task, app *apps.App) (*c
 		Port:    task.Ports[0],
 		Address: serviceAddress,
 		Tags:    c.marathonLabelsToConsulTags(app.Labels),
-		Check:   c.marathonToConsulCheck(task, app.HealthChecks, serviceAddress),
+		Checks:  c.marathonToConsulChecks(task, app.HealthChecks, serviceAddress),
 	}, nil
 }
 
-func (c *Consul) marathonToConsulCheck(task *apps.Task, healthChecks []apps.HealthCheck, serviceAddress string) *consulapi.AgentServiceCheck {
-	//	TODO: Handle all types of checks
+func (c *Consul) marathonToConsulChecks(task *apps.Task, healthChecks []apps.HealthCheck, serviceAddress string) consulapi.AgentServiceChecks {
+	var checks consulapi.AgentServiceChecks = make(consulapi.AgentServiceChecks, 0, len(healthChecks))
+
 	for _, check := range healthChecks {
-		if check.Protocol == "HTTP" {
-			return &consulapi.AgentServiceCheck{
-				HTTP: (&url.URL{
-					Scheme: "http",
-					Host:   fmt.Sprintf("%s:%d", serviceAddress, task.Ports[check.PortIndex]),
-					Path:   check.Path,
-				}).String(),
+		switch check.Protocol {
+		case "HTTP", "HTTPS":
+			if parsedUrl, err := url.ParseRequestURI(check.Path); err == nil {
+				parsedUrl.Scheme = strings.ToLower(check.Protocol)
+				parsedUrl.Host = fmt.Sprintf("%s:%d", serviceAddress, task.Ports[check.PortIndex])
+
+				checks = append(checks, &consulapi.AgentServiceCheck{
+					HTTP:     parsedUrl.String(),
+					Interval: fmt.Sprintf("%ds", check.IntervalSeconds),
+					Timeout:  fmt.Sprintf("%ds", check.TimeoutSeconds),
+				})
+			} else {
+				log.WithError(err).
+					WithField("Id", task.AppID.String()).
+					WithField("Address", serviceAddress).
+					Warn(fmt.Sprintf("Could not parse provided path: %s", check.Path))
+			}
+		case "TCP":
+			checks = append(checks, &consulapi.AgentServiceCheck{
+				TCP:      fmt.Sprintf("%s:%d", serviceAddress, task.Ports[check.PortIndex]),
 				Interval: fmt.Sprintf("%ds", check.IntervalSeconds),
 				Timeout:  fmt.Sprintf("%ds", check.TimeoutSeconds),
-			}
+			})
+		case "COMMAND":
+			checks = append(checks, &consulapi.AgentServiceCheck{
+				Script:   check.Command.Value,
+				Interval: fmt.Sprintf("%ds", check.IntervalSeconds),
+				Timeout:  fmt.Sprintf("%ds", check.TimeoutSeconds),
+			})
+		default:
+			log.WithField("Id", task.AppID.String()).WithField("Address", serviceAddress).
+				Warn(fmt.Sprintf("Unrecognized check protocol %s", check.Protocol))
 		}
 	}
-	return nil
+	return checks
 }
 
 func (c *Consul) marathonLabelsToConsulTags(labels map[string]string) []string {

diff --git a/consul/consul_test.go b/consul/consul_test.go
@@ -507,63 +507,62 @@ func TestDeregisterServices_shouldReturnErrorOnFailure(t *testing.T) {
 	assert.Error(t, err)
 }
 
-func TestMarathonTaskToConsulServiceMapping_WithNoHttpChecks(t *testing.T) {
+func TestMarathonTaskToConsulServiceMapping(t *testing.T) {
 	t.Parallel()
 
 	// given
-	consul := New(ConsulConfig{})
-
+	consul := New(ConsulConfig{Tag: "marathon"})
 	app := &apps.App{
 		ID: "someApp",
 		HealthChecks: []apps.HealthCheck{
 			{
-				Path:                   "/",
-				Protocol:               "TCP",
+				Path:                   "/api/health?with=query",
+				Protocol:               "HTTP",
 				PortIndex:              0,
 				IntervalSeconds:        60,
 				TimeoutSeconds:         20,
 				MaxConsecutiveFailures: 3,
 			},
-		},
-		Labels: map[string]string{
-			"consul": "true",
-			"public": "tag",
-		},
-	}
-	task := &apps.Task{
-		ID:    "someTask",
-		AppID: app.ID,
-		Host:  "127.0.0.6",
-		Ports: []int{8090, 8443},
-	}
-
-	// when
-	service, err := consul.marathonTaskToConsulService(task, app)
-
-	// then
-	assert.NoError(t, err)
-	assert.Equal(t, "127.0.0.6", service.Address)
-	assert.Equal(t, 8090, service.Port)
-	assert.Nil(t, service.Check)
-	assert.Empty(t, service.Checks)
-}
-
-func TestMarathonTaskToConsulServiceMapping(t *testing.T) {
-	t.Parallel()
-
-	// given
-	consul := New(ConsulConfig{Tag: "marathon"})
-	app := &apps.App{
-		ID: "someApp",
-		HealthChecks: []apps.HealthCheck{
 			{
-				Path:                   "/api/health",
+				Path:                   "",
 				Protocol:               "HTTP",
 				PortIndex:              0,
 				IntervalSeconds:        60,
 				TimeoutSeconds:         20,
 				MaxConsecutiveFailures: 3,
 			},
+			{
+				Path:                   "/api/health?with=query",
+				Protocol:               "INVALID_PROTOCOL",
+				PortIndex:              0,
+				IntervalSeconds:        60,
+				TimeoutSeconds:         20,
+				MaxConsecutiveFailures: 3,
+			},
+			{
+				Path:                   "/secure/health?with=query",
+				Protocol:               "HTTPS",
+				PortIndex:              0,
+				IntervalSeconds:        50,
+				TimeoutSeconds:         20,
+				MaxConsecutiveFailures: 3,
+			},
+			{
+				Protocol:               "TCP",
+				PortIndex:              1,
+				IntervalSeconds:        40,
+				TimeoutSeconds:         20,
+				MaxConsecutiveFailures: 3,
+			},
+			{
+				Protocol: "COMMAND",
+				Command: struct {
+					Value string `json:"value`
+				}{Value: "echo 1"},
+				IntervalSeconds:        30,
+				TimeoutSeconds:         20,
+				MaxConsecutiveFailures: 3,
+			},
 		},
 		Labels: map[string]string{
 			"consul": "true",
@@ -585,10 +584,31 @@ func TestMarathonTaskToConsulServiceMapping(t *testing.T) {
 	assert.Equal(t, "127.0.0.6", service.Address)
 	assert.Equal(t, []string{"marathon", "public"}, service.Tags)
 	assert.Equal(t, 8090, service.Port)
-	assert.NotNil(t, "http://127.0.0.6:8090/api/health", service.Check)
-	assert.Empty(t, service.Checks)
-	assert.Equal(t, "http://127.0.0.6:8090/api/health", service.Check.HTTP)
-	assert.Equal(t, "60s", service.Check.Interval)
+	assert.Nil(t, service.Check)
+	assert.Equal(t, 4, len(service.Checks))
+
+	assert.Equal(t, consulapi.AgentServiceChecks{
+		{
+			HTTP:     "http://127.0.0.6:8090/api/health?with=query",
+			Interval: "60s",
+			Timeout:  "20s",
+		},
+		{
+			HTTP:     "https://127.0.0.6:8090/secure/health?with=query",
+			Interval: "50s",
+			Timeout:  "20s",
+		},
+		{
+			TCP:      "127.0.0.6:8443",
+			Interval: "40s",
+			Timeout:  "20s",
+		},
+		{
+			Script:   "echo 1",
+			Interval: "30s",
+			Timeout:  "20s",
+		},
+	}, service.Checks)
 }
 
 func TestMarathonTaskToConsulServiceMapping_NotResolvableTaskHost(t *testing.T) {