Register a service and its check atomically in the local state

Prevent race between register and deregister requests by saving them
together in the local state on registration.
Also adds more cleaning in case of failure when registering services
/ checks.

Author: Thibault Gilles

agent/agent.go

@@ -20,11 +20,11 @@ import (
 
 	"google.golang.org/grpc"
 
-	"github.com/armon/go-metrics"
+	metrics "github.com/armon/go-metrics"
 	"github.com/hashicorp/consul/acl"
 	"github.com/hashicorp/consul/agent/ae"
 	"github.com/hashicorp/consul/agent/cache"
-	"github.com/hashicorp/consul/agent/cache-types"
+	cachetype "github.com/hashicorp/consul/agent/cache-types"
 	"github.com/hashicorp/consul/agent/checks"
 	"github.com/hashicorp/consul/agent/config"
 	"github.com/hashicorp/consul/agent/consul"
@@ -42,8 +42,8 @@ import (
 	"github.com/hashicorp/consul/logger"
 	"github.com/hashicorp/consul/types"
 	"github.com/hashicorp/consul/watch"
-	"github.com/hashicorp/go-multierror"
-	"github.com/hashicorp/go-uuid"
+	multierror "github.com/hashicorp/go-multierror"
+	uuid "github.com/hashicorp/go-uuid"
 	"github.com/hashicorp/memberlist"
 	"github.com/hashicorp/raft"
 	"github.com/hashicorp/serf/serf"
@@ -1911,15 +1911,7 @@ func (a *Agent) AddService(service *structs.NodeService, chkTypes []*structs.Che
 	snap := a.snapshotCheckState()
 	defer a.restoreCheckState(snap)
 
-	// Add the service
-	a.State.AddService(service, token)
-
-	// Persist the service to a file
-	if persist && a.config.DataDir != "" {
-		if err := a.persistService(service); err != nil {
-			return err
-		}
-	}
+	var checks []*structs.HealthCheck
 
 	// Create an associated health check
 	for i, chkType := range chkTypes {
@@ -1947,14 +1939,81 @@ func (a *Agent) AddService(service *structs.NodeService, chkTypes []*structs.Che
 		if chkType.Status != "" {
 			check.Status = chkType.Status
 		}
-		if err := a.AddCheck(check, chkType, persist, token, source); err != nil {
+
+		checks = append(checks, check)
+	}
+
+	// cleanup, store the ids of services and checks that weren't previously
+	// registered so we clean them up if somthing fails halfway through the
+	// process.
+	var cleanupServices []string
+	var cleanupChecks []types.CheckID
+
+	if s := a.State.Service(service.ID); s == nil {
+		cleanupServices = append(cleanupServices, service.ID)
+	}
+
+	for _, check := range checks {
+		if c := a.State.Check(check.CheckID); c == nil {
+			cleanupChecks = append(cleanupChecks, check.CheckID)
+		}
+	}
+
+	err := a.State.AddServiceWithChecks(service, checks, token)
+	if err != nil {
+		a.cleanupRegistration(cleanupServices, cleanupChecks)
+		return err
+	}
+
+	for i := range checks {
+		if err := a.addCheck(checks[i], chkTypes[i], service, persist, token, source); err != nil {
+			a.cleanupRegistration(cleanupServices, cleanupChecks)
+			return err
+		}
+
+		if persist && a.config.DataDir != "" {
+			if err := a.persistCheck(checks[i], chkTypes[i]); err != nil {
+				a.cleanupRegistration(cleanupServices, cleanupChecks)
+				return err
+
+			}
+		}
+	}
+
+	// Persist the service to a file
+	if persist && a.config.DataDir != "" {
+		if err := a.persistService(service); err != nil {
+			a.cleanupRegistration(cleanupServices, cleanupChecks)
 			return err
 		}
 	}
 
 	return nil
 }
 
+// cleanupRegistration is called on  registration error to ensure no there are no
+// leftovers after a partial failure
+func (a *Agent) cleanupRegistration(serviceIDs []string, checksIDs []types.CheckID) {
+	for _, s := range serviceIDs {
+		if err := a.State.RemoveService(s); err != nil {
+			a.logger.Printf("[ERR] consul: service registration: cleanup: failed to remove service %s: %s", s, err)
+		}
+		if err := a.purgeService(s); err != nil {
+			a.logger.Printf("[ERR] consul: service registration: cleanup: failed to purge service %s file: %s", s, err)
+		}
+	}
+
+	for _, c := range checksIDs {
+		a.cancelCheckMonitors(c)
+		if err := a.State.RemoveCheck(c); err != nil {
+			a.logger.Printf("[ERR] consul: service registration: cleanup: failed to remove check %s: %s", c, err)
+		}
+		if err := a.purgeCheck(c); err != nil {
+			a.logger.Printf("[ERR] consul: service registration: cleanup: failed to purge check %s file: %s", c, err)
+		}
+	}
+}
+
 // RemoveService is used to remove a service entry.
 // The agent will make a best effort to ensure it is deregistered
 func (a *Agent) RemoveService(serviceID string, persist bool) error {
@@ -2018,6 +2077,44 @@ func (a *Agent) RemoveService(serviceID string, persist bool) error {
 // ensure it is registered. The Check may include a CheckType which
 // is used to automatically update the check status
 func (a *Agent) AddCheck(check *structs.HealthCheck, chkType *structs.CheckType, persist bool, token string, source configSource) error {
+	var service *structs.NodeService
+
+	if check.ServiceID != "" {
+		service = a.State.Service(check.ServiceID)
+		if service == nil {
+			return fmt.Errorf("ServiceID %q does not exist", check.ServiceID)
+		}
+	}
+
+	// snapshot the current state of the health check to avoid potential flapping
+	existing := a.State.Check(check.CheckID)
+	defer func() {
+		if existing != nil {
+			a.State.UpdateCheck(check.CheckID, existing.Status, existing.Output)
+		}
+	}()
+
+	err := a.addCheck(check, chkType, service, persist, token, source)
+	if err != nil {
+		a.State.RemoveCheck(check.CheckID)
+		return err
+	}
+
+	// Add to the local state for anti-entropy
+	err = a.State.AddCheck(check, token)
+	if err != nil {
+		return err
+	}
+
+	// Persist the check
+	if persist && a.config.DataDir != "" {
+		return a.persistCheck(check, chkType)
+	}
+
+	return nil
+}
+
+func (a *Agent) addCheck(check *structs.HealthCheck, chkType *structs.CheckType, service *structs.NodeService, persist bool, token string, source configSource) error {
 	if check.CheckID == "" {
 		return fmt.Errorf("CheckID missing")
 	}
@@ -2039,12 +2136,8 @@ func (a *Agent) AddCheck(check *structs.HealthCheck, chkType *structs.CheckType,
 	}
 
 	if check.ServiceID != "" {
-		s := a.State.Service(check.ServiceID)
-		if s == nil {
-			return fmt.Errorf("ServiceID %q does not exist", check.ServiceID)
-		}
-		check.ServiceName = s.Service
-		check.ServiceTags = s.Tags
+		check.ServiceName = service.Service
+		check.ServiceTags = service.Tags
 	}
 
 	a.checkLock.Lock()
@@ -2265,18 +2358,6 @@ func (a *Agent) AddCheck(check *structs.HealthCheck, chkType *structs.CheckType,
 		}
 	}
 
-	// Add to the local state for anti-entropy
-	err := a.State.AddCheck(check, token)
-	if err != nil {
-		a.cancelCheckMonitors(check.CheckID)
-		return err
-	}
-
-	// Persist the check
-	if persist && a.config.DataDir != "" {
-		return a.persistCheck(check, chkType)
-	}
-
 	return nil
 }
 

agent/agent_test.go

@@ -1839,7 +1839,7 @@ func TestAgent_PersistCheck(t *testing.T) {
 		t.Fatalf("err: %s", err)
 	}
 	if !bytes.Equal(expected, content) {
-		t.Fatalf("bad: %s", string(content))
+		t.Fatalf("bad: %s != %s", string(content), expected)
 	}
 
 	// Updates the check definition on disk

agent/local/state.go

@@ -265,6 +265,12 @@ func (l *State) serviceToken(id string) string {
 // This entry is persistent and the agent will make a best effort to
 // ensure it is registered
 func (l *State) AddService(service *structs.NodeService, token string) error {
+	l.Lock()
+	defer l.Unlock()
+	return l.addServiceLocked(service, token)
+}
+
+func (l *State) addServiceLocked(service *structs.NodeService, token string) error {
 	if service == nil {
 		return fmt.Errorf("no service")
 	}
@@ -274,13 +280,31 @@ func (l *State) AddService(service *structs.NodeService, token string) error {
 		service.ID = service.Service
 	}
 
-	l.SetServiceState(&ServiceState{
+	l.setServiceStateLocked(&ServiceState{
 		Service: service,
 		Token:   token,
 	})
 	return nil
 }
 
+// AddServiceWithChecks adds a service in its check in the local state atomically
+func (l *State) AddServiceWithChecks(service *structs.NodeService, checks []*structs.HealthCheck, token string) error {
+	l.Lock()
+	defer l.Unlock()
+
+	if err := l.addServiceLocked(service, token); err != nil {
+		return err
+	}
+
+	for _, check := range checks {
+		if err := l.addCheckLocked(check, token); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // RemoveService is used to remove a service entry from the local state.
 // The agent will make a best effort to ensure it is deregistered.
 func (l *State) RemoveService(id string) error {
@@ -358,6 +382,10 @@ func (l *State) SetServiceState(s *ServiceState) {
 	l.Lock()
 	defer l.Unlock()
 
+	l.setServiceStateLocked(s)
+}
+
+func (l *State) setServiceStateLocked(s *ServiceState) {
 	s.WatchCh = make(chan struct{})
 
 	old, hasOld := l.services[s.Service.ID]
@@ -414,6 +442,13 @@ func (l *State) checkToken(id types.CheckID) string {
 // This entry is persistent and the agent will make a best effort to
 // ensure it is registered
 func (l *State) AddCheck(check *structs.HealthCheck, token string) error {
+	l.Lock()
+	defer l.Unlock()
+
+	return l.addCheckLocked(check, token)
+}
+
+func (l *State) addCheckLocked(check *structs.HealthCheck, token string) error {
 	if check == nil {
 		return fmt.Errorf("no check")
 	}
@@ -427,14 +462,14 @@ func (l *State) AddCheck(check *structs.HealthCheck, token string) error {
 
 	// if there is a serviceID associated with the check, make sure it exists before adding it
 	// NOTE - This logic may be moved to be handled within the Agent's Addcheck method after a refactor
-	if check.ServiceID != "" && l.Service(check.ServiceID) == nil {
+	if _, ok := l.services[check.ServiceID]; check.ServiceID != "" && !ok {
 		return fmt.Errorf("Check %q refers to non-existent service %q", check.CheckID, check.ServiceID)
 	}
 
 	// hard-set the node name
 	check.Node = l.config.NodeName
 
-	l.SetCheckState(&CheckState{
+	l.setCheckStateLocked(&CheckState{
 		Check: check,
 		Token: token,
 	})
@@ -620,6 +655,10 @@ func (l *State) SetCheckState(c *CheckState) {
 	l.Lock()
 	defer l.Unlock()
 
+	l.setCheckStateLocked(c)
+}
+
+func (l *State) setCheckStateLocked(c *CheckState) {
 	l.checks[c.Check.CheckID] = c
 	l.TriggerSyncChanges()
 }