state: withhold HA probe results only for routers tied to a deferred one

hscontrol/servertest/ha_health_test.go

@@ -393,6 +393,61 @@ 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()
 

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/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"
@@ -362,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())
+		})
+	}
+}