Cache subnetIDs in cachedState (#4402)

Signed-off-by: Michael Kaplan <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: michaelkaplan13

snow/validators/state.go

@@ -13,7 +13,10 @@ import (
 	"github.com/ava-labs/avalanchego/ids"
 )
 
-const validatorSetsCacheSize = 8
+const (
+	validatorSetsCacheSize = 8
+	subnetIDsCacheSize     = 4096 // At 64 bytes per entry, this is ~256 KB
+)
 
 var (
 	_ State = (*lockedState)(nil)
@@ -189,10 +192,15 @@ type cachedState struct {
 	// long as it is cached.
 	activation time.Time
 
+	// Caches the subnet ID for given blockchain IDs.
+	// Key: blockchain ID
+	// Value: subnet ID
+	subnetIDsCache cache.Cacher[ids.ID, ids.ID]
+
 	// Caches validators for all subnets at various heights.
 	// Key: height
 	// Value: mapping subnet ID -> validator set
-	cache cache.Cacher[uint64, map[ids.ID]WarpSet]
+	validatorSetsCache cache.Cacher[uint64, map[ids.ID]WarpSet]
 }
 
 // TODO: Remove the graniteActivation parameter once all networks have
@@ -202,25 +210,38 @@ func NewCachedState(
 	graniteActivation time.Time,
 ) State {
 	return &cachedState{
-		State:      state,
-		activation: graniteActivation,
-		cache:      lru.NewCache[uint64, map[ids.ID]WarpSet](validatorSetsCacheSize),
+		State:              state,
+		activation:         graniteActivation,
+		validatorSetsCache: lru.NewCache[uint64, map[ids.ID]WarpSet](validatorSetsCacheSize),
+		subnetIDsCache:     lru.NewCache[ids.ID, ids.ID](subnetIDsCacheSize),
 	}
 }
 
+func (c *cachedState) GetSubnetID(ctx context.Context, chainID ids.ID) (ids.ID, error) {
+	if s, ok := c.subnetIDsCache.Get(chainID); ok {
+		return s, nil
+	}
+	s, err := c.State.GetSubnetID(ctx, chainID)
+	if err != nil {
+		return ids.Empty, err
+	}
+	c.subnetIDsCache.Put(chainID, s)
+	return s, nil
+}
+
 func (c *cachedState) GetWarpValidatorSets(
 	ctx context.Context,
 	height uint64,
 ) (map[ids.ID]WarpSet, error) {
-	if s, ok := c.cache.Get(height); ok {
+	if s, ok := c.validatorSetsCache.Get(height); ok {
 		return s, nil
 	}
 
 	s, err := c.State.GetWarpValidatorSets(ctx, height)
 	if err != nil {
 		return nil, err
 	}
-	c.cache.Put(height, s)
+	c.validatorSetsCache.Put(height, s)
 	return s, nil
 }
 

snow/validators/state_test.go

@@ -21,6 +21,88 @@ import (
 
 var errTest = errors.New("test error")
 
+func TestCachedState_GetSubnetID(t *testing.T) {
+	require := require.New(t)
+
+	uncached := &validatorstest.State{
+		T:               t,
+		CantGetSubnetID: true,
+	}
+	cached := NewCachedState(uncached, upgrade.InitiallyActiveTime)
+
+	blockchainID1 := ids.GenerateTestID()
+	blockchainID2 := ids.GenerateTestID()
+	subnetIDs := map[ids.ID]ids.ID{
+		blockchainID1: ids.GenerateTestID(),
+		blockchainID2: ids.GenerateTestID(),
+	}
+	tests := []struct {
+		name         string
+		expectCached bool
+		blockchainID ids.ID
+		want         ids.ID
+		wantErr      error
+	}{
+		{
+			name:         "populate_initial_entry",
+			expectCached: false,
+			blockchainID: blockchainID1,
+			want:         subnetIDs[blockchainID1],
+		},
+		{
+			name:         "initial_entry_was_cached",
+			expectCached: true,
+			blockchainID: blockchainID1,
+			want:         subnetIDs[blockchainID1],
+		},
+		{
+			name:         "cache_miss_error",
+			expectCached: false,
+			blockchainID: blockchainID2,
+			wantErr:      errTest,
+		},
+		{
+			name:         "cache_after_miss_error",
+			expectCached: false,
+			blockchainID: blockchainID2,
+			want:         subnetIDs[blockchainID2],
+		},
+		{
+			name:         "second_cache_hit",
+			expectCached: true,
+			blockchainID: blockchainID2,
+			want:         subnetIDs[blockchainID2],
+		},
+		{
+			name:         "cache_multiple",
+			expectCached: true,
+			blockchainID: blockchainID1,
+			want:         subnetIDs[blockchainID1],
+		},
+	}
+	for _, test := range tests {
+		t.Logf("starting test: %s", test.name)
+
+		var cacheMiss bool
+		if !test.expectCached {
+			uncached.GetSubnetIDF = func(_ context.Context, chainID ids.ID) (ids.ID, error) {
+				require.Equal(test.blockchainID, chainID)
+				require.False(cacheMiss)
+
+				cacheMiss = true
+				return test.want, test.wantErr
+			}
+		} else {
+			uncached.GetSubnetIDF = nil
+		}
+
+		got, err := cached.GetSubnetID(t.Context(), test.blockchainID)
+		require.ErrorIs(err, test.wantErr)
+		require.Equal(test.want, got)
+		require.Equal(test.expectCached, !cacheMiss)
+	}
+}
+
 func TestCachedState_GetWarpValidatorSets(t *testing.T) {
 	require := require.New(t)