Add support for Consul TTLs and multiple checks.

CHANGELOG.md

@@ -21,6 +21,8 @@ All notable changes to this project will be documented in this file.
 ### Changed
 - Overall refactoring and cleanup
 - Decoupled registries into subpackages using extpoints
+- Add full TTL support for Consul 0.5.0.
+- Support multiple Consul checks per service (such as a TTL and script)
 
 
 ## [v5] - 2015-02-18

README.md

@@ -30,7 +30,7 @@ If the argument `-internal` is passed, registrator will register the docker0 int
 
 The `-resync` argument controls how often registrator will query Docker for all containers and reregister all services.  This allows registrator and the service registry to get back in sync if they fall out of sync.  The time is measured in seconds, and if set to zero will not resync.
 
-The consul backend does not support automatic expiry of stale registrations after some TTL. Instead, TTL checks must be configured (see below). For backends that do support TTL expiry, registrator can be started with the `-ttl` and `-ttl-refresh` arguments (both disabled by default).
+For backends that support TTL expiry, registrator can be started with the `-ttl` and `-ttl-refresh` arguments (both disabled by default).
 
 ### Host mode services
 
@@ -219,7 +219,7 @@ Then add a factory which accepts a uri and returns the registry adapter, and reg
 
 > All health checking integration is going to change soon, so consider these features deprecated.
 
-When using the Consul's service catalog backend, you can specify a health check associated with a service. Registrator can pull this from your container environment data if provided. Here are some examples:
+When using the Consul's service catalog backend, you can specify several health check to be associated with a service. Registrator can pull this from your container environment data if provided. Here are some examples:
 
 #### Basic HTTP health check
 

consul/consul.go

@@ -4,7 +4,9 @@ import (
 	"fmt"
 	"log"
 	"net/url"
+	"sort"
 	"strings"
+	"time"
 
 	"github.com/gliderlabs/registrator/bridge"
 	consulapi "github.com/hashicorp/consul/api"
@@ -59,37 +61,112 @@ func (r *ConsulAdapter) Register(service *bridge.Service) error {
 	registration.Port = service.Port
 	registration.Tags = service.Tags
 	registration.Address = service.IP
-	registration.Check = r.buildCheck(service)
+	registration.Checks = r.buildChecks(service)
 	return r.client.Agent().ServiceRegister(registration)
 }
 
-func (r *ConsulAdapter) buildCheck(service *bridge.Service) *consulapi.AgentServiceCheck {
-	check := new(consulapi.AgentServiceCheck)
-	if path := service.Attrs["check_http"]; path != "" {
-		check.Script = fmt.Sprintf("check-http %s %s %s", service.Origin.ContainerID[:12], service.Origin.ExposedPort, path)
-	} else if cmd := service.Attrs["check_cmd"]; cmd != "" {
-		check.Script = fmt.Sprintf("check-cmd %s %s %s", service.Origin.ContainerID[:12], service.Origin.ExposedPort, cmd)
-	} else if script := service.Attrs["check_script"]; script != "" {
-		check.Script = r.interpolateService(script, service)
-	} else if ttl := service.Attrs["check_ttl"]; ttl != "" {
-		check.TTL = ttl
-	} else {
-		return nil
-	}
-	if check.Script != "" {
-		if interval := service.Attrs["check_interval"]; interval != "" {
-			check.Interval = interval
-		} else {
-			check.Interval = DefaultInterval
+type Check struct {
+	Service *bridge.Service
+	Type    string
+	Value   string
+}
+type Checks []*Check
+
+func (c Checks) Len() int           { return len(c) }
+func (c Checks) Swap(i, j int)      { c[i], c[j] = c[j], c[i] }
+func (c Checks) Less(i, j int) bool { return c[i].Type < c[j].Type }
+
+/*
+ * Checks are ordered in Consul. There needs to be a deterministic way of
+ * re-attaching them for refreshes. The simplest solution is to order them
+ * alphabetically.
+ */
+func filterChecks(service *bridge.Service) Checks {
+	checks := make(Checks, 0)
+	for k, v := range service.Attrs {
+		if strings.Index(k, "check_") == 0 && k != "check_interval" && v != "" {
+			check := new(Check)
+			check.Service = service
+			check.Type = k
+			check.Value = v
+			checks = append(checks, check)
 		}
 	}
+	sort.Sort(checks)
+	return checks
+}
+
+func interval(service *bridge.Service) string {
+	interval := DefaultInterval
+	if service.Attrs["check_interval"] != "" {
+		interval = service.Attrs["check_interval"]
+	}
+	return interval
+}
+
+func scriptCheck(interval string, script string) *consulapi.AgentServiceCheck {
+	check := new(consulapi.AgentServiceCheck)
+	check.Script = script
+	check.Interval = interval
+	return check
+}
+
+func ttlCheck(ttl string) *consulapi.AgentServiceCheck {
+	check := new(consulapi.AgentServiceCheck)
+	check.TTL = ttl
 	return check
 }
 
+func (r *ConsulAdapter) buildChecks(service *bridge.Service) consulapi.AgentServiceChecks {
+	interval := interval(service)
+	checks := make(consulapi.AgentServiceChecks, 0)
+
+	for _, c := range filterChecks(service) {
+		switch c.Type {
+		case "check_http":
+			checks = append(checks, scriptCheck(interval, fmt.Sprintf("check-http %s %s %s", service.Origin.ContainerID[:12], service.Origin.ExposedPort, c.Value)))
+		case "check_cmd":
+			checks = append(checks, scriptCheck(interval, fmt.Sprintf("check-cmd %s %s %s", service.Origin.ContainerID[:12], service.Origin.ExposedPort, c.Value)))
+		case "check_script":
+			checks = append(checks, scriptCheck(interval, r.interpolateService(c.Value, service)))
+		case "check_ttl":
+			checks = append(checks, ttlCheck(c.Value))
+		default:
+			break
+		}
+	}
+
+	if len(checks) < 1 {
+		return nil
+	}
+	return checks
+}
+
 func (r *ConsulAdapter) Deregister(service *bridge.Service) error {
 	return r.client.Agent().ServiceDeregister(service.ID)
 }
 
-func (r *ConsulAdapter) Refresh(service *bridge.Service) error {
+// Used for testing.
+type refresher func(string, string) error
+
+func (r *ConsulAdapter) refresh(service *bridge.Service, refresher refresher) error {
+	checks := filterChecks(service)
+	count := len(checks)
+	for i, c := range checks {
+		if c.Type == "check_ttl" {
+			checkId := ""
+			if count > 1 {
+				checkId = fmt.Sprintf(":%d", i+1)
+			}
+			ttl := fmt.Sprintf("service:%s%s", service.ID, checkId)
+			log.Println("refreshing:", ttl)
+			// Because checks are passed in a map, there will only be one TTL check.
+			return refresher(ttl, time.Now().Format(time.RFC850))
+		}
+	}
 	return nil
 }
+
+func (r *ConsulAdapter) Refresh(service *bridge.Service) error {
+	return r.refresh(service, r.client.Agent().PassTTL)
+}

consul/consul_test.go

@@ -0,0 +1,88 @@
+package consul
+
+import (
+	"testing"
+
+	"github.com/gliderlabs/registrator/bridge"
+)
+
+func TestBuildChecksInOrder(t *testing.T) {
+	service := service()
+	service.Attrs["check_http"] = "/test"
+	service.Attrs["check_cmd"] = "echo 1"
+	service.Attrs["check_ttl"] = "30s"
+	service.Attrs["check_script"] = "some script"
+
+	checks := new(ConsulAdapter).buildChecks(service)
+	if checks[0].Script != "check-cmd 123456789012 1 echo 1" {
+		t.Error("Expected check-cmd but got", checks[0])
+	} else if checks[1].Script != "check-http 123456789012 1 /test" {
+		t.Error("Expected check-http but got", checks[1])
+	} else if checks[2].Script != "some script" {
+		t.Error("Expected check script but got", checks[2])
+	} else if checks[3].TTL != "30s" {
+		t.Error("Expected TTL but got", checks[3])
+	}
+}
+
+func TestBuildChecksInterpolates(t *testing.T) {
+	service := service()
+	service.Attrs["check_script"] = "$SERVICE_IP:$SERVICE_PORT"
+
+	script := new(ConsulAdapter).buildChecks(service)[0].Script
+	if script != "127.0.0.1:10" {
+		t.Error("Unexpected result,", script)
+	}
+}
+
+func TestIntervalSpecification(t *testing.T) {
+	service := service()
+	service.Attrs["check_interval"] = "15s"
+	service.Attrs["check_script"] = "something"
+
+	interval := new(ConsulAdapter).buildChecks(service)[0].Interval
+	if interval != "15s" {
+		t.Error("Unexpected result,", interval)
+	}
+}
+
+func TestRefreshOnlyTTL(t *testing.T) {
+	service := service()
+	service.Attrs["check_ttl"] = "30s"
+	service.Attrs["check_interval"] = "25s"
+	verifyTTLCall(t, service, "service:tests")
+}
+
+func TestRefreshMultipleChecks(t *testing.T) {
+	service := service()
+	service.Attrs["check_script"] = "something"
+	service.Attrs["check_ttl"] = "30s"
+	verifyTTLCall(t, service, "service:tests:2")
+}
+
+func verifyTTLCall(t *testing.T, service *bridge.Service, expected string) {
+	callCheck := ""
+	callNotes := ""
+	adapter := ConsulAdapter{}
+	adapter.refresh(service, func(check string, notes string) error {
+		callCheck = check
+		callNotes = notes
+		return nil
+	})
+	if callCheck != expected {
+		t.Error("Actual service check called:", callCheck)
+	} else if callNotes == "" {
+		t.Error("No notes passed.")
+	}
+}
+
+func service() *bridge.Service {
+	service := new(bridge.Service)
+	service.ID = "tests"
+	service.Origin.ExposedPort = "1"
+	service.Origin.ContainerID = "1234567890123"
+	service.Origin.HostIP = "127.0.0.1"
+	service.Origin.HostPort = "10"
+	service.Attrs = make(map[string]string)
+	return service
+}