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

@@ -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.

hscontrol/servertest/lifecycle_test.go

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