diff --git a/agent/local/state.go b/agent/local/state.go index 9f2aebb29..69e136816 100644 --- a/agent/local/state.go +++ b/agent/local/state.go @@ -206,21 +206,13 @@ func (l *State) AddService(service *structs.NodeService, token string) error { service.ID = service.Service } - l.AddServiceState(&ServiceState{ + l.SetServiceState(&ServiceState{ Service: service, Token: token, }) return nil } -func (l *State) AddServiceState(s *ServiceState) { - l.Lock() - defer l.Unlock() - - l.services[s.Service.ID] = s - l.TriggerSyncChanges() -} - // 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 { @@ -285,6 +277,17 @@ func (l *State) ServiceState(id string) *ServiceState { return s.Clone() } +// SetServiceState is used to overwrite a raw service state with the given +// state. This method is safe to be called concurrently but should only be used +// during testing. You should most likely call AddService instead. +func (l *State) SetServiceState(s *ServiceState) { + l.Lock() + defer l.Unlock() + + l.services[s.Service.ID] = s + l.TriggerSyncChanges() +} + // ServiceStates returns a shallow copy of all service state records. // The service record still points to the original service record and // must not be modified. @@ -345,21 +348,13 @@ func (l *State) AddCheck(check *structs.HealthCheck, token string) error { // hard-set the node name check.Node = l.config.NodeName - l.AddCheckState(&CheckState{ + l.SetCheckState(&CheckState{ Check: check, Token: token, }) return nil } -func (l *State) AddCheckState(c *CheckState) { - l.Lock() - defer l.Unlock() - - l.checks[c.Check.CheckID] = c - l.TriggerSyncChanges() -} - // RemoveCheck is used to remove a health check from the local state. // The agent will make a best effort to ensure it is deregistered // todo(fs): RemoveService returns an error for a non-existant service. RemoveCheck should as well. @@ -484,6 +479,17 @@ func (l *State) CheckState(id types.CheckID) *CheckState { return c.Clone() } +// SetCheckState is used to overwrite a raw check state with the given +// state. This method is safe to be called concurrently but should only be used +// during testing. You should most likely call AddCheck instead. +func (l *State) SetCheckState(c *CheckState) { + l.Lock() + defer l.Unlock() + + l.checks[c.Check.CheckID] = c + l.TriggerSyncChanges() +} + // CheckStates returns a shallow copy of all health check state records. // The health check records and the deferred checks still point to // the original values and must not be modified. diff --git a/agent/local/state_test.go b/agent/local/state_test.go index b092998aa..d40300a4c 100644 --- a/agent/local/state_test.go +++ b/agent/local/state_test.go @@ -108,7 +108,7 @@ func TestAgentAntiEntropy_Services(t *testing.T) { Tags: []string{}, Port: 11211, } - a.State.AddServiceState(&local.ServiceState{ + a.State.SetServiceState(&local.ServiceState{ Service: srv6, InSync: true, }) @@ -679,7 +679,7 @@ func TestAgentAntiEntropy_Checks(t *testing.T) { Name: "cache", Status: api.HealthPassing, } - a.State.AddCheckState(&local.CheckState{ + a.State.SetCheckState(&local.CheckState{ Check: chk5, InSync: true, })