[HA Agent] Refactor is_leader to state (DataDog#31777)

comp/haagent/def/component.go

@@ -16,8 +16,8 @@ type Component interface {
 	// GetGroup returns the value of ha_agent.group
 	GetGroup() string
 
-	// IsLeader returns true if the current Agent is leader
-	IsLeader() bool
+	// GetState returns current HA agent state
+	GetState() State
 
 	// SetLeader takes the leader agent hostname as input, if it matches the current agent hostname,
 	// the isLeader state is set to true, otherwise false.

comp/haagent/impl/utils_test.go → comp/haagent/def/state.go

@@ -3,15 +3,16 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/).
 // Copyright 2024-present Datadog, Inc.
 
-package haagentimpl
+package haagent
 
-import (
-	"testing"
+// State type for HA Agent State
+type State string
 
-	"github.com/stretchr/testify/assert"
+const (
+	// Active HA Agent state
+	Active State = "active"
+	// Standby HA Agent state
+	Standby State = "standby"
+	// Unknown HA Agent state
+	Unknown State = "unknown"
 )
-
-func Test_leaderStatusToRole(t *testing.T) {
-	assert.Equal(t, "leader", leaderStateToRole(true))
-	assert.Equal(t, "follower", leaderStateToRole(false))
-}

comp/haagent/impl/haagent.go

@@ -10,6 +10,7 @@ import (
 	"encoding/json"
 
 	log "github.com/DataDog/datadog-agent/comp/core/log/def"
+	haagent "github.com/DataDog/datadog-agent/comp/haagent/def"
 	"github.com/DataDog/datadog-agent/pkg/remoteconfig/state"
 	"github.com/DataDog/datadog-agent/pkg/util/hostname"
 	"go.uber.org/atomic"
@@ -18,14 +19,14 @@ import (
 type haAgentImpl struct {
 	log            log.Component
 	haAgentConfigs *haAgentConfigs
-	isLeader       *atomic.Bool
+	state          *atomic.String
 }
 
 func newHaAgentImpl(log log.Component, haAgentConfigs *haAgentConfigs) *haAgentImpl {
 	return &haAgentImpl{
 		log:            log,
 		haAgentConfigs: haAgentConfigs,
-		isLeader:       atomic.NewBool(false),
+		state:          atomic.NewString(string(haagent.Unknown)),
 	}
 }
 
@@ -37,8 +38,8 @@ func (h *haAgentImpl) GetGroup() string {
 	return h.haAgentConfigs.group
 }
 
-func (h *haAgentImpl) IsLeader() bool {
-	return h.isLeader.Load()
+func (h *haAgentImpl) GetState() haagent.State {
+	return haagent.State(h.state.Load())
 }
 
 func (h *haAgentImpl) SetLeader(leaderAgentHostname string) {
@@ -47,21 +48,29 @@ func (h *haAgentImpl) SetLeader(leaderAgentHostname string) {
 		h.log.Warnf("error getting the hostname: %v", err)
 		return
 	}
-	newIsLeader := agentHostname == leaderAgentHostname
-	prevIsLeader := h.isLeader.Load()
-	if newIsLeader != prevIsLeader {
-		h.log.Infof("agent role switched from %s to %s", leaderStateToRole(prevIsLeader), leaderStateToRole(newIsLeader))
-		h.isLeader.Store(newIsLeader)
+
+	var newState haagent.State
+	if agentHostname == leaderAgentHostname {
+		newState = haagent.Active
+	} else {
+		newState = haagent.Standby
+	}
+
+	prevState := h.GetState()
+
+	if newState != prevState {
+		h.log.Infof("agent state switched from %s to %s", prevState, newState)
+		h.state.Store(string(newState))
 	} else {
-		h.log.Debugf("agent role not changed (current role: %s)", leaderStateToRole(prevIsLeader))
+		h.log.Debugf("agent state not changed (current state: %s)", prevState)
 	}
 }
 
 // ShouldRunIntegration return true if the agent integrations should to run.
 // When ha-agent is disabled, the agent behave as standalone agent (non HA) and will always run all integrations.
 func (h *haAgentImpl) ShouldRunIntegration(integrationName string) bool {
 	if h.Enabled() && validHaIntegrations[integrationName] {
-		return h.isLeader.Load()
+		return h.GetState() == haagent.Active
 	}
 	return true
 }

comp/haagent/impl/haagent_test.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/DataDog/datadog-agent/comp/core/config"
 	logmock "github.com/DataDog/datadog-agent/comp/core/log/mock"
+	haagent "github.com/DataDog/datadog-agent/comp/haagent/def"
 	"github.com/DataDog/datadog-agent/pkg/remoteconfig/state"
 	"github.com/DataDog/datadog-agent/pkg/util/fxutil"
 	"github.com/stretchr/testify/assert"
@@ -62,11 +63,13 @@ func Test_IsLeader_SetLeader(t *testing.T) {
 	}
 	haAgent := newTestHaAgentComponent(t, agentConfigs).Comp
 
+	assert.Equal(t, haagent.Unknown, haAgent.GetState())
+
 	haAgent.SetLeader("another-agent")
-	assert.False(t, haAgent.IsLeader())
+	assert.Equal(t, haagent.Standby, haAgent.GetState())
 
 	haAgent.SetLeader("my-agent-hostname")
-	assert.True(t, haAgent.IsLeader())
+	assert.Equal(t, haagent.Active, haAgent.GetState())
 }
 
 func Test_RCListener(t *testing.T) {

comp/haagent/mock/mock.go

@@ -21,6 +21,7 @@ type mockHaAgent struct {
 
 	group   string
 	enabled bool
+	state   haagent.State
 }
 
 func (m *mockHaAgent) GetGroup() string {
@@ -34,7 +35,7 @@ func (m *mockHaAgent) Enabled() bool {
 func (m *mockHaAgent) SetLeader(_ string) {
 }
 
-func (m *mockHaAgent) IsLeader() bool { return false }
+func (m *mockHaAgent) GetState() haagent.State { return haagent.Standby }
 
 func (m *mockHaAgent) SetGroup(group string) {
 	m.group = group
@@ -43,6 +44,9 @@ func (m *mockHaAgent) SetGroup(group string) {
 func (m *mockHaAgent) SetEnabled(enabled bool) {
 	m.enabled = enabled
 }
+func (m *mockHaAgent) SetState(state haagent.State) {
+	m.state = state
+}
 
 func (m *mockHaAgent) ShouldRunIntegration(_ string) bool {
 	return true
@@ -54,6 +58,7 @@ type Component interface {
 
 	SetGroup(string)
 	SetEnabled(bool)
+	SetState(haagent.State)
 }
 
 // NewMockHaAgent returns a new Mock