mapper,policy: add reconnect-storm and lock-concurrency regression tests

TestInitialMapNotStarvedByReconnectStorm reproduces the #3346 stall;
TestPolicyManagerConcurrentReads guards the RLock cache access under -race.

Updates #3346

Author: kradalby

hscontrol/mapper/initialmap_storm_test.go

@@ -0,0 +1,187 @@
+//go:build !race
+
+// This is a timing-sensitive performance regression test; the race detector's
+// ~10x slowdown makes its wall-clock assertion meaningless, so it is excluded
+// from -race builds. The concurrency correctness of the policy lock change it
+// guards is covered under -race by TestPolicyManagerConcurrentReads in
+// hscontrol/policy/v2.
+
+package mapper
+
+import (
+	"net/netip"
+	"runtime"
+	"slices"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/juanfont/headscale/hscontrol/db"
+	"github.com/juanfont/headscale/hscontrol/derp"
+	"github.com/juanfont/headscale/hscontrol/state"
+	"github.com/juanfont/headscale/hscontrol/types"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"tailscale.com/tailcfg"
+)
+
+// setupStormBatcher builds a real state+batcher with production-default
+// NodeStore batching so the reconnect-storm contention is realistic. It mirrors
+// setupBatcherWithTestData but lets the test control BatcherWorkers and the
+// policy.
+func setupStormBatcher(tb testing.TB, nodeCount, workers int, policy string) (*TestData, func()) {
+	tb.Helper()
+
+	tmpDir := tb.TempDir()
+	prefixV4 := netip.MustParsePrefix("100.64.0.0/10")
+	prefixV6 := netip.MustParsePrefix("fd7a:115c:a1e0::/48")
+
+	cfg := &types.Config{
+		Database: types.DatabaseConfig{
+			Type:   types.DatabaseSqlite,
+			Sqlite: types.SqliteConfig{Path: tmpDir + "/headscale_test.db"},
+		},
+		PrefixV4:     &prefixV4,
+		PrefixV6:     &prefixV6,
+		IPAllocation: types.IPAllocationStrategySequential,
+		BaseDomain:   "headscale.test",
+		Policy:       types.PolicyConfig{Mode: types.PolicyModeDB},
+		DERP: types.DERPConfig{
+			ServerEnabled: false,
+			DERPMap: &tailcfg.DERPMap{
+				Regions: map[int]*tailcfg.DERPRegion{999: {RegionID: 999}},
+			},
+		},
+		Tuning: types.Tuning{
+			BatchChangeDelay: 10 * time.Millisecond,
+			BatcherWorkers:   workers,
+			// Production defaults: coalesce writes so the storm is not
+			// exaggerated by an unrealistically small NodeStore batch.
+			NodeStoreBatchSize:    100,
+			NodeStoreBatchTimeout: 500 * time.Millisecond,
+		},
+	}
+
+	database, err := db.NewHeadscaleDatabase(cfg)
+	require.NoError(tb, err)
+
+	users := database.CreateUsersForTest(1, "testuser")
+	dbNodes := database.CreateRegisteredNodesForTest(users[0], nodeCount, "node")
+
+	allNodes := make([]node, 0, nodeCount)
+	for i := range dbNodes {
+		allNodes = append(allNodes, node{
+			n:  dbNodes[i],
+			ch: make(chan *tailcfg.MapResponse, normalBufferSize),
+		})
+	}
+
+	st, err := state.NewState(cfg)
+	require.NoError(tb, err)
+
+	derpMap, err := derp.GetDERPMap(cfg.DERP)
+	require.NoError(tb, err)
+	st.SetDERPMap(derpMap)
+
+	_, err = st.SetPolicy([]byte(policy))
+	require.NoError(tb, err)
+
+	batcher := wrapBatcherForTest(NewBatcherAndMapper(cfg, st), st)
+	batcher.Start()
+
+	td := &TestData{
+		Database: database,
+		Users:    users,
+		Nodes:    allNodes,
+		State:    st,
+		Config:   cfg,
+		Batcher:  batcher,
+	}
+
+	return td, func() {
+		batcher.Close()
+		st.Close()
+		database.Close()
+	}
+}
+
+// TestInitialMapNotStarvedByReconnectStorm reproduces juanfont/headscale#3346.
+//
+// When every node redials at once (e.g. after a server upgrade restart), each
+// connection writes the NodeStore (UpdateNodeFromMapRequest + Connect) and the
+// batcher generates its initial map. All of that reads the policy through the
+// PolicyManager. Before the fix the PolicyManager guarded every read with a
+// single exclusive mutex, so the NodeStore writer's O(n^2) BuildPeerMap and
+// every node's FilterForNode serialised against each other. On a per-node
+// filter policy (autogroup:self, via, relay grants) each hold is expensive, so
+// under the storm time-to-initial-map grew without bound.
+//
+// On the production server in #3346 this drove the batcher's per-node
+// total.duration from ~4s to ~76s; tailscale clients aborted the map POST
+// first and reported
+//
+//	PollNetMap: Post ".../machine/map": unexpected EOF
+//
+// then redialled, feeding the storm so it never converged. An allow-all policy
+// does NOT reproduce this — BuildPeerMap is cheap there; the per-node filter
+// path is what makes it expensive, matching a real deployment's ACLs.
+//
+// The fix makes PolicyManager reads take a shared RLock so map generation runs
+// concurrently. AddNode blocks until the initial map is generated and handed to
+// the node channel, so its wall-clock duration is the time-to-initial-map the
+// client experiences. Without the fix this test's slowest node takes ~10s+ at
+// this scale (lock-bound, and more workers do not help); with it, generation
+// parallelises across workers and stays well within a client's patience.
+func TestInitialMapNotStarvedByReconnectStorm(t *testing.T) {
+	if testing.Short() {
+		t.Skip("timing-sensitive storm regression; skipped in -short")
+	}
+
+	const (
+		nodeCount = 300
+
+		// A per-node-filter policy: forces BuildPeerMap and FilterForNode onto
+		// the slow path that recompiles filter rules per node, the same shape
+		// as a real ACL using autogroup:self / via / relay grants.
+		perNodeFilterPolicy = `{"acls":[{"action":"accept","src":["autogroup:member"],"dst":["autogroup:self:*"]}]}`
+
+		// Generous bound: post-fix the slowest node lands at ~2-5s depending on
+		// core count; pre-fix it is lock-bound at ~10s+ regardless of workers.
+		maxAcceptableLatency = 8 * time.Second
+	)
+
+	// Use the real available parallelism, as production does; the fix's win is
+	// that generation scales with it instead of serialising on the policy lock.
+	workers := runtime.NumCPU()
+
+	td, cleanup := setupStormBatcher(t, nodeCount, workers, perNodeFilterPolicy)
+	defer cleanup()
+
+	latencies := make([]time.Duration, nodeCount)
+
+	var wg sync.WaitGroup
+
+	for i := range td.Nodes {
+		wg.Go(func() {
+			n := &td.Nodes[i]
+
+			start := time.Now()
+			err := td.Batcher.AddNode(n.n.ID, n.ch, tailcfg.CapabilityVersion(100), nil)
+			latencies[i] = time.Since(start)
+
+			assert.NoError(t, err) //nolint:testifylint // assert (not require) is correct off the test goroutine
+		})
+	}
+
+	wg.Wait()
+
+	slices.Sort(latencies)
+	p50 := latencies[len(latencies)/2]
+	p95 := latencies[len(latencies)*95/100]
+	maxLatency := latencies[len(latencies)-1]
+	t.Logf("initial-map latency over %d nodes (workers=%d): p50=%s p95=%s max=%s",
+		nodeCount, workers, p50, p95, maxLatency)
+
+	require.Less(t, maxLatency, maxAcceptableLatency,
+		"slowest initial map took %s: policy reads are serialising instead of running concurrently (issue #3346)", maxLatency)
+}

hscontrol/policy/v2/policy_concurrency_test.go

@@ -0,0 +1,103 @@
+package v2
+
+import (
+	"fmt"
+	"sync"
+	"testing"
+
+	"github.com/juanfont/headscale/hscontrol/types"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
+)
+
+// TestPolicyManagerConcurrentReads is the correctness guard for the #3346 fix:
+// PolicyManager read methods take a shared RLock and populate their per-node
+// caches (filterRulesMap, matchersForNodeMap) concurrently. This test hammers
+// those reads from many goroutines while a writer mutates the node set, so the
+// race detector catches any unsafe access to the shared caches or policy state.
+//
+// It uses an autogroup:self policy so reads take the per-node filter slow path
+// — the same path that made #3346's reconnect storm expensive — which is where
+// the lazy caches are written.
+func TestPolicyManagerConcurrentReads(t *testing.T) {
+	users := types.Users{
+		{Model: gorm.Model{ID: 1}, Name: "user1", Email: "user1@headscale.net"},
+		{Model: gorm.Model{ID: 2}, Name: "user2", Email: "user2@headscale.net"},
+		{Model: gorm.Model{ID: 3}, Name: "user3", Email: "user3@headscale.net"},
+	}
+
+	policy := `{
+		"acls": [
+			{
+				"action": "accept",
+				"src": ["autogroup:member"],
+				"dst": ["autogroup:self:*"]
+			}
+		]
+	}`
+
+	const nodeCount = 60
+
+	nodes := make(types.Nodes, 0, nodeCount)
+	for i := range nodeCount {
+		n := node(
+			fmt.Sprintf("node%d", i),
+			fmt.Sprintf("100.64.0.%d", i+1),
+			fmt.Sprintf("fd7a:115c:a1e0::%d", i+1),
+			users[i%len(users)],
+		)
+		n.ID = types.NodeID(i + 1) //nolint:gosec // safe in test
+		nodes = append(nodes, n)
+	}
+
+	pm, err := NewPolicyManager([]byte(policy), users, nodes.ViewSlice())
+	require.NoError(t, err)
+
+	const (
+		readers        = 16
+		iterations     = 60
+		mutatorReloads = 30
+	)
+
+	var wg sync.WaitGroup
+
+	// Concurrent readers exercise every converted RLock read path, including
+	// the two lazily populated per-node caches. Assertions inside the
+	// goroutines use assert (not require) so a failure does not call
+	// t.FailNow from a non-test goroutine.
+	for r := range readers {
+		wg.Go(func() {
+			for i := range iterations {
+				nv := nodes[(r+i)%len(nodes)].View()
+
+				rules, err := pm.FilterForNode(nv)
+				assert.NoError(t, err) //nolint:testifylint // assert (not require) is correct off the test goroutine
+				assert.NotNil(t, rules)
+
+				_, err = pm.MatchersForNode(nv)
+				assert.NoError(t, err) //nolint:testifylint // assert (not require) is correct off the test goroutine
+
+				pm.Filter()
+				pm.NodeCapMap(nv.ID())
+
+				// BuildPeerMap is the O(n^2) writer-side read; exercise it
+				// under RLock too, but not every iteration.
+				if i%8 == 0 {
+					assert.NotNil(t, pm.BuildPeerMap(nodes.ViewSlice()))
+				}
+			}
+		})
+	}
+
+	// A writer repeatedly re-sets the node set, invalidating and racing the
+	// caches the readers are populating.
+	wg.Go(func() {
+		for range mutatorReloads {
+			_, err := pm.SetNodes(nodes.ViewSlice())
+			assert.NoError(t, err) //nolint:testifylint // assert (not require) is correct off the test goroutine
+		}
+	})
+
+	wg.Wait()
+}