Add Fast Copy of Trie

beacon-chain/blockchain/head.go

@@ -127,17 +127,17 @@ func (s *Service) saveHeadNoDB(ctx context.Context, b *ethpb.SignedBeaconBlock,
 			return errors.Wrap(err, "could not retrieve head state in DB")
 		}
 	} else {
-		headState, err = s.beaconDB.State(ctx, r)
-		if err != nil {
-			return errors.Wrap(err, "could not retrieve head state in DB")
+		s.initSyncStateLock.RLock()
+		cachedHeadState, ok := s.initSyncState[r]
+		if ok {
+			headState = cachedHeadState
 		}
+		s.initSyncStateLock.RUnlock()
 		if headState == nil {
-			s.initSyncStateLock.RLock()
-			cachedHeadState, ok := s.initSyncState[r]
-			if ok {
-				headState = cachedHeadState
+			headState, err = s.beaconDB.State(ctx, r)
+			if err != nil {
+				return errors.Wrap(err, "could not retrieve head state in DB")
 			}
-			s.initSyncStateLock.RUnlock()
 		}
 	}
 	if headState == nil {

beacon-chain/blockchain/process_block.go

@@ -229,7 +229,7 @@ func (s *Service) onBlockInitialSyncStateTransition(ctx context.Context, signed
 	} else {
 		s.initSyncStateLock.Lock()
 		defer s.initSyncStateLock.Unlock()
-		s.initSyncState[root] = postState.Copy()
+		s.initSyncState[root] = postState.FastCopy()
 		s.filterBoundaryCandidates(ctx, root, postState)
 	}
 

beacon-chain/blockchain/process_block_helpers.go

@@ -99,7 +99,7 @@ func (s *Service) verifyBlkPreState(ctx context.Context, b *ethpb.BeaconBlock) (
 		}
 		return preState, nil // No copy needed from newly hydrated DB object.
 	}
-	return preState.Copy(), nil
+	return preState.FastCopy(), nil
 }
 
 // verifyBlkDescendant validates input block root is a descendant of the

beacon-chain/state/state_trie.go

@@ -177,6 +177,29 @@ func (b *BeaconState) Copy() *BeaconState {
 	return dst
 }
 
+// FastCopy is used to copy all the field trie contents and
+// empty out the references to it in the source state. This
+// is used when we have a state that we know will most
+// likely not be used anytime in the future.
+func (b *BeaconState) FastCopy() *BeaconState {
+	newState := b.Copy()
+	for fldIdx, fieldTrie := range b.stateFieldLeaves {
+		if fieldTrie.reference != nil {
+			fieldTrie.Lock()
+			fieldTrie.MinusRef()
+			fieldTrie.Unlock()
+			b.rebuildTrie[fldIdx] = true
+			b.dirtyFields[fldIdx] = true
+			b.stateFieldLeaves[fldIdx] = &FieldTrie{
+				field:     fldIdx,
+				reference: &reference{1},
+				Mutex:     new(sync.Mutex),
+			}
+		}
+	}
+	return newState
+}
+
 // HashTreeRoot of the beacon state retrieves the Merkle root of the trie
 // representation of the beacon state based on the eth2 Simple Serialize specification.
 func (b *BeaconState) HashTreeRoot(ctx context.Context) ([32]byte, error) {