From 8b48f3847d91d9a309b9593dcd17d7fe6aae1291 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Mon, 15 Sep 2025 15:49:56 -0700
Subject: [PATCH] net/netmon, wgengine/magicsock: simplify LinkChangeLogLimiter
 signature

Remove the need for the caller to hold on to and call an unregister
function. Both two callers (one real, one test) already have a context
they can use. Use context.AfterFunc instead. There are no observable
side effects from scheduling too late if the goroutine doesn't run sync.

Updates #17148

Change-Id: Ie697dae0e797494fa8ef27fbafa193bfe5ceb307
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 net/netmon/loghelper.go         | 12 +++++++-----
 net/netmon/loghelper_test.go    | 19 ++++++++++++++-----
 wgengine/magicsock/magicsock.go | 12 ++++--------
 3 files changed, 25 insertions(+), 18 deletions(-)

diff --git a/net/netmon/loghelper.go b/net/netmon/loghelper.go
index 824faeef0..96991644c 100644
--- a/net/netmon/loghelper.go
+++ b/net/netmon/loghelper.go
@@ -4,6 +4,7 @@
 package netmon
 
 import (
+	"context"
 	"sync"
 
 	"tailscale.com/types/logger"
@@ -12,16 +13,17 @@ import (
 // LinkChangeLogLimiter returns a new [logger.Logf] that logs each unique
 // format string to the underlying logger only once per major LinkChange event.
 //
-// The returned function should be called when the logger is no longer needed,
-// to release resources from the Monitor.
-func LinkChangeLogLimiter(logf logger.Logf, nm *Monitor) (_ logger.Logf, unregister func()) {
+// The logger stops tracking seen format strings when the provided context is
+// done.
+func LinkChangeLogLimiter(ctx context.Context, logf logger.Logf, nm *Monitor) logger.Logf {
 	var formatSeen sync.Map // map[string]bool
-	unregister = nm.RegisterChangeCallback(func(cd *ChangeDelta) {
+	unregister := nm.RegisterChangeCallback(func(cd *ChangeDelta) {
 		// If we're in a major change or a time jump, clear the seen map.
 		if cd.Major || cd.TimeJumped {
 			formatSeen.Clear()
 		}
 	})
+	context.AfterFunc(ctx, unregister)
 
 	return func(format string, args ...any) {
 		// We only store 'true' in the map, so if it's present then it
@@ -38,5 +40,5 @@ func LinkChangeLogLimiter(logf logger.Logf, nm *Monitor) (_ logger.Logf, unregis
 		}
 
 		logf(format, args...)
-	}, unregister
+	}
 }
diff --git a/net/netmon/loghelper_test.go b/net/netmon/loghelper_test.go
index 44aa46783..aeac9f031 100644
--- a/net/netmon/loghelper_test.go
+++ b/net/netmon/loghelper_test.go
@@ -5,13 +5,17 @@ package netmon
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"testing"
+	"testing/synctest"
 
 	"tailscale.com/util/eventbus"
 )
 
-func TestLinkChangeLogLimiter(t *testing.T) {
+func TestLinkChangeLogLimiter(t *testing.T) { synctest.Test(t, syncTestLinkChangeLogLimiter) }
+
+func syncTestLinkChangeLogLimiter(t *testing.T) {
 	bus := eventbus.New()
 	defer bus.Close()
 	mon, err := New(bus, t.Logf)
@@ -30,8 +34,10 @@ func TestLinkChangeLogLimiter(t *testing.T) {
 		fmt.Fprintf(&logBuffer, format, args...)
 	}
 
-	logf, unregister := LinkChangeLogLimiter(logf, mon)
-	defer unregister()
+	ctx, cancel := context.WithCancel(t.Context())
+	defer cancel()
+
+	logf = LinkChangeLogLimiter(ctx, logf, mon)
 
 	// Log once, which should write to our log buffer.
 	logf("hello %s", "world")
@@ -72,8 +78,11 @@ func TestLinkChangeLogLimiter(t *testing.T) {
 		t.Errorf("unexpected log buffer contents: %q", got)
 	}
 
-	// Unregistering the callback should clear our 'cbs' set.
-	unregister()
+	// Canceling the context we passed to LinkChangeLogLimiter should
+	// unregister the callback from the netmon.
+	cancel()
+	synctest.Wait()
+
 	mon.mu.Lock()
 	if len(mon.cbs) != 0 {
 		t.Errorf("expected no callbacks, got %v", mon.cbs)
diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index fa1f1f88f..36402122c 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -209,10 +209,6 @@ type Conn struct {
 	// port mappings from NAT devices.
 	portMapper *portmapper.Client
 
-	// portMapperLogfUnregister is the function to call to unregister
-	// the portmapper log limiter.
-	portMapperLogfUnregister func()
-
 	// derpRecvCh is used by receiveDERP to read DERP messages.
 	// It must have buffer size > 0; see issue 3736.
 	derpRecvCh chan derpReadResult
@@ -748,10 +744,13 @@ func NewConn(opts Options) (*Conn, error) {
 	c.subsDoneCh = make(chan struct{})
 	go c.consumeEventbusTopics()
 
+	c.connCtx, c.connCtxCancel = context.WithCancel(context.Background())
+	c.donec = c.connCtx.Done()
+
 	// Don't log the same log messages possibly every few seconds in our
 	// portmapper.
 	portmapperLogf := logger.WithPrefix(c.logf, "portmapper: ")
-	portmapperLogf, c.portMapperLogfUnregister = netmon.LinkChangeLogLimiter(portmapperLogf, opts.NetMon)
+	portmapperLogf = netmon.LinkChangeLogLimiter(c.connCtx, portmapperLogf, opts.NetMon)
 	portMapOpts := &portmapper.DebugKnobs{
 		DisableAll: func() bool { return opts.DisablePortMapper || c.onlyTCP443.Load() },
 	}
@@ -772,8 +771,6 @@ func NewConn(opts Options) (*Conn, error) {
 		return nil, err
 	}
 
-	c.connCtx, c.connCtxCancel = context.WithCancel(context.Background())
-	c.donec = c.connCtx.Done()
 	c.netChecker = &netcheck.Client{
 		Logf:                logger.WithPrefix(c.logf, "netcheck: "),
 		NetMon:              c.netMon,
@@ -3330,7 +3327,6 @@ func (c *Conn) Close() error {
 	}
 	c.stopPeriodicReSTUNTimerLocked()
 	c.portMapper.Close()
-	c.portMapperLogfUnregister()
 
 	c.peerMap.forEachEndpoint(func(ep *endpoint) {
 		ep.stopAndReset()