state: keep HA health probes active without delaying unrelated routers

Probe lone unhealthy advertisers. Defer results only for routers whose
routes tie them to a router still awaiting its probe.

hscontrol/servertest/ha_health_test.go

@@ -392,3 +392,101 @@ func TestHAHealthProbe_NoHARoutes(t *testing.T) {
 	})
 	assert.False(t, dispatched, "no HA routes should produce no changes")
 }
+
+func TestHAHealthProbe_DeferralScopedToPrefix(t *testing.T) {
+	t.Parallel()
+
+	srv := servertest.NewServer(t)
+	user := srv.CreateUser(t, "ha-scoped")
+	officeRoute := netip.MustParsePrefix("10.103.0.0/24")
+	labRoute := netip.MustParsePrefix("10.104.0.0/24")
+
+	c1 := servertest.NewClient(t, srv, "ha-sc-office1", servertest.WithUser(user))
+	c2 := servertest.NewClient(t, srv, "ha-sc-office2", servertest.WithUser(user))
+	c3 := servertest.NewClient(t, srv, "ha-sc-lab1", servertest.WithUser(user))
+	c4 := servertest.NewClient(t, srv, "ha-sc-lab2", servertest.WithUser(user))
+
+	for _, c := range []*servertest.TestClient{c1, c2, c3, c4} {
+		c.WaitForPeers(t, 3, 10*time.Second)
+	}
+
+	officeID1 := advertiseAndApproveRoute(t, srv, c1, officeRoute)
+	advertiseAndApproveRoute(t, srv, c2, officeRoute)
+	labID1 := advertiseAndApproveRoute(t, srv, c3, labRoute)
+	labID2 := advertiseAndApproveRoute(t, srv, c4, labRoute)
+
+	require.True(t, srv.State().SetNodeHealth(labID1, false))
+	require.Contains(t, srv.State().GetNodePrimaryRoutes(labID2), labRoute)
+
+	prober := state.NewHAHealthProber(
+		srv.State(),
+		types.HARouteConfig{
+			ProbeInterval: 30 * time.Second,
+			ProbeTimeout:  2 * time.Second,
+		},
+		srv.URL,
+		srv.App.MapBatcher().IsConnected,
+	)
+
+	ctx, cancel := context.WithTimeout(t.Context(), 15*time.Second)
+	defer cancel()
+
+	// Only the fresh office session misses its probe; lab recovery can proceed.
+	prober.ProbeOnce(ctx, func(changes ...change.Change) {
+		for _, c := range changes {
+			if c.PingRequest != nil && c.TargetNode == officeID1 {
+				continue
+			}
+
+			srv.App.Change(c)
+		}
+	})
+
+	assert.True(t, srv.State().IsNodeHealthy(officeID1))
+	assert.True(t, srv.State().IsNodeHealthy(labID1),
+		"a deferral on another prefix must not withhold the lab node's recovery")
+	assert.Contains(t, srv.State().GetNodePrimaryRoutes(labID2), labRoute)
+}
+
+func TestHAHealthProbe_LastAdvertiserRecovers(t *testing.T) {
+	t.Parallel()
+
+	srv := servertest.NewServer(t)
+	user := srv.CreateUser(t, "ha-last")
+	route := netip.MustParsePrefix("10.107.0.0/24")
+	c1 := servertest.NewClient(t, srv, "ha-la-r1", servertest.WithUser(user))
+	c2 := servertest.NewClient(t, srv, "ha-la-r2", servertest.WithUser(user))
+
+	c1.WaitForPeers(t, 1, 10*time.Second)
+	c2.WaitForPeers(t, 1, 10*time.Second)
+
+	nodeID1 := advertiseAndApproveRoute(t, srv, c1, route)
+	nodeID2 := advertiseAndApproveRoute(t, srv, c2, route)
+
+	srv.State().SetNodeHealth(nodeID1, false)
+	require.Contains(t, srv.State().GetNodePrimaryRoutes(nodeID2), route)
+
+	// The remaining advertiser needs a probe before its route can return.
+	_, rc, err := srv.State().SetApprovedRoutes(nodeID2, nil)
+	require.NoError(t, err)
+	srv.App.Change(rc)
+	require.Empty(t, srv.State().GetNodePrimaryRoutes(nodeID1))
+
+	prober := state.NewHAHealthProber(
+		srv.State(),
+		types.HARouteConfig{
+			ProbeInterval: 30 * time.Second,
+			ProbeTimeout:  5 * time.Second,
+		},
+		srv.URL,
+		srv.App.MapBatcher().IsConnected,
+	)
+
+	ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
+	defer cancel()
+
+	prober.ProbeOnce(ctx, srv.App.Change)
+
+	assert.True(t, srv.State().IsNodeHealthy(nodeID1))
+	assert.Contains(t, srv.State().GetNodePrimaryRoutes(nodeID1), route)
+}

hscontrol/state/ha_health.go

@@ -2,8 +2,9 @@ package state
 
 import (
 	"context"
+	"net/netip"
+	"slices"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/juanfont/headscale/hscontrol/types"
@@ -60,14 +61,10 @@ func (p *HAHealthProber) forgetSession(id types.NodeID) {
 	p.lastStableSession.Delete(id)
 }
 
-// ProbeOnce pings every HA subnet router and applies the cycle's
-// results in one batch so the election sees a single transition.
-// Per-result snapshots could otherwise elect a node that the next
-// snapshot demotes again, flipping primary onto an unreachable peer.
-// A timeout that fires after the node reconnected, or against a
-// session younger than one probe cycle, is dropped so wgengine has
-// time to apply the new netmap before silence is read as
-// unreachability.
+// ProbeOnce applies HA probe results in one batch to avoid transient elections.
+// Fresh or reconnected sessions defer their whole group of routers linked
+// through shared prefixes, giving wgengine time to apply the new netmap.
+// Other groups can still apply their complete results.
 func (p *HAHealthProber) ProbeOnce(
 	ctx context.Context,
 	dispatch func(...change.Change),
@@ -111,7 +108,7 @@ func (p *HAHealthProber) ProbeOnce(
 	var (
 		wg       sync.WaitGroup
 		results  = xsync.NewMap[types.NodeID, bool]()
-		deferred atomic.Bool
+		deferred = xsync.NewMap[types.NodeID, struct{}]()
 	)
 
 	for _, id := range nodeIDs {
@@ -176,7 +173,7 @@ func (p *HAHealthProber) ProbeOnce(
 						Uint64("current_session", curr.SessionEpoch()).
 						Msg("HA probe: node reconnected during probe, skipping")
 
-					deferred.Store(true)
+					deferred.Store(id, struct{}{})
 
 					return
 				}
@@ -187,7 +184,7 @@ func (p *HAHealthProber) ProbeOnce(
 						Uint64("probe_session", probeSession).
 						Msg("HA probe: probe of fresh session timed out, deferring to next cycle")
 
-					deferred.Store(true)
+					deferred.Store(id, struct{}{})
 
 					return
 				}
@@ -207,18 +204,14 @@ func (p *HAHealthProber) ProbeOnce(
 
 	wg.Wait()
 
-	// When any probe in the cycle was deferred (fresh session or
-	// reconnected mid-probe), drop the whole cycle's results: a partial
-	// batch lets the election pick a node whose connectivity is still
-	// unknown. The next cycle will run with stable sessions for every
-	// candidate and can decide on a complete picture.
-	if deferred.Load() {
-		return
-	}
+	// Withhold incomplete HA groups; independent groups can apply their results.
+	undecided := undecidedHANodes(haNodes, deferred)
 
 	healthByNode := make(map[types.NodeID]bool, results.Size())
 	results.Range(func(id types.NodeID, healthy bool) bool {
-		healthByNode[id] = healthy
+		if !undecided.Contains(id) {
+			healthByNode[id] = healthy
+		}
 
 		return true
 	})
@@ -231,3 +224,38 @@ func (p *HAHealthProber) ProbeOnce(
 			Msg("HA probe: health changed, triggering failover/recovery")
 	}
 }
+
+// undecidedHANodes follows shared prefixes because one node's health affects
+// every prefix it advertises. Each remaining group must have a complete batch.
+func undecidedHANodes(
+	haNodes map[netip.Prefix][]types.NodeID,
+	deferred *xsync.Map[types.NodeID, struct{}],
+) set.Set[types.NodeID] {
+	undecided := make(set.Set[types.NodeID])
+
+	deferred.Range(func(id types.NodeID, _ struct{}) bool {
+		undecided.Add(id)
+
+		return true
+	})
+
+	for grown := true; grown; {
+		grown = false
+
+		for _, ids := range haNodes {
+			if !slices.ContainsFunc(ids, undecided.Contains) {
+				continue
+			}
+
+			for _, id := range ids {
+				if !undecided.Contains(id) {
+					undecided.Add(id)
+
+					grown = true
+				}
+			}
+		}
+	}
+
+	return undecided
+}

hscontrol/state/node_store.go

@@ -1046,8 +1046,8 @@ func (s *NodeStore) PrimaryRoutesForNode(id types.NodeID) []netip.Prefix {
 	return out
 }
 
-// HANodes returns the prefixes with two or more online advertisers, the
-// candidate set the HA prober needs to monitor.
+// HANodes returns prefixes with multiple online advertisers and lone unhealthy
+// advertisers that need another probe before their route can return.
 func (s *NodeStore) HANodes() map[netip.Prefix][]types.NodeID {
 	snap := s.data.Load()
 
@@ -1056,10 +1056,14 @@ func (s *NodeStore) HANodes() map[netip.Prefix][]types.NodeID {
 		slices.Sorted(maps.Keys(snap.nodesByID)),
 	)
 
+	unhealthy := func(id types.NodeID) bool {
+		return snap.nodesByID[id].Unhealthy
+	}
+
 	out := make(map[netip.Prefix][]types.NodeID)
 
 	for p, ids := range advertisers {
-		if len(ids) < 2 {
+		if len(ids) < 2 && !slices.ContainsFunc(ids, unhealthy) {
 			continue
 		}
 

hscontrol/state/primaries_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"github.com/juanfont/headscale/hscontrol/types"
+	"github.com/puzpuzpuz/xsync/v4"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"tailscale.com/tailcfg"
@@ -281,6 +282,29 @@ func TestPrimaries_HANodes(t *testing.T) {
 			},
 			want: map[netip.Prefix][]types.NodeID{},
 		},
+		{
+			name: "unhealthy-node-left-alone-is-still-monitored",
+			setup: func(f *primariesFixture) {
+				f.advertise(1, mp("192.168.1.0/24"))
+				f.advertise(2, mp("192.168.1.0/24"))
+				f.unhealthy(1)
+				f.approveRoutes(2)
+			},
+			want: map[netip.Prefix][]types.NodeID{
+				mp("192.168.1.0/24"): {1},
+			},
+		},
+		{
+			name: "node-left-alone-not-ha-once-healthy",
+			setup: func(f *primariesFixture) {
+				f.advertise(1, mp("192.168.1.0/24"))
+				f.advertise(2, mp("192.168.1.0/24"))
+				f.unhealthy(1)
+				f.approveRoutes(2)
+				f.healthy(1)
+			},
+			want: map[netip.Prefix][]types.NodeID{},
+		},
 		{
 			name: "two-nodes-same-prefix-is-ha",
 			setup: func(f *primariesFixture) {
@@ -339,3 +363,55 @@ func TestPrimaries_HANodes(t *testing.T) {
 		})
 	}
 }
+
+func TestUndecidedHANodes(t *testing.T) {
+	office := mp("10.70.0.0/24")
+	lab := mp("10.80.0.0/24")
+	shared := mp("10.90.0.0/24")
+
+	tests := []struct {
+		name     string
+		haNodes  map[netip.Prefix][]types.NodeID
+		deferred []types.NodeID
+		want     []types.NodeID
+	}{
+		{
+			name: "nothing deferred",
+			haNodes: map[netip.Prefix][]types.NodeID{
+				office: {1, 2},
+				lab:    {3, 4},
+			},
+		},
+		{
+			name: "disjoint prefixes",
+			haNodes: map[netip.Prefix][]types.NodeID{
+				office: {1, 2},
+				lab:    {3, 4},
+			},
+			deferred: []types.NodeID{3},
+			want:     []types.NodeID{3, 4},
+		},
+		{
+			name: "shared router carries the deferral across",
+			haNodes: map[netip.Prefix][]types.NodeID{
+				office: {1, 2},
+				shared: {2, 5},
+				lab:    {3, 4},
+			},
+			deferred: []types.NodeID{5},
+			want:     []types.NodeID{1, 2, 5},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			deferred := xsync.NewMap[types.NodeID, struct{}]()
+			for _, id := range tt.deferred {
+				deferred.Store(id, struct{}{})
+			}
+
+			got := undecidedHANodes(tt.haNodes, deferred)
+			assert.ElementsMatch(t, tt.want, got.Slice())
+		})
+	}
+}