Support multiple health checks per service

api/agent.go

@@ -41,18 +41,20 @@ type AgentMember struct {
 
 // AgentServiceRegistration is used to register a new service
 type AgentServiceRegistration struct {
-	ID    string   `json:",omitempty"`
-	Name  string   `json:",omitempty"`
-	Tags  []string `json:",omitempty"`
-	Port  int      `json:",omitempty"`
-	Check *AgentServiceCheck
+	ID     string   `json:",omitempty"`
+	Name   string   `json:",omitempty"`
+	Tags   []string `json:",omitempty"`
+	Port   int      `json:",omitempty"`
+	Check  *AgentServiceCheck
+	Checks AgentServiceChecks
 }
 
 // AgentCheckRegistration is used to register a new check
 type AgentCheckRegistration struct {
-	ID    string `json:",omitempty"`
-	Name  string `json:",omitempty"`
-	Notes string `json:",omitempty"`
+	ID        string `json:",omitempty"`
+	Name      string `json:",omitempty"`
+	Notes     string `json:",omitempty"`
+	ServiceID string `json:",omitempty"`
 	AgentServiceCheck
 }
 
@@ -63,6 +65,7 @@ type AgentServiceCheck struct {
 	Interval string `json:",omitempty"`
 	TTL      string `json:",omitempty"`
 }
+type AgentServiceChecks []*AgentServiceCheck
 
 // Agent can be used to query the Agent endpoints
 type Agent struct {

api/agent_test.go

@@ -76,6 +76,49 @@ func TestAgent_Services(t *testing.T) {
 	}
 }
 
+func TestAgent_Services_MultipleChecks(t *testing.T) {
+	c, s := makeClient(t)
+	defer s.stop()
+
+	agent := c.Agent()
+
+	reg := &AgentServiceRegistration{
+		Name: "foo",
+		Tags: []string{"bar", "baz"},
+		Port: 8000,
+		Checks: AgentServiceChecks{
+			&AgentServiceCheck{
+				TTL: "15s",
+			},
+			&AgentServiceCheck{
+				TTL: "30s",
+			},
+		},
+	}
+	if err := agent.ServiceRegister(reg); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	services, err := agent.Services()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if _, ok := services["foo"]; !ok {
+		t.Fatalf("missing service: %v", services)
+	}
+
+	checks, err := agent.Checks()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if _, ok := checks["service:foo:1"]; !ok {
+		t.Fatalf("missing check: %v", checks)
+	}
+	if _, ok := checks["service:foo:2"]; !ok {
+		t.Fatalf("missing check: %v", checks)
+	}
+}
+
 func TestAgent_SetTTLStatus(t *testing.T) {
 	c, s := makeClient(t)
 	defer s.stop()
@@ -143,6 +186,44 @@ func TestAgent_Checks(t *testing.T) {
 	}
 }
 
+func TestAgent_Checks_serviceBound(t *testing.T) {
+	c, s := makeClient(t)
+	defer s.stop()
+
+	agent := c.Agent()
+
+	// First register a service
+	serviceReg := &AgentServiceRegistration{
+		Name: "redis",
+	}
+	if err := agent.ServiceRegister(serviceReg); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Register a check bound to the service
+	reg := &AgentCheckRegistration{
+		Name:      "redischeck",
+		ServiceID: "redis",
+	}
+	reg.TTL = "15s"
+	if err := agent.CheckRegister(reg); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	checks, err := agent.Checks()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	check, ok := checks["redischeck"]
+	if !ok {
+		t.Fatalf("missing check: %v", checks)
+	}
+	if check.ServiceID != "redis" {
+		t.Fatalf("missing service association for check: %v", check)
+	}
+}
+
 func TestAgent_Join(t *testing.T) {
 	c, s := makeClient(t)
 	defer s.stop()

command/agent/agent.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"path/filepath"
 	"strconv"
+	"strings"
 	"sync"
 
 	"github.com/hashicorp/consul/consul"
@@ -582,15 +583,17 @@ func (a *Agent) purgeCheck(checkID string) error {
 // AddService is used to add a service entry.
 // This entry is persistent and the agent will make a best effort to
 // ensure it is registered
-func (a *Agent) AddService(service *structs.NodeService, chkType *CheckType, persist bool) error {
+func (a *Agent) AddService(service *structs.NodeService, chkTypes CheckTypes, persist bool) error {
 	if service.Service == "" {
 		return fmt.Errorf("Service name missing")
 	}
 	if service.ID == "" && service.Service != "" {
 		service.ID = service.Service
 	}
-	if chkType != nil && !chkType.Valid() {
-		return fmt.Errorf("Check type is not valid")
+	for _, check := range chkTypes {
+		if !check.Valid() {
+			return fmt.Errorf("Check type is not valid")
+		}
 	}
 
 	// Add the service
@@ -604,10 +607,14 @@ func (a *Agent) AddService(service *structs.NodeService, chkType *CheckType, per
 	}
 
 	// Create an associated health check
-	if chkType != nil {
+	for i, chkType := range chkTypes {
+		checkID := fmt.Sprintf("service:%s", service.ID)
+		if len(chkTypes) > 1 {
+			checkID += fmt.Sprintf(":%d", i+1)
+		}
 		check := &structs.HealthCheck{
 			Node:        a.config.NodeName,
-			CheckID:     fmt.Sprintf("service:%s", service.ID),
+			CheckID:     checkID,
 			Name:        fmt.Sprintf("Service '%s' check", service.Service),
 			Status:      structs.HealthCritical,
 			Notes:       chkType.Notes,
@@ -642,9 +649,14 @@ func (a *Agent) RemoveService(serviceID string, persist bool) error {
 	}
 
 	// Deregister any associated health checks
-	checkID := fmt.Sprintf("service:%s", serviceID)
-	if err := a.RemoveCheck(checkID, persist); err != nil {
-		return err
+	for checkID, _ := range a.state.Checks() {
+		prefix := "service:" + serviceID
+		if checkID != prefix && !strings.HasPrefix(checkID, prefix+":") {
+			continue
+		}
+		if err := a.RemoveCheck(checkID, persist); err != nil {
+			return err
+		}
 	}
 
 	log.Printf("[DEBUG] agent: removed service %q", serviceID)
@@ -663,6 +675,14 @@ func (a *Agent) AddCheck(check *structs.HealthCheck, chkType *CheckType, persist
 		return fmt.Errorf("Check type is not valid")
 	}
 
+	if check.ServiceID != "" {
+		svc, ok := a.state.Services()[check.ServiceID]
+		if !ok {
+			return fmt.Errorf("ServiceID %q does not exist", check.ServiceID)
+		}
+		check.ServiceName = svc.Service
+	}
+
 	a.checkLock.Lock()
 	defer a.checkLock.Unlock()
 
@@ -864,8 +884,8 @@ func (a *Agent) loadServices(conf *Config) error {
 	// Register the services from config
 	for _, service := range conf.Services {
 		ns := service.NodeService()
-		chkType := service.CheckType()
-		if err := a.AddService(ns, chkType, false); err != nil {
+		chkTypes := service.CheckTypes()
+		if err := a.AddService(ns, chkTypes, false); err != nil {
 			return fmt.Errorf("Failed to register service '%s': %v", service.ID, err)
 		}
 	}

command/agent/agent_endpoint.go

@@ -132,17 +132,25 @@ func (s *HTTPServer) AgentRegisterService(resp http.ResponseWriter, req *http.Re
 			return nil
 		}
 
-		var check interface{}
 		for k, v := range rawMap {
-			if strings.ToLower(k) == "check" {
-				check = v
+			switch strings.ToLower(k) {
+			case "check":
+				if err := FixupCheckType(v); err != nil {
+					return err
+				}
+			case "checks":
+				chkTypes, ok := v.([]interface{})
+				if !ok {
+					return nil
+				}
+				for _, chkType := range chkTypes {
+					if err := FixupCheckType(chkType); err != nil {
+						return err
+					}
+				}
 			}
 		}
-		if check == nil {
-			return nil
-		}
-
-		return FixupCheckType(check)
+		return nil
 	}
 	if err := decodeBody(req, &args, decodeCB); err != nil {
 		resp.WriteHeader(400)
@@ -161,15 +169,17 @@ func (s *HTTPServer) AgentRegisterService(resp http.ResponseWriter, req *http.Re
 	ns := args.NodeService()
 
 	// Verify the check type
-	chkType := args.CheckType()
-	if chkType != nil && !chkType.Valid() {
-		resp.WriteHeader(400)
-		resp.Write([]byte("Must provide TTL or Script and Interval!"))
-		return nil, nil
+	chkTypes := args.CheckTypes()
+	for _, check := range chkTypes {
+		if !check.Valid() {
+			resp.WriteHeader(400)
+			resp.Write([]byte("Must provide TTL or Script and Interval!"))
+			return nil, nil
+		}
 	}
 
 	// Add the check
-	return nil, s.agent.AddService(ns, chkType, true)
+	return nil, s.agent.AddService(ns, chkTypes, true)
 }
 
 func (s *HTTPServer) AgentDeregisterService(resp http.ResponseWriter, req *http.Request) (interface{}, error) {

command/agent/agent_endpoint_test.go

@@ -430,6 +430,14 @@ func TestHTTPAgentRegisterService(t *testing.T) {
 		Check: CheckType{
 			TTL: 15 * time.Second,
 		},
+		Checks: CheckTypes{
+			&CheckType{
+				TTL: 20 * time.Second,
+			},
+			&CheckType{
+				TTL: 30 * time.Second,
+			},
+		},
 	}
 	req.Body = encodeReq(args)
 
@@ -447,12 +455,13 @@ func TestHTTPAgentRegisterService(t *testing.T) {
 	}
 
 	// Ensure we have a check mapping
-	if _, ok := srv.agent.state.Checks()["service:test"]; !ok {
-		t.Fatalf("missing test check")
+	checks := srv.agent.state.Checks()
+	if len(checks) != 3 {
+		t.Fatalf("bad: %v", checks)
 	}
 
-	if _, ok := srv.agent.checkTTLs["service:test"]; !ok {
-		t.Fatalf("missing test check ttl")
+	if len(srv.agent.checkTTLs) != 3 {
+		t.Fatalf("missing test check ttls: %v", srv.agent.checkTTLs)
 	}
 }
 
@@ -683,3 +692,60 @@ func TestHTTPAgent_DisableNodeMaintenance(t *testing.T) {
 		t.Fatalf("should have removed maintenance check")
 	}
 }
+
+func TestHTTPAgentRegisterServiceCheck(t *testing.T) {
+	dir, srv := makeHTTPServer(t)
+	defer os.RemoveAll(dir)
+	defer srv.Shutdown()
+	defer srv.agent.Shutdown()
+
+	// First register the service
+	req, err := http.NewRequest("GET", "/v1/agent/service/register", nil)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	args := &ServiceDefinition{
+		Name: "memcache",
+		Port: 8000,
+		Check: CheckType{
+			TTL: 15 * time.Second,
+		},
+	}
+	req.Body = encodeReq(args)
+
+	if _, err := srv.AgentRegisterService(nil, req); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Now register an additional check
+	req, err = http.NewRequest("GET", "/v1/agent/check/register", nil)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	checkArgs := &CheckDefinition{
+		Name:      "memcache_check2",
+		ServiceID: "memcache",
+		CheckType: CheckType{
+			TTL: 15 * time.Second,
+		},
+	}
+	req.Body = encodeReq(checkArgs)
+
+	if _, err := srv.AgentRegisterCheck(nil, req); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Ensure we have a check mapping
+	result := srv.agent.state.Checks()
+	if _, ok := result["service:memcache"]; !ok {
+		t.Fatalf("missing memcached check")
+	}
+	if _, ok := result["memcache_check2"]; !ok {
+		t.Fatalf("missing memcache_check2 check")
+	}
+
+	// Make sure the new check is associated with the service
+	if result["memcache_check2"].ServiceID != "memcache" {
+		t.Fatalf("bad: %#v", result["memcached_check2"])
+	}
+}