From c29ad761a7cf098cbcd404b5d1ae1e559a5ddc53 Mon Sep 17 00:00:00 2001
From: Stephen Buttolph <stephen@avalabs.org>
Date: Wed, 9 Aug 2023 22:18:05 -0400
Subject: [PATCH] Check P-chain ShouldPrune during Initialize (#1836)

---
 vms/platformvm/state/mock_state.go | 15 +++++++++++++++
 vms/platformvm/state/state.go      | 28 +++++++++++-----------------
 vms/platformvm/vm.go               | 13 +++++++++++++
 3 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/vms/platformvm/state/mock_state.go b/vms/platformvm/state/mock_state.go
index 883cc3f84f00..7e08bedb44b7 100644
--- a/vms/platformvm/state/mock_state.go
+++ b/vms/platformvm/state/mock_state.go
@@ -726,6 +726,21 @@ func (mr *MockStateMockRecorder) SetUptime(arg0, arg1, arg2, arg3 interface{}) *
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetUptime", reflect.TypeOf((*MockState)(nil).SetUptime), arg0, arg1, arg2, arg3)
 }
 
+// ShouldPrune mocks base method.
+func (m *MockState) ShouldPrune() (bool, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "ShouldPrune")
+	ret0, _ := ret[0].(bool)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// ShouldPrune indicates an expected call of ShouldPrune.
+func (mr *MockStateMockRecorder) ShouldPrune() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ShouldPrune", reflect.TypeOf((*MockState)(nil).ShouldPrune))
+}
+
 // UTXOIDs mocks base method.
 func (m *MockState) UTXOIDs(arg0 []byte, arg1 ids.ID, arg2 int) ([]ids.ID, error) {
 	m.ctrl.T.Helper()
diff --git a/vms/platformvm/state/state.go b/vms/platformvm/state/state.go
index 5626696eed30..b70e0cb6b53e 100644
--- a/vms/platformvm/state/state.go
+++ b/vms/platformvm/state/state.go
@@ -183,6 +183,12 @@ type State interface {
 	// Discard uncommitted changes to the database.
 	Abort()
 
+	// Returns if the state should be pruned and indexed to remove rejected
+	// blocks and generate the block height index.
+	//
+	// TODO: Remove after v1.11.x is activated
+	ShouldPrune() (bool, error)
+
 	// Removes rejected blocks from disk and indexes accepted blocks by height. This
 	// function supports being (and is recommended to be) called asynchronously.
 	//
@@ -467,11 +473,14 @@ func New(
 	// to be run.
 	//
 	// TODO: Cleanup after v1.11.x is activated
-	shouldPrune, err := s.shouldPrune()
+	shouldPrune, err := s.ShouldPrune()
 	if err != nil {
 		return nil, err
 	}
 	if shouldPrune {
+		// If the pruned key is on disk, we must delete it to ensure our disk
+		// can't get into a partially pruned state if the node restarts mid-way
+		// through pruning.
 		if err := s.singletonDB.Delete(prunedKey); err != nil {
 			return nil, fmt.Errorf("failed to remove prunedKey from singletonDB: %w", err)
 		}
@@ -738,7 +747,7 @@ func (s *state) doneInit() error {
 	return s.singletonDB.Put(initializedKey, nil)
 }
 
-func (s *state) shouldPrune() (bool, error) {
+func (s *state) ShouldPrune() (bool, error) {
 	has, err := s.singletonDB.Has(prunedKey)
 	if err != nil {
 		return true, err
@@ -2343,21 +2352,6 @@ func parseStoredBlock(blkBytes []byte) (blocks.Block, choices.Status, bool, erro
 
 func (s *state) PruneAndIndex(lock sync.Locker, log logging.Logger) error {
 	lock.Lock()
-	shouldPrune, err := s.shouldPrune()
-	if err != nil {
-		lock.Unlock()
-		return fmt.Errorf(
-			"failed to check if the database should be pruned: %w",
-			err,
-		)
-	}
-	if !shouldPrune {
-		lock.Unlock()
-
-		log.Info("state already pruned and indexed")
-		return nil
-	}
-
 	// It is possible that new blocks are added after grabbing this iterator. New
 	// blocks are guaranteed to be accepted and height-indexed, so we don't need to
 	// check them.
diff --git a/vms/platformvm/vm.go b/vms/platformvm/vm.go
index d554faa11bf0..2d051c71d81f 100644
--- a/vms/platformvm/vm.go
+++ b/vms/platformvm/vm.go
@@ -219,6 +219,19 @@ func (vm *VM) Initialize(
 		return err
 	}
 
+	shouldPrune, err := vm.state.ShouldPrune()
+	if err != nil {
+		return fmt.Errorf(
+			"failed to check if the database should be pruned: %w",
+			err,
+		)
+	}
+	if !shouldPrune {
+		chainCtx.Log.Info("state already pruned and indexed")
+		vm.pruned.Set(true)
+		return nil
+	}
+
 	go func() {
 		err := vm.state.PruneAndIndex(&vm.ctx.Lock, vm.ctx.Log)
 		if err != nil {