Merge pull request #281 from rstudio/feature/cluster-health-callbacks

Add cluster-health observability callbacks to leader and task handler

Author: jonyoder

pkg/rselection/impls/pgx/leader.go

@@ -63,6 +63,9 @@ type PgxLeader struct {
 	// lead() goroutine.
 	unhealthySince time.Time
 
+	onIntegrityResult func(err error)
+	onPingResult      func(success bool, t time.Time)
+
 	// now returns the current time; overridable in tests. Use timeNow().
 	now func() time.Time
 
@@ -90,25 +93,38 @@ type PgxLeaderConfig struct {
 	SweepInterval   time.Duration
 	MaxPingAge      time.Duration
 	StepDownTimeout time.Duration
+	// OnIntegrityResult, if set, is called with the result of each cluster
+	// integrity evaluation (nil means healthy). It may run concurrently with
+	// OnPingResult, so it must be non-blocking, panic-free, and safe for
+	// concurrent use.
+	OnIntegrityResult func(err error)
+	// OnPingResult, if set, is called after each leader ping attempt with
+	// whether it succeeded and the time it ran. Periodic pings are dispatched on
+	// their own goroutines, so it may run concurrently with itself and with
+	// OnIntegrityResult; it must be non-blocking, panic-free, and safe for
+	// concurrent use.
+	OnPingResult func(success bool, t time.Time)
 }
 
 func NewPgxLeader(cfg PgxLeaderConfig) *PgxLeader {
 	return &PgxLeader{
-		store:           cfg.Store,
-		awb:             cfg.Broadcaster,
-		notify:          cfg.Notifier,
-		taskHandler:     cfg.TaskHandler,
-		chLeader:        cfg.LeaderChannel,
-		chFollower:      cfg.FollowerChannel,
-		chMessages:      cfg.MessagesChannel,
-		address:         cfg.Address,
-		stop:            cfg.StopChan,
-		ping:            cfg.PingInterval,
-		sweep:           cfg.SweepInterval,
-		maxPingAge:      cfg.MaxPingAge,
-		stepDownTimeout: cfg.StepDownTimeout,
-		now:             time.Now,
-		mutex:           sync.RWMutex{},
+		store:             cfg.Store,
+		awb:               cfg.Broadcaster,
+		notify:            cfg.Notifier,
+		taskHandler:       cfg.TaskHandler,
+		chLeader:          cfg.LeaderChannel,
+		chFollower:        cfg.FollowerChannel,
+		chMessages:        cfg.MessagesChannel,
+		address:           cfg.Address,
+		stop:              cfg.StopChan,
+		ping:              cfg.PingInterval,
+		sweep:             cfg.SweepInterval,
+		maxPingAge:        cfg.MaxPingAge,
+		stepDownTimeout:   cfg.StepDownTimeout,
+		onIntegrityResult: cfg.OnIntegrityResult,
+		onPingResult:      cfg.OnPingResult,
+		now:               time.Now,
+		mutex:             sync.RWMutex{},
 	}
 }
 
@@ -276,6 +292,9 @@ func (p *PgxLeader) clusterIntegrityErr() error {
 // (when that timeout is non-zero). Called from the lead() goroutine only.
 func (p *PgxLeader) evaluateHealth() bool {
 	err := p.clusterIntegrityErr()
+	if p.onIntegrityResult != nil {
+		p.onIntegrityResult(err)
+	}
 	if err == nil {
 		if !p.unhealthySince.IsZero() {
 			slog.Info("Cluster integrity recovered", "degradedFor", p.timeNow().Sub(p.unhealthySince))
@@ -320,6 +339,20 @@ func (p *PgxLeader) verify(vCh chan bool) {
 // pingNodes sends a ping request out on the follower channel. All online cluster
 // nodes should receive this message and respond with a ping response.
 func (p *PgxLeader) pingNodes(ctx context.Context) {
+	success := p.sendPings(ctx)
+
+	p.pingSuccessLock.Lock()
+	p.pingSuccess = success
+	p.pingSuccessLock.Unlock()
+
+	if p.onPingResult != nil {
+		p.onPingResult(success, p.timeNow())
+	}
+}
+
+// sendPings publishes the ping notifications on the follower and leader
+// channels and reports whether both sends succeeded.
+func (p *PgxLeader) sendPings(ctx context.Context) bool {
 	req := &electiontypes.ClusterNotification{
 		GuidVal:     uuid.New().String(),
 		MessageType: electiontypes.ClusterMessageTypePing,
@@ -328,30 +361,23 @@ func (p *PgxLeader) pingNodes(ctx context.Context) {
 	b, err := json.Marshal(req)
 	if err != nil {
 		slog.Error("Error marshaling notification to JSON", "error", err)
-		return
+		return false
 	}
 
-	p.pingSuccessLock.Lock()
-	defer p.pingSuccessLock.Unlock()
-	p.pingSuccess = true
-
 	slog.Log(ctx, LevelTrace, fmt.Sprintf("Leader pinging nodes on follower channel %s", p.chFollower))
-	err = p.notify.Notify(ctx, p.chFollower, b)
-	if err != nil {
-		p.pingSuccess = false
+	if err = p.notify.Notify(ctx, p.chFollower, b); err != nil {
 		slog.Error("Leader error pinging followers", "error", err)
-		return
+		return false
 	}
 
 	// This will ensure that the leader is tracked as part of the nodes in the cluster, and it will force
 	// duplicate leaders to surrender the position.
 	slog.Log(ctx, LevelTrace, fmt.Sprintf("Leader pinging itself on leader channel %s", p.chLeader))
-	err = p.notify.Notify(ctx, p.chLeader, b)
-	if err != nil {
-		p.pingSuccess = false
+	if err = p.notify.Notify(ctx, p.chLeader, b); err != nil {
 		slog.Error("Leader error pinging leaders", "error", err)
-		return
+		return false
 	}
+	return true
 }
 
 func (p *PgxLeader) handleNodesRequest(ctx context.Context, cn *electiontypes.ClusterNodesRequest) {

pkg/rselection/impls/pgx/leader_recovery_test.go

@@ -3,6 +3,7 @@ package pgxelection
 // Copyright (C) 2026 by Posit Software, PBC
 
 import (
+	"context"
 	"errors"
 	"time"
 
@@ -141,3 +142,55 @@ func (s *RecoverySuite) TestStepDownDisabledByZeroTimeout(c *check.C) {
 	leader.setClockTEST(start.Add(100 * time.Hour))
 	c.Assert(leader.evaluateHealth(), check.Equals, false)
 }
+
+func (s *RecoverySuite) TestEvaluateHealthFiresIntegrityCallback(c *check.C) {
+	var got []error
+	store := &fakeStore{nodes: map[string]*electiontypes.ClusterNode{
+		"a": {Name: "a", IP: "10.0.0.1"},
+	}}
+	leader := newTestLeader(store, map[string]*electiontypes.ClusterNode{
+		"a_10.0.0.1": {Name: "a", IP: "10.0.0.1"},
+	})
+	leader.onIntegrityResult = func(err error) { got = append(got, err) }
+
+	// Healthy check -> callback receives nil.
+	leader.evaluateHealth()
+	c.Assert(got, check.HasLen, 1)
+	c.Assert(got[0], check.IsNil)
+
+	// Unhealthy check -> callback receives a non-nil error.
+	leader.pingSuccess = false
+	leader.evaluateHealth()
+	c.Assert(got, check.HasLen, 2)
+	c.Assert(got[1], check.NotNil)
+}
+
+func (s *RecoverySuite) TestPingNodesFiresPingCallback(c *check.C) {
+	var results []bool
+	var times []time.Time
+	fixed := time.Unix(1_700_000_000, 0)
+	leader := &PgxLeader{
+		address:    "leader",
+		chLeader:   "leader",
+		chFollower: "follower",
+		notify:     FakePgNotifier{},
+		now:        func() time.Time { return fixed },
+	}
+	leader.onPingResult = func(success bool, t time.Time) {
+		results = append(results, success)
+		times = append(times, t)
+	}
+	ctx := context.Background()
+
+	// Successful ping.
+	leader.pingNodes(ctx)
+	c.Assert(results, check.DeepEquals, []bool{true})
+
+	// Failed ping.
+	leader.notify = FakePgNotifier{notifyErr: errors.New("down")}
+	leader.pingNodes(ctx)
+	c.Assert(results, check.DeepEquals, []bool{true, false})
+
+	// The callback receives the time the ping ran, from the injectable clock.
+	c.Assert(times, check.DeepEquals, []time.Time{fixed, fixed})
+}

pkg/rselection/taskhandler.go

@@ -50,20 +50,27 @@ type Schedule interface {
 }
 
 type GenericTaskHandler struct {
-	tasks  map[string]Task
-	cancel context.CancelFunc
-	verify chan chan bool
-	mutex  sync.RWMutex
-	role   string
+	tasks     map[string]Task
+	cancel    context.CancelFunc
+	verify    chan chan bool
+	mutex     sync.RWMutex
+	role      string
+	onTaskRun func(name string)
 }
 
 type GenericTaskHandlerConfig struct {
+	// OnTaskRun, if set, is called with the task name whenever a scheduled task
+	// is about to run (after the cluster verify gate passes). Each scheduled
+	// task runs on its own goroutine, so it may be called concurrently across
+	// tasks; it must be non-blocking, panic-free, and safe for concurrent use.
+	OnTaskRun func(name string)
 }
 
 func NewGenericTaskHandler(cfg GenericTaskHandlerConfig) *GenericTaskHandler {
 	return &GenericTaskHandler{
-		tasks: make(map[string]Task),
-		role:  "Leader",
+		tasks:     make(map[string]Task),
+		role:      "Leader",
+		onTaskRun: cfg.OnTaskRun,
 	}
 }
 
@@ -108,6 +115,9 @@ func (h *GenericTaskHandler) runScheduled(t Task, schedule Schedule, ctx context
 			vCh := make(chan bool)
 			h.verify <- vCh
 			if ok := <-vCh; ok {
+				if h.onTaskRun != nil {
+					h.onTaskRun(t.Name())
+				}
 				slog.Debug(fmt.Sprintf("%s running task %s", h.role, t.Name()))
 				go t.Run(ctx, b)
 			}

pkg/rselection/taskhandler_test.go

@@ -155,3 +155,47 @@ func (s *TaskHandlerSuite) TestGenericTaskHandler(c *check.C) {
 	handler.Stop()
 	<-persistentTaskDone
 }
+
+func (s *TaskHandlerSuite) TestOnTaskRunFires(c *check.C) {
+	defer leaktest.Check(c)
+
+	var ran int
+	taskDone := make(chan bool)
+	defer close(taskDone)
+	tick := make(chan time.Time)
+	scheduledTask := &TestTask{
+		GenericTask: GenericTask{
+			TaskName:     "one",
+			TaskType:     TaskTypeScheduled,
+			TaskSchedule: &IntervalSchedule{Ticker: tick},
+		},
+		ran: &ran,
+		n:   taskDone,
+	}
+
+	names := make(chan string, 4)
+	handler := NewGenericTaskHandler(GenericTaskHandlerConfig{
+		OnTaskRun: func(name string) { names <- name },
+	})
+	handler.Register("one", scheduledTask)
+	handler.Handle(nil)
+
+	end := make(chan struct{})
+	defer close(end)
+	go func() {
+		for {
+			select {
+			case ch := <-handler.Verify():
+				ch <- true
+			case <-end:
+				return
+			}
+		}
+	}()
+
+	tick <- time.Now()
+	<-taskDone
+	c.Assert(<-names, check.Equals, "one")
+
+	handler.Stop()
+}