5 changed files with 133 additions and 37 deletions
@@ -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) {
@@ -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)
 	}

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