From 2b7b2475fd38c3e00fd0389973bb69fd4e5d6a49 Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Tue, 20 Oct 2026 20:14:36 +0000 Subject: [PATCH] servertest: count map frames per client Counts ControlTime events, so frames that leave the netmap as it was, or that a delta client applies in place, still show. --- hscontrol/servertest/client.go | 34 +++++++++++++ hscontrol/servertest/lifecycle_test.go | 67 ++++++++++++++++++++++++++ 2 files changed, 101 insertions(+) diff --git a/hscontrol/servertest/client.go b/hscontrol/servertest/client.go index 1b66231c9..1f45f9f1a 100644 --- a/hscontrol/servertest/client.go +++ b/hscontrol/servertest/client.go @@ -5,6 +5,7 @@ import ( "fmt" "path" "slices" + "strings" "sync" "testing" "time" @@ -76,6 +77,11 @@ type TestClient struct { // registers: see [TestClient.RegisterAttempts]. registers []testcapture.RegisterAttempt + // frames: see [TestClient.Frames]. stamped is set while the response + // in hand carried a ControlTime, so its rebuild is already counted. + frames int + stamped bool + // updates is a buffered channel that receives a signal // each time a new [netmap.NetworkMap] arrives. updates chan *netmap.NetworkMap @@ -260,6 +266,12 @@ func TryNewClient(tb testing.TB, server *TestServer, name string, opts ...Client tc.mu.Lock() tc.registers, _ = testcapture.AppendRegisterLine(tc.registers, line) + + // The poll logs a response's ControlTime before it applies it. + if strings.Contains(line, `{"controltime":`) { + tc.frames++ + tc.stamped = true + } tc.mu.Unlock() } @@ -405,6 +417,7 @@ func (c *TestClient) startPollLoop() { c.mu.Lock() c.streams++ c.streamMaps = 0 + c.stamped = false c.mu.Unlock() go func() { @@ -447,6 +460,16 @@ func (c *TestClient) UpdateFullNetmap(nm *netmap.NetworkMap) { if !c.noHistory { c.history = append(c.history, nm) } + + // Removals sent ahead of their batch and a deleted node's expiry carry + // no ControlTime, so a rebuild with none before it is a response too. + // shortcut: one right after a response a delta client applied in place + // goes uncounted; count at the decoder if headscale sends one alone. + if !c.stamped { + c.frames++ + } + + c.stamped = false c.mu.Unlock() // Non-blocking send to the updates channel. @@ -919,6 +942,17 @@ func (c *TestClient) UpdateCount() int { return c.updatesReceived } +// Frames returns how many map responses other than keep-alives the client +// has been sent: each one with a ControlTime, and each netmap rebuild from +// one without. Unlike [TestClient.UpdateCount] it also counts responses that +// a [WithDeltaUpdates] client applies without rebuilding its netmap. +func (c *TestClient) Frames() int { + c.mu.RLock() + defer c.mu.RUnlock() + + return c.frames +} + // Stream returns how many map streams the client has opened, then how many // netmaps the current one delivered. [TestClient.Netmap] outlives a stream, // so only this tells a new stream that was answered from one that was not. diff --git a/hscontrol/servertest/lifecycle_test.go b/hscontrol/servertest/lifecycle_test.go index 9168635c5..2ffce9f8e 100644 --- a/hscontrol/servertest/lifecycle_test.go +++ b/hscontrol/servertest/lifecycle_test.go @@ -549,6 +549,73 @@ func TestPathsDiscoAndExpiryHelpers(t *testing.T) { "the poll goes on while the login waits") } +// TestFramesCountResponsesWithoutControlTime checks that Frames counts the +// responses headscale sends without a ControlTime: a removal ahead of its +// batch, and the expiry a deleted node's new stream opens with. A client +// that rebuilds its netmap on every response counts one frame per rebuild. +func TestFramesCountResponsesWithoutControlTime(t *testing.T) { + t.Parallel() + + requireFramePerRebuild := func(t *testing.T, cl *servertest.TestClient) { + t.Helper() + + assert.EventuallyWithT(t, func(c *assert.CollectT) { + assert.Equal(c, cl.UpdateCount(), cl.Frames()) + }, 10*time.Second, 50*time.Millisecond) + } + + t.Run("removal", func(t *testing.T) { + t.Parallel() + + srv := servertest.NewServer(t) + user1 := srv.CreateUser(t, "frames-user1") + user2 := srv.CreateUser(t, "frames-user2") + viewer := servertest.NewClient(t, srv, "frames-viewer", servertest.WithUser(user1)) + servertest.NewClient(t, srv, "frames-peer", servertest.WithUser(user2)) + viewer.WaitForPeers(t, 1, 15*time.Second) + + changed, err := srv.State().SetPolicy([]byte(`{ + "acls": [ + {"action": "accept", "src": ["frames-user1@"], "dst": ["frames-user1@:*"]}, + {"action": "accept", "src": ["frames-user2@"], "dst": ["frames-user2@:*"]} + ] + }`)) + require.NoError(t, err) + require.True(t, changed) + + changes, err := srv.State().ReloadPolicy() + require.NoError(t, err) + srv.App.Change(changes...) + + viewer.WaitForPeerCount(t, 0, 15*time.Second) + requireFramePerRebuild(t, viewer) + }) + + t.Run("deleted-node-expiry", func(t *testing.T) { + t.Parallel() + + srv := servertest.NewServer(t) + cl := servertest.NewClient(t, srv, "frames-deleted") + cl.WaitForCondition(t, "initial netmap", 15*time.Second, + func(nm *netmap.NetworkMap) bool { return nm != nil }) + + node, ok := srv.State().GetNodeByID(findNodeID(t, srv, "frames-deleted")) + require.True(t, ok) + + changes, err := srv.State().DeleteNode(node) + require.NoError(t, err) + srv.App.Change(changes...) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + require.NoError(t, cl.RestartPoll(ctx)) + cl.WaitForCondition(t, "expired self", 15*time.Second, + func(nm *netmap.NetworkMap) bool { return nm != nil && nm.SelfNode.Expired() }) + requireFramePerRebuild(t, cl) + }) +} + func TestNodeExpiryRouteFailover(t *testing.T) { t.Parallel()