Fix getAttPreState memory usage for new-state-mgmt (#5584)

* Starting

* Revert "update commit hash (#5554)"

This reverts commit c69f561.

* Various memory fixes

* Revert "Revert "update commit hash (#5554)""

This reverts commit c9a174b.

* Better error message

* Provide getters

* Using it in DomainData

* Fixed test

* Tests

* Comment

* Update cache for regular mode

Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>

beacon-chain/blockchain/process_attestation_helpers.go

@@ -31,54 +31,78 @@ func (s *Service) getAttPreState(ctx context.Context, c *ethpb.Checkpoint) (*sta
 		return cachedState, nil
 	}
 
-	var baseState *stateTrie.BeaconState
 	if featureconfig.Get().NewStateMgmt {
 		if !s.stateGen.HasState(ctx, bytesutil.ToBytes32(c.Root)) {
 			if err := s.beaconDB.SaveBlocks(ctx, s.getInitSyncBlocks()); err != nil {
 				return nil, errors.Wrap(err, "could not save initial sync blocks")
 			}
 			s.clearInitSyncBlocks()
 		}
-		baseState, err = s.stateGen.StateByRoot(ctx, bytesutil.ToBytes32(c.Root))
+
+		baseState, err := s.stateGen.StateByRoot(ctx, bytesutil.ToBytes32(c.Root))
 		if err != nil {
 			return nil, errors.Wrapf(err, "could not get pre state for slot %d", helpers.StartSlot(c.Epoch))
 		}
-	} else {
-		if featureconfig.Get().CheckHeadState {
-			headRoot, err := s.HeadRoot(ctx)
+
+		if helpers.StartSlot(c.Epoch) > baseState.Slot() {
+			baseState = baseState.Copy()
+			baseState, err = state.ProcessSlots(ctx, baseState, helpers.StartSlot(c.Epoch))
 			if err != nil {
-				return nil, errors.Wrapf(err, "could not get head root")
-			}
-			if bytes.Equal(headRoot, c.Root) {
-				st, err := s.HeadState(ctx)
-				if err != nil {
-					return nil, errors.Wrapf(err, "could not get head state")
-				}
-				if err := s.checkpointState.AddCheckpointState(&cache.CheckpointState{
-					Checkpoint: c,
-					State:      st.Copy(),
-				}); err != nil {
-					return nil, errors.Wrap(err, "could not saved checkpoint state to cache")
-				}
-				return st, nil
+				return nil, errors.Wrapf(err, "could not process slots up to %d", helpers.StartSlot(c.Epoch))
 			}
 		}
 
-		baseState, err = s.beaconDB.State(ctx, bytesutil.ToBytes32(c.Root))
+		if err := s.checkpointState.AddCheckpointState(&cache.CheckpointState{
+			Checkpoint: c,
+			State:      baseState,
+		}); err != nil {
+			return nil, errors.Wrap(err, "could not saved checkpoint state to cache")
+		}
+
+		return baseState, nil
+	}
+
+	if featureconfig.Get().CheckHeadState {
+		headRoot, err := s.HeadRoot(ctx)
 		if err != nil {
-			return nil, errors.Wrapf(err, "could not get pre state for slot %d", helpers.StartSlot(c.Epoch))
+			return nil, errors.Wrapf(err, "could not get head root")
+		}
+		if bytes.Equal(headRoot, c.Root) {
+			st, err := s.HeadState(ctx)
+			if err != nil {
+				return nil, errors.Wrapf(err, "could not get head state")
+			}
+			if err := s.checkpointState.AddCheckpointState(&cache.CheckpointState{
+				Checkpoint: c,
+				State:      st.Copy(),
+			}); err != nil {
+				return nil, errors.Wrap(err, "could not saved checkpoint state to cache")
+			}
+			return st, nil
 		}
 	}
 
+	baseState, err := s.beaconDB.State(ctx, bytesutil.ToBytes32(c.Root))
+	if err != nil {
+		return nil, errors.Wrapf(err, "could not get pre state for slot %d", helpers.StartSlot(c.Epoch))
+	}
 	if baseState == nil {
 		return nil, fmt.Errorf("pre state of target block %d does not exist", helpers.StartSlot(c.Epoch))
 	}
 
 	if helpers.StartSlot(c.Epoch) > baseState.Slot() {
-		baseState, err = state.ProcessSlots(ctx, baseState, helpers.StartSlot(c.Epoch))
+		savedState := baseState.Copy()
+		savedState, err = state.ProcessSlots(ctx, savedState, helpers.StartSlot(c.Epoch))
 		if err != nil {
 			return nil, errors.Wrapf(err, "could not process slots up to %d", helpers.StartSlot(c.Epoch))
 		}
+		if err := s.checkpointState.AddCheckpointState(&cache.CheckpointState{
+			Checkpoint: c,
+			State:      savedState.Copy(),
+		}); err != nil {
+			return nil, errors.Wrap(err, "could not saved checkpoint state to cache")
+		}
+		return savedState, nil
 	}
 
 	if err := s.checkpointState.AddCheckpointState(&cache.CheckpointState{

beacon-chain/blockchain/receive_attestation.go

@@ -60,7 +60,7 @@ func (s *Service) ReceiveAttestationNoPubsub(ctx context.Context, att *ethpb.Att
 func (s *Service) IsValidAttestation(ctx context.Context, att *ethpb.Attestation) bool {
 	baseState, err := s.getAttPreState(ctx, att.Data.Target)
 	if err != nil {
-		log.WithError(err).Error("Failed to validate attestation")
+		log.WithError(err).Error("Failed to get attestation pre state")
 		return false
 	}
 

beacon-chain/cache/checkpoint_state.go

@@ -93,7 +93,8 @@ func (c *CheckpointStateCache) StateByCheckpoint(cp *ethpb.Checkpoint) (*stateTr
 		return nil, ErrNotCheckpointState
 	}
 
-	return info.State.Copy(), nil
+	// Copy here is unnecessary since the return will only be used to verify attestation signature.
+	return info.State, nil
 }
 
 // AddCheckpointState adds CheckpointState object to the cache. This method also trims the least