ipn/ipnlocal/netmapcache: add UpdateSelfOnly method (#19818)

Some netmap updates are guaranteed to affect only the "static" parts of the
netmap, and so should not require us to walk through all the peers and user
profiles when updating the cache. To support this, the new UpdateSelfOnly
method updates only the Self node and other tailnet settings that are not
dependent on the peers and profiles.

Use this when updating the cache on DERP home changes.

Updates #12542

Change-Id: Ifed522b29d579fb76e010b4ff738cc4e0a72d27f
Signed-off-by: M. J. Fromberger <fromberger@tailscale.com>

ipn/ipnlocal/diskcache.go

@@ -21,21 +21,29 @@ type diskCache struct {
 	cache *netmapcache.Cache
 }
 
-func (b *LocalBackend) writeNetmapToDiskLocked(nm *netmap.NetworkMap) error {
+// writeNetmapToDiskLockedWithoutPeers updates nm in the cache, excluding peers and profiles.
+func (b *LocalBackend) writeNetmapToDiskLockedWithoutPeers(nm *netmap.NetworkMap) error {
 	if !buildfeatures.HasCacheNetMap || nm == nil || nm.Cached {
 		return nil
-	}
-	b.logf("writing netmap to disk cache")
-
-	dir, err := b.profileMkdirAllLocked(b.pm.CurrentProfile().ID(), "netmap-cache")
-	if err != nil {
+	} else if err := b.ensureDiskCacheLocked(); err != nil {
 		return err
 	}
-	if c := b.diskCache; c.cache == nil || c.dir != dir {
-		b.diskCache.cache = netmapcache.NewCache(netmapcache.FileStore(dir))
-		b.diskCache.dir = dir
-	}
+	b.logf("updating netmap in disk cache")
+	return b.diskCache.cache.UpdateSelfOnly(b.currentNode().Context(), b.patchNetmapHomeDERPLocked(nm))
+}
 
+// writeNetmapToDiskLockedWithPeers writes nm into the cache, including peers and profiles.
+func (b *LocalBackend) writeNetmapToDiskLockedWithPeers(nm *netmap.NetworkMap) error {
+	if !buildfeatures.HasCacheNetMap || nm == nil || nm.Cached {
+		return nil
+	} else if err := b.ensureDiskCacheLocked(); err != nil {
+		return err
+	}
+	b.logf("writing netmap to disk cache")
+	return b.diskCache.cache.Store(b.currentNode().Context(), b.patchNetmapHomeDERPLocked(nm))
+}
+
+func (b *LocalBackend) patchNetmapHomeDERPLocked(nm *netmap.NetworkMap) *netmap.NetworkMap {
 	// Set the homeDERP on the self node before saving. The self node homeDERP is
 	// generally not used since the homeDERP for self is stored in magicsock, but
 	// to be able to load it during loading the cache, we use the existing field
@@ -46,8 +54,19 @@ func (b *LocalBackend) writeNetmapToDiskLocked(nm *netmap.NetworkMap) error {
 	selfNode := nm.SelfNode.AsStruct()
 	selfNode.HomeDERP = int(b.currentNode().homeDERP.Load())
 	nmCopy.SelfNode = selfNode.View()
+	return &nmCopy
+}
 
-	return b.diskCache.cache.Store(b.currentNode().Context(), &nmCopy)
+func (b *LocalBackend) ensureDiskCacheLocked() error {
+	dir, err := b.profileMkdirAllLocked(b.pm.CurrentProfile().ID(), "netmap-cache")
+	if err != nil {
+		return err
+	}
+	if c := b.diskCache; c.cache == nil || c.dir != dir {
+		b.diskCache.cache = netmapcache.NewCache(netmapcache.FileStore(dir))
+		b.diskCache.dir = dir
+	}
+	return nil
 }
 
 func (b *LocalBackend) loadDiskCacheLocked() (om *netmap.NetworkMap, ok bool) {

ipn/ipnlocal/diskcache_test.go

@@ -61,7 +61,7 @@ func TestWriteAndLoadHomeDERP(t *testing.T) {
 	b.mu.Lock()
 	defer b.mu.Unlock()
 
-	if err := b.writeNetmapToDiskLocked(nm); err != nil {
+	if err := b.writeNetmapToDiskLockedWithPeers(nm); err != nil {
 		t.Fatalf("writeNetmapToDiskLocked: %v", err)
 	}
 
@@ -127,7 +127,7 @@ func TestOnHomeDERPUpdate(t *testing.T) {
 
 		// Write an initial cache entry so we can verify it is not overwritten.
 		b.mu.Lock()
-		if err := b.writeNetmapToDiskLocked(nm); err != nil {
+		if err := b.writeNetmapToDiskLockedWithPeers(nm); err != nil {
 			b.mu.Unlock()
 			t.Fatalf("setup writeNetmapToDiskLocked: %v", err)
 		}
@@ -172,7 +172,7 @@ func TestOnHomeDERPUpdate(t *testing.T) {
 
 		// Write an initial cache entry so we can verify it is not overwritten.
 		b.mu.Lock()
-		if err := b.writeNetmapToDiskLocked(nm); err != nil {
+		if err := b.writeNetmapToDiskLockedWithPeers(nm); err != nil {
 			b.mu.Unlock()
 			t.Fatalf("setup writeNetmapToDiskLocked: %v", err)
 		}
@@ -218,7 +218,7 @@ func TestWriteNetmapDoesNotMutateOriginal(t *testing.T) {
 	b.mu.Lock()
 	defer b.mu.Unlock()
 
-	if err := b.writeNetmapToDiskLocked(nm); err != nil {
+	if err := b.writeNetmapToDiskLockedWithPeers(nm); err != nil {
 		t.Fatalf("writeNetmapToDiskLocked: %v", err)
 	}
 

ipn/ipnlocal/local.go

@@ -703,10 +703,10 @@ func (b *LocalBackend) onHomeDERPUpdateLocked(du magicsock.HomeDERPChanged) {
 		return
 	}
 
-	// Persist the full netmap (including up-to-date Peers) to disk for
-	// fast restart.
-	if err := b.writeNetmapToDiskLocked(b.NetMapWithPeers()); err != nil {
-		b.logf("write netmap to cache: %v", err)
+	// Update the relevant parts of the cached copy of the network map, notably
+	// not including the peers.
+	if err := b.writeNetmapToDiskLockedWithoutPeers(b.NetMapNoPeers()); err != nil {
+		b.logf("update netmap cache: %v", err)
 	}
 }
 
@@ -6710,7 +6710,7 @@ func (b *LocalBackend) setNetMapLocked(nm *netmap.NetworkMap) {
 	// the node starts up.
 	if nm != nil {
 		if b.currentNode().SelfHasCap(tailcfg.NodeAttrCacheNetworkMaps) && envknob.BoolDefaultTrue("TS_USE_CACHED_NETMAP") {
-			if err := b.writeNetmapToDiskLocked(nm); err != nil {
+			if err := b.writeNetmapToDiskLockedWithPeers(nm); err != nil {
 				b.logf("write netmap to cache: %v", err)
 			}
 		} else {

ipn/ipnlocal/netmapcache/netmapcache.go

@@ -83,8 +83,8 @@ func (c *Cache) writeJSON(ctx context.Context, key cacheKey, v any) error {
 		return fmt.Errorf("JSON marshalling %q: %w", key, err)
 	}
 
-	// TODO(creachadair): Maybe use a hash instead of the contents? Do we need
-	// this at all?
+	// If the digest of this value has not changed since it was last written,
+	// record that the key is of interest, but do not actually perform a write.
 	last, ok := c.lastWrote[key]
 	if ok && cacheDigest(j) == last.digest {
 		c.wantKeys.Add(key)
@@ -193,6 +193,23 @@ const (
 	packetFilterKey cacheKey = "filter"
 )
 
+// UpdateSelfOnly updates nm in the cache, replacing any previously cached
+// values for nm.Self and other tailnet metadata for the current node, but
+// skipping the Peers and UserProfiles fields. Any existing peer and profile
+// data are left unmodified.
+func (c *Cache) UpdateSelfOnly(ctx context.Context, nm *netmap.NetworkMap) error {
+	if !buildfeatures.HasCacheNetMap || nm == nil || nm.Cached {
+		return nil
+	}
+	if selfID := nm.User(); selfID == 0 {
+		return errors.New("no user in netmap")
+	}
+	// Since we are not modifying peers or users (and in particular, not
+	// removing any), we do not need to do any garbage collection on the storage
+	// keys.
+	return c.updateSelfOnly(ctx, nm)
+}
+
 // Store records nm in the cache, replacing any previously-cached values.
 func (c *Cache) Store(ctx context.Context, nm *netmap.NetworkMap) error {
 	if !buildfeatures.HasCacheNetMap || nm == nil || nm.Cached {
@@ -202,7 +219,32 @@ func (c *Cache) Store(ctx context.Context, nm *netmap.NetworkMap) error {
 		return errors.New("no user in netmap")
 	}
 
+	// Because we may modify which peers and user profiles are valid, we need to
+	// keep track of which storage keys we actually touch during the store.
+	// This is used by c.removeUnwantedKeys to clean up keys that are no longer
+	// referenced after peers and/or profiles are removed from nm.
 	clear(c.wantKeys)
+	if err := c.updateSelfOnly(ctx, nm); err != nil {
+		return err
+	}
+	for _, p := range nm.Peers {
+		key := peerKeyPrefix + cacheKey(p.StableID())
+		if err := c.writeJSON(ctx, key, netmapNode{Node: &p}); err != nil {
+			return err
+		}
+	}
+	for uid, u := range nm.UserProfiles {
+		key := fmt.Sprintf("%s%d", userKeyPrefix, uid)
+		if err := c.writeJSON(ctx, cacheKey(key), netmapUserProfile{UserProfile: &u}); err != nil {
+			return err
+		}
+	}
+	return c.removeUnwantedKeys(ctx)
+}
+
+// updateSelfOnly updates the "static" parts of the netmap in the cache.  It is
+// shared between [Cache.UpdateSelfOnly] and [Cache.Store].
+func (c *Cache) updateSelfOnly(ctx context.Context, nm *netmap.NetworkMap) error {
 	if err := c.writeJSON(ctx, miscKey, netmapMisc{
 		MachineKey:       &nm.MachineKey,
 		CollectServices:  &nm.CollectServices,
@@ -226,29 +268,15 @@ func (c *Cache) Store(ctx context.Context, nm *netmap.NetworkMap) error {
 		// N.B. The NodeKey and AllCaps fields can be recovered from SelfNode on
 		// load, and do not need to be stored separately.
 	}
-	for _, p := range nm.Peers {
-		key := peerKeyPrefix + cacheKey(p.StableID())
-		if err := c.writeJSON(ctx, key, netmapNode{Node: &p}); err != nil {
-			return err
-		}
-	}
-	for uid, u := range nm.UserProfiles {
-		key := fmt.Sprintf("%s%d", userKeyPrefix, uid)
-		if err := c.writeJSON(ctx, cacheKey(key), netmapUserProfile{UserProfile: &u}); err != nil {
-			return err
-		}
-	}
 	if err := c.writeJSON(ctx, packetFilterKey, netmapPacketFilter{Rules: &nm.PacketFilterRules}); err != nil {
 		return err
 	}
-
 	if buildfeatures.HasSSH && nm.SSHPolicy != nil {
 		if err := c.writeJSON(ctx, sshPolicyKey, netmapSSH{SSHPolicy: &nm.SSHPolicy}); err != nil {
 			return err
 		}
 	}
-
-	return c.removeUnwantedKeys(ctx)
+	return nil
 }
 
 // Load loads the cached [netmap.NetworkMap] value stored in c, if one is available.

ipn/ipnlocal/netmapcache/netmapcache_test.go

@@ -258,6 +258,55 @@ func TestInvalidCache(t *testing.T) {
 	})
 }
 
+func TestUpdateSelfOnly(t *testing.T) {
+	s := make(testStore)
+	c := netmapcache.NewCache(s)
+
+	// Initialize the cache with the test map so we get a baseline.
+	if err := c.Store(t.Context(), testMap); err != nil {
+		t.Fatalf("Store initial netmap: %v", err)
+	}
+
+	// Modify a shallow copy of the map so we can perform an update and verify
+	// that it round-trips through a Load after calling UpdateSelfOnly.
+	newSelf := &tailcfg.Node{
+		ID:           23456,
+		StableID:     "n23456FAKE",
+		User:         8675309,
+		Name:         "alt.example.com.",
+		Key:          testNodeKey,
+		HomeDERP:     6174,
+		Capabilities: []tailcfg.NodeCapability{"cap1", "cap3"},
+	}
+	updated := *testMap // shallow copy
+	updated.SelfNode = newSelf.View()
+	updated.AllCaps = set.Of[tailcfg.NodeCapability]("cap1", "cap3")
+	updated.DNS = tailcfg.DNSConfig{Domains: []string{"example3.org", "example4.horse"}}
+
+	// Empty the peers and profiles so that we can verify the update does not
+	// attempt to use them or GC based on their absence.
+	updated.Peers = nil
+	updated.UserProfiles = nil
+
+	if err := c.UpdateSelfOnly(t.Context(), &updated); err != nil {
+		t.Fatalf("UpdateSelfOnly failed: %v", err)
+	}
+
+	// Verify we got the same results back. Importantly, we expect the same
+	// peers and profiles as before, to enforce that the self-only update did
+	// not prune
+	updated.Peers = testMap.Peers
+	updated.UserProfiles = testMap.UserProfiles
+
+	got, err := c.Load(t.Context())
+	if err != nil {
+		t.Fatalf("Load netmap failed: %v", err)
+	}
+	if diff := diffNetMaps(got, &updated); diff != "" {
+		t.Fatalf("Updated map differs (-got, +want):\n%s", diff)
+	}
+}
+
 // skippedMapFields are the names of fields that should not be considered by
 // network map caching, and thus skipped when comparing test results.
 var skippedMapFields = []string{