Add and use HasStateInCache (#7691)

beacon-chain/blockchain/process_attestation_helpers.go

@@ -63,7 +63,10 @@ func (s *Service) getAttPreState(ctx context.Context, c *ethpb.Checkpoint) (*sta
 		}
 		return baseState, nil
 	}
-	has, err := s.stateGen.HasState(ctx, bytesutil.ToBytes32(c.Root))
+
+	// To avoid sharing the same state across checkpoint state cache and hot state cache,
+	// we don't add the state to check point cache.
+	has, err := s.stateGen.HasStateInCache(ctx, bytesutil.ToBytes32(c.Root))
 	if err != nil {
 		return nil, err
 	}

beacon-chain/state/stategen/getter.go

@@ -13,10 +13,7 @@ import (
 
 // HasState returns true if the state exists in cache or in DB.
 func (s *State) HasState(ctx context.Context, blockRoot [32]byte) (bool, error) {
-	if s.hotStateCache.Has(blockRoot) {
-		return true, nil
-	}
-	_, has, err := s.epochBoundaryStateCache.getByRoot(blockRoot)
+	has, err := s.HasStateInCache(ctx, blockRoot)
 	if err != nil {
 		return false, err
 	}
@@ -26,6 +23,18 @@ func (s *State) HasState(ctx context.Context, blockRoot [32]byte) (bool, error)
 	return s.beaconDB.HasState(ctx, blockRoot), nil
 }
 
+// HasStateInCache returns true if the state exists in cache.
+func (s *State) HasStateInCache(ctx context.Context, blockRoot [32]byte) (bool, error) {
+	if s.hotStateCache.Has(blockRoot) {
+		return true, nil
+	}
+	_, has, err := s.epochBoundaryStateCache.getByRoot(blockRoot)
+	if err != nil {
+		return false, err
+	}
+	return has, nil
+}
+
 // StateByRoot retrieves the state using input block root.
 func (s *State) StateByRoot(ctx context.Context, blockRoot [32]byte) (*state.BeaconState, error) {
 	ctx, span := trace.StartSpan(ctx, "stateGen.StateByRoot")

beacon-chain/state/stategen/getter_test.go

@@ -437,3 +437,60 @@ func TestLastAncestorState_CanGetUsingCache(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, b1State.Slot(), lastState.Slot(), "Did not get wanted state")
 }
+
+func TestState_HasState(t *testing.T) {
+	ctx := context.Background()
+	db, _ := testDB.SetupDB(t)
+	service := New(db, cache.NewStateSummaryCache())
+	s := testutil.NewBeaconState()
+	rHit1 := [32]byte{1}
+	rHit2 := [32]byte{2}
+	rMiss := [32]byte{3}
+	service.hotStateCache.Put(rHit1, s)
+	require.NoError(t, service.epochBoundaryStateCache.put(rHit2, s))
+
+	b := testutil.NewBeaconBlock()
+	rHit3, err := b.Block.HashTreeRoot()
+	require.NoError(t, err)
+	require.NoError(t, service.beaconDB.SaveState(ctx, s, rHit3))
+	tt := []struct {
+		root [32]byte
+		want bool
+	}{
+		{rHit1, true},
+		{rHit2, true},
+		{rMiss, false},
+		{rHit3, true},
+	}
+	for _, tc := range tt {
+		got, err := service.HasState(ctx, tc.root)
+		require.NoError(t, err)
+		require.Equal(t, tc.want, got)
+	}
+}
+
+func TestState_HasStateInCache(t *testing.T) {
+	ctx := context.Background()
+	db, _ := testDB.SetupDB(t)
+	service := New(db, cache.NewStateSummaryCache())
+	s := testutil.NewBeaconState()
+	rHit1 := [32]byte{1}
+	rHit2 := [32]byte{2}
+	rMiss := [32]byte{3}
+	service.hotStateCache.Put(rHit1, s)
+	require.NoError(t, service.epochBoundaryStateCache.put(rHit2, s))
+
+	tt := []struct {
+		root [32]byte
+		want bool
+	}{
+		{rHit1, true},
+		{rHit2, true},
+		{rMiss, false},
+	}
+	for _, tc := range tt {
+		got, err := service.HasStateInCache(ctx, tc.root)
+		require.NoError(t, err)
+		require.Equal(t, tc.want, got)
+	}
+}