adds the command execution back

Author: ncode

README.md

@@ -86,6 +86,5 @@ election:
 
 - Write more tests
 - Add more examples
-- Re-enable the hooks on state change
 - Allow to pre-define the preferred leader
 - Update the docks with the lock delays and timeouts

examples/config/ballot1.yaml

@@ -11,5 +11,5 @@ election:
       serviceChecks:
         - service:election1
       primaryTag: primary
-      execOnPromote: '/usr/bin/say I am da boss'
-      execOnDemote: '/usr/bin/say I am no longer da boss'
+      execOnPromote: '/usr/bin/say primary'
+      execOnDemote: '/usr/bin/say secondary'

examples/config/ballot2.yaml

@@ -11,5 +11,5 @@ election:
       serviceChecks:
         - service:election2
       primaryTag: primary
-      execOnPromote: '/usr/bin/say I am da boss'
-      execOnDemote: '/usr/bin/say I am no longer da boss'
+      execOnPromote: '/usr/bin/say primary'
+      execOnDemote: '/usr/bin/say secondary'

examples/consul/my-service1.json

@@ -11,9 +11,5 @@
     "Interval": "10s",
     "Timeout": "5s"
   },
-  "EnableTagOverride": true,
-  "Weights": {
-    "Passing": 10,
-    "Warning": 1
-  }
+  "EnableTagOverride": true
 }

examples/consul/my-service2.json

@@ -11,9 +11,5 @@
     "Interval": "10s",
     "Timeout": "5s"
   },
-  "EnableTagOverride": true,
-  "Weights": {
-    "Passing": 10,
-    "Warning": 1
-  }
+  "EnableTagOverride": true
 }

internal/ballot/ballot.go

@@ -33,6 +33,12 @@ type ElectionPayload struct {
 	SessionID string
 }
 
+type commandExecutor struct{}
+
+func (c *commandExecutor) Command(name string, arg ...string) *exec.Cmd {
+	return exec.Command(name, arg...)
+}
+
 // New returns a new Ballot instance.
 func New(ctx context.Context, name string) (b *Ballot, err error) {
 	if ctx == nil {
@@ -55,6 +61,7 @@ func New(ctx context.Context, name string) (b *Ballot, err error) {
 	b.leader.Store(false)
 	b.Token = consulConfig.Token
 	b.ctx = ctx
+	b.executor = &commandExecutor{}
 
 	b.Name = name
 	if b.LockDelay == 0 {
@@ -84,7 +91,7 @@ type Ballot struct {
 	leader        atomic.Bool     `mapstructure:"-"`
 	client        ConsulClient    `mapstructure:"-"`
 	ctx           context.Context `mapstructure:"-"`
-	exec          CommandExecutor `mapstructure:"-"`
+	executor      CommandExecutor `mapstructure:"-"`
 }
 
 // Copy *api.AgentService to *api.AgentServiceRegistration
@@ -154,15 +161,15 @@ func (b *Ballot) runCommand(command string, electionPayload *ElectionPayload) ([
 	if err != nil {
 		return nil, err
 	}
-	cmd := b.exec.Command(args[0], args[1:]...)
+	cmd := b.executor.Command(args[0], args[1:]...)
 	cmd.Env = append(cmd.Env, fmt.Sprintf("ADDRESS=%s", electionPayload.Address))
 	cmd.Env = append(cmd.Env, fmt.Sprintf("PORT=%d", electionPayload.Port))
 	cmd.Env = append(cmd.Env, fmt.Sprintf("SESSIONID=%s", electionPayload.SessionID))
 	return cmd.Output()
 }
 
 // updateServiceTags updates the service tags.
-func (b *Ballot) updateServiceTags() error {
+func (b *Ballot) updateServiceTags(isLeader bool) error {
 	service, _, err := b.getService()
 	if err != nil {
 		return err
@@ -175,10 +182,10 @@ func (b *Ballot) updateServiceTags() error {
 	hasPrimaryTag := slices.Contains(registration.Tags, b.PrimaryTag)
 
 	// Update tags based on leadership status
-	if b.IsLeader() && !hasPrimaryTag {
+	if isLeader && !hasPrimaryTag {
 		// Add primary tag if not present and this node is the leader
 		registration.Tags = append(registration.Tags, b.PrimaryTag)
-	} else if !b.IsLeader() && hasPrimaryTag {
+	} else if !isLeader && hasPrimaryTag {
 		// Remove primary tag if present and this node is not the leader
 		index := slices.Index(registration.Tags, b.PrimaryTag)
 		registration.Tags = append(registration.Tags[:index], registration.Tags[index+1:]...)
@@ -187,6 +194,35 @@ func (b *Ballot) updateServiceTags() error {
 		return nil
 	}
 
+	// Run the command associated with the new leadership status
+	var command string
+	if isLeader {
+		command = b.ExecOnPromote
+	} else {
+		command = b.ExecOnDemote
+	}
+	if command != "" && b.executor != nil {
+		go func(isLeader bool, command string) {
+			// Run the command in a separate goroutine
+			ctx, cancel := context.WithTimeout(b.ctx, (b.TTL+b.LockDelay)*2)
+			defer cancel()
+			payload, err := b.waitForNextValidSessionData(ctx)
+			output, err := b.runCommand(command, payload)
+			if err != nil {
+				log.WithFields(log.Fields{
+					"caller":   "updateLeadershipStatus",
+					"isLeader": isLeader,
+					"error":    err,
+				}).Error("failed to run command")
+			}
+			log.WithFields(log.Fields{
+				"caller":   "updateLeadershipStatus",
+				"isLeader": isLeader,
+				"output":   string(output),
+			}).Info("ran command")
+		}(isLeader, command)
+	}
+
 	// Log the updated tags
 	log.WithFields(log.Fields{
 		"caller":  "updateServiceTags",
@@ -345,7 +381,7 @@ func (b *Ballot) updateLeadershipStatus(isLeader bool) error {
 	b.leader.Store(isLeader)
 
 	// Update service tags based on leadership status
-	err := b.updateServiceTags()
+	err := b.updateServiceTags(isLeader)
 	if err != nil {
 		return err
 	}
@@ -447,6 +483,25 @@ func (b *Ballot) IsLeader() bool {
 	return b.leader.Load() && b.sessionID.Load() != nil
 }
 
+func (b *Ballot) waitForNextValidSessionData(ctx context.Context) (data *ElectionPayload, err error) {
+	ticker := time.NewTicker(1 * time.Second)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			data, err := b.getSessionData()
+			if err != nil {
+				return data, err
+			}
+			if data != nil {
+				return data, nil
+			}
+		case <-ctx.Done():
+			return data, ctx.Err()
+		}
+	}
+}
+
 func (b *Ballot) getSessionData() (data *ElectionPayload, err error) {
 	sessionKey, _, err := b.client.KV().Get(b.Key, nil)
 	if err != nil {

internal/ballot/ballot_test.go

@@ -103,7 +103,7 @@ func TestRunCommand(t *testing.T) {
 
 	// Create a Ballot instance with the mock executor
 	b := &Ballot{
-		exec: mockExecutor,
+		executor: mockExecutor,
 	}
 
 	// Define the command to run