stmgr: Handle state forks with cron on null blocks correctly

chain/stmgr/forks.go

@@ -3,23 +3,20 @@ package stmgr
 import (
 	"context"
 
+	"github.com/filecoin-project/lotus/chain/types"
 	"github.com/filecoin-project/specs-actors/actors/abi"
-	"github.com/ipfs/go-cid"
 )
 
-var ForksAtHeight = map[abi.ChainEpoch]func(context.Context, *StateManager, cid.Cid) (cid.Cid, error){}
+var ForksAtHeight = map[abi.ChainEpoch]func(context.Context, *StateManager, types.StateTree) error{}
 
-func (sm *StateManager) handleStateForks(ctx context.Context, pstate cid.Cid, height, parentH abi.ChainEpoch) (_ cid.Cid, err error) {
+func (sm *StateManager) handleStateForks(ctx context.Context, st types.StateTree, height abi.ChainEpoch) (err error) {
-	for i := parentH; i < height; i++ {
+	f, ok := ForksAtHeight[height]
-		f, ok := ForksAtHeight[i]
 	if ok {
-			nstate, err := f(ctx, sm, pstate)
+		err := f(ctx, sm, st)
 		if err != nil {
-				return cid.Undef, err
+			return err
-			}
-			pstate = nstate
 		}
 	}
 
-	return pstate, nil
+	return nil
 }

chain/stmgr/forks_test.go

@@ -21,7 +21,6 @@ import (
 	"github.com/filecoin-project/lotus/chain/actors"
 	"github.com/filecoin-project/lotus/chain/actors/aerrors"
 	"github.com/filecoin-project/lotus/chain/gen"
-	"github.com/filecoin-project/lotus/chain/state"
 	"github.com/filecoin-project/lotus/chain/stmgr"
 	. "github.com/filecoin-project/lotus/chain/stmgr"
 	"github.com/filecoin-project/lotus/chain/types"
@@ -122,37 +121,33 @@ func TestForkHeightTriggers(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	stmgr.ForksAtHeight[testForkHeight] = func(ctx context.Context, sm *StateManager, pstate cid.Cid) (cid.Cid, error) {
+	stmgr.ForksAtHeight[testForkHeight] = func(ctx context.Context, sm *StateManager, st types.StateTree) error {
 		cst := cbor.NewCborStore(sm.ChainStore().Blockstore())
-		st, err := state.LoadStateTree(cst, pstate)
-		if err != nil {
-			return cid.Undef, err
-		}
 
 		act, err := st.GetActor(taddr)
 		if err != nil {
-			return cid.Undef, err
+			return err
 		}
 
 		var tas testActorState
 		if err := cst.Get(ctx, act.Head, &tas); err != nil {
-			return cid.Undef, xerrors.Errorf("in fork handler, failed to run get: %w", err)
+			return xerrors.Errorf("in fork handler, failed to run get: %w", err)
 		}
 
 		tas.HasUpgraded = 55
 
 		ns, err := cst.Put(ctx, &tas)
 		if err != nil {
-			return cid.Undef, err
+			return err
 		}
 
 		act.Head = ns
 
 		if err := st.SetActor(taddr, act); err != nil {
-			return cid.Undef, err
+			return err
 		}
 
-		return st.Flush(ctx)
+		return nil
 	}
 
 	inv.Register(builtin.PaymentChannelActorCodeID, &testActor{}, &testActorState{})

chain/stmgr/stmgr.go

@@ -147,7 +147,7 @@ type BlockMessages struct {
 type ExecCallback func(cid.Cid, *types.Message, *vm.ApplyRet) error
 
 func (sm *StateManager) ApplyBlocks(ctx context.Context, parentEpoch abi.ChainEpoch, pstate cid.Cid, bms []BlockMessages, epoch abi.ChainEpoch, r vm.Rand, cb ExecCallback) (cid.Cid, cid.Cid, error) {
-	vmi, err := sm.newVM(pstate, parentEpoch+1, r, sm.cs.Blockstore(), sm.cs.VMSys())
+	vmi, err := sm.newVM(pstate, parentEpoch, r, sm.cs.Blockstore(), sm.cs.VMSys())
 	if err != nil {
 		return cid.Undef, cid.Undef, xerrors.Errorf("instantiating VM failed: %w", err)
 	}
@@ -185,11 +185,19 @@ func (sm *StateManager) ApplyBlocks(ctx context.Context, parentEpoch abi.ChainEp
 		return nil
 	}
 
+	for i := parentEpoch; i < epoch; i++ {
+		// handle state forks
+		err = sm.handleStateForks(ctx, vmi.StateTree(), i)
+		if err != nil {
+			return cid.Undef, cid.Undef, xerrors.Errorf("error handling state forks: %w", err)
+		}
+
+		if i > parentEpoch {
 			// run cron for null rounds if any
-	for i := parentEpoch + 1; i < epoch; i++ {
 			if err := runCron(); err != nil {
 				return cid.Cid{}, cid.Cid{}, err
 			}
+		}
 
 		vmi.SetBlockHeight(i + 1)
 	}
@@ -302,18 +310,13 @@ func (sm *StateManager) computeTipSetState(ctx context.Context, blks []*types.Bl
 
 	var parentEpoch abi.ChainEpoch
 	pstate := blks[0].ParentStateRoot
-	if len(blks[0].Parents) > 0 { // don't support forks on genesis
+	if len(blks[0].Parents) > 0 {
 		parent, err := sm.cs.GetBlock(blks[0].Parents[0])
 		if err != nil {
 			return cid.Undef, cid.Undef, xerrors.Errorf("getting parent block: %w", err)
 		}
 
 		parentEpoch = parent.Height
-
-		pstate, err = sm.handleStateForks(ctx, blks[0].ParentStateRoot, blks[0].Height, parent.Height)
-		if err != nil {
-			return cid.Undef, cid.Undef, xerrors.Errorf("error handling state forks: %w", err)
-		}
 	}
 
 	cids := make([]cid.Cid, len(blks))

chain/stmgr/utils.go

@@ -439,15 +439,20 @@ func ComputeState(ctx context.Context, sm *StateManager, height abi.ChainEpoch,
 		return cid.Undef, nil, err
 	}
 
-	fstate, err := sm.handleStateForks(ctx, base, height, ts.Height())
+	r := store.NewChainRand(sm.cs, ts.Cids(), height)
+	vmi, err := vm.NewVM(base, height, r, sm.cs.Blockstore(), sm.cs.VMSys())
 	if err != nil {
 		return cid.Undef, nil, err
 	}
 
-	r := store.NewChainRand(sm.cs, ts.Cids(), height)
+	for i := ts.Height(); i < height; i++ {
-	vmi, err := vm.NewVM(fstate, height, r, sm.cs.Blockstore(), sm.cs.VMSys())
+		// handle state forks
+		err = sm.handleStateForks(ctx, vmi.StateTree(), i)
 		if err != nil {
-		return cid.Undef, nil, err
+			return cid.Undef, nil, xerrors.Errorf("error handling state forks: %w", err)
+		}
+
+		// TODO: should we also run cron here?
 	}
 
 	for i, msg := range msgs {