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/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{