Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 54 additions & 28 deletions pkg/rselection/impls/pgx/leader.go
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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{},
}
}

Expand Down Expand Up @@ -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)
}
Comment on lines +295 to +297

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this also check if err is not nil to avoid unnecessary reporting of errors that are actually nil?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good question — this is intentional. OnIntegrityResult is a result hook, not a failure hook: it reports the outcome of every integrity check, where nil means healthy (as noted in its doc comment). That mirrors OnPingResult, which likewise fires on every ping with success bool (both true and false), so the two leader callbacks stay consistent.

Firing on nil also lets a consumer observe healthy checks, not just failures — e.g. recording a "last successful integrity check" timestamp or resetting state on recovery. Guarding with err != nil here would turn it into a failure-only hook and foreclose that.

The current consumer (PPM) treats the healthy case as a no-op (if err != nil { ClusterIntegrityCheckFailed() }), so a nil result never increments anything, and the per-sweep nil call is negligible. Happy to revisit if you'd prefer a narrower failure-only hook, but I'd lean toward keeping it symmetric with OnPingResult.

if err == nil {
if !p.unhealthySince.IsZero() {
slog.Info("Cluster integrity recovered", "degradedFor", p.timeNow().Sub(p.unhealthySince))
Expand Down Expand Up @@ -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,
Expand All @@ -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) {
Expand Down
53 changes: 53 additions & 0 deletions pkg/rselection/impls/pgx/leader_recovery_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package pgxelection
// Copyright (C) 2026 by Posit Software, PBC

import (
"context"
"errors"
"time"

Expand Down Expand Up @@ -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})
}
24 changes: 17 additions & 7 deletions pkg/rselection/taskhandler.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
}

Expand Down Expand Up @@ -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)
}
Expand Down
44 changes: 44 additions & 0 deletions pkg/rselection/taskhandler_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}
Loading