Load cold state by root (#5086)

beacon-chain/state/stategen/cold.go

@@ -37,6 +37,66 @@ func (s *State) saveColdState(ctx context.Context, blockRoot [32]byte, state *st
 	return nil
 }
 
+// This loads the cold state by block root, it decides whether to load from archived point (faster) or
+// somewhere between archived points (slower) because it requires replaying blocks.
+// This method is more efficient than load cold state by slot.
+func (s *State) loadColdStateByRoot(ctx context.Context, blockRoot [32]byte) (*state.BeaconState, error) {
+	ctx, span := trace.StartSpan(ctx, "stateGen.loadColdStateByRoot")
+	defer span.End()
+
+	summary, err := s.beaconDB.StateSummary(ctx, blockRoot)
+	if err != nil {
+		return nil, err
+	}
+	if summary == nil {
+		return nil, errUnknownStateSummary
+	}
+
+	// Use the archived point state if the summary slot lies on top of the archived point.
+	if summary.Slot%s.slotsPerArchivedPoint == 0 {
+		archivedPoint := summary.Slot / s.slotsPerArchivedPoint
+		s, err := s.loadColdStateByArchivedPoint(ctx, archivedPoint)
+		if err != nil {
+			return nil, errors.Wrap(err, "could not get cold state using archived index")
+		}
+		if s == nil {
+			return nil, errUnknownArchivedState
+		}
+		return s, nil
+	}
+
+	return s.loadColdIntermediateStateByRoot(ctx, summary.Slot, blockRoot)
+}
+
+// This loads the cold state for the input archived point.
+func (s *State) loadColdStateByArchivedPoint(ctx context.Context, archivedPoint uint64) (*state.BeaconState, error) {
+	return s.beaconDB.ArchivedPointState(ctx, archivedPoint)
+}
+
+// This loads a cold state by slot and block root combinations.
+// This is a faster implementation than by slot given the input block root is provided.
+func (s *State) loadColdIntermediateStateByRoot(ctx context.Context, slot uint64, blockRoot [32]byte) (*state.BeaconState, error) {
+	ctx, span := trace.StartSpan(ctx, "stateGen.loadColdIntermediateStateByRoot")
+	defer span.End()
+
+	// Load the archive point for lower side of the intermediate state.
+	lowArchivedPointIdx := slot / s.slotsPerArchivedPoint
+	lowArchivedPointState, err := s.archivedPointByIndex(ctx, lowArchivedPointIdx)
+	if err != nil {
+		return nil, errors.Wrap(err, "could not get lower archived state using index")
+	}
+	if lowArchivedPointState == nil {
+		return nil, errUnknownArchivedState
+	}
+
+	replayBlks, err := s.LoadBlocks(ctx, lowArchivedPointState.Slot()+1, slot, blockRoot)
+	if err != nil {
+		return nil, errors.Wrap(err, "could not get load blocks for cold state using slot")
+	}
+
+	return s.ReplayBlocks(ctx, lowArchivedPointState, replayBlks, slot)
+}
+
 // Given the archive index, this returns the archived cold state in the DB.
 // If the archived state does not exist in the state, it'll compute it and save it.
 func (s *State) archivedPointByIndex(ctx context.Context, archiveIndex uint64) (*state.BeaconState, error) {

beacon-chain/state/stategen/cold_test.go

@@ -59,6 +59,79 @@ func TestSaveColdState_CanSave(t *testing.T) {
 	}
 }
 
+func TestLoadColdStateByRoot_NoStateSummary(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+	if _, err := service.loadColdStateByRoot(ctx, [32]byte{'a'}); err != errUnknownStateSummary {
+		t.Fatal("Did not get correct error")
+	}
+}
+
+func TestLoadColdStateByRoot_ByArchivedPoint(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+	service.slotsPerArchivedPoint = 1
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	if err := service.beaconDB.SaveArchivedPointState(ctx, beaconState, 1); err != nil {
+		t.Fatal(err)
+	}
+	r := [32]byte{'a'}
+	if err := service.beaconDB.SaveStateSummary(ctx, &pb.StateSummary{
+		Root: r[:],
+		Slot: 1,
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	loadedState, err := service.loadColdStateByRoot(ctx, r)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !proto.Equal(loadedState.InnerStateUnsafe(), beaconState.InnerStateUnsafe()) {
+		t.Error("Did not correctly save state")
+	}
+}
+
+func TestLoadColdStateByRoot_IntermediatePlayback(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+	service.slotsPerArchivedPoint = 2
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	if err := service.beaconDB.SaveArchivedPointState(ctx, beaconState, 1); err != nil {
+		t.Fatal(err)
+	}
+	if err := service.beaconDB.SaveArchivedPointRoot(ctx, [32]byte{}, 1); err != nil {
+		t.Fatal(err)
+	}
+	r := [32]byte{'a'}
+	slot := uint64(3)
+	if err := service.beaconDB.SaveStateSummary(ctx, &pb.StateSummary{
+		Root: r[:],
+		Slot: slot,
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	loadedState, err := service.loadColdStateByRoot(ctx, r)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if loadedState.Slot() != slot {
+		t.Error("Did not correctly save state")
+	}
+}
+
 func TestArchivedPointByIndex_HasPoint(t *testing.T) {
 	ctx := context.Background()
 	db := testDB.SetupDB(t)