From fe912223bdfa000b008c143b21f2cca1c0ac7388 Mon Sep 17 00:00:00 2001
From: Steven Allen <steven@stebalien.com>
Date: Wed, 7 Oct 2020 13:56:38 -0700
Subject: [PATCH] pass an explicit upgrade height to migrations

The tipset height may not be the correct one, given null blocks.
---
 chain/stmgr/forks.go      | 23 +++++++++++++----------
 chain/stmgr/forks_test.go |  2 +-
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/chain/stmgr/forks.go b/chain/stmgr/forks.go
index 52c764ba2..b897c86eb 100644
--- a/chain/stmgr/forks.go
+++ b/chain/stmgr/forks.go
@@ -33,7 +33,14 @@ import (
 	"github.com/filecoin-project/lotus/chain/vm"
 )
 
-type UpgradeFunc func(context.Context, *StateManager, ExecCallback, cid.Cid, *types.TipSet) (cid.Cid, error)
+// UpgradeFunc is a migration function run ate very upgrade.
+//
+// - The oldState is the state produced by the upgrade epoch.
+// - The returned newState is the new state that will be used by the next epoch.
+// - The height is the upgrade epoch height (already executed).
+// - The tipset is the tipset for the last non-null block before the upgrade. Do
+//   not assume that ts.Height() is the upgrade height.
+type UpgradeFunc func(ctx context.Context, sm *StateManager, cb ExecCallback, oldState cid.Cid, height abi.ChainEpoch, ts *types.TipSet) (newState cid.Cid, err error)
 
 type Upgrade struct {
 	Height    abi.ChainEpoch
@@ -108,7 +115,7 @@ func (sm *StateManager) handleStateForks(ctx context.Context, root cid.Cid, heig
 	var err error
 	f, ok := sm.stateMigrations[height]
 	if ok {
-		retCid, err = f(ctx, sm, cb, root, ts)
+		retCid, err = f(ctx, sm, cb, root, height, ts)
 		if err != nil {
 			return cid.Undef, err
 		}
@@ -179,7 +186,7 @@ func doTransfer(cb ExecCallback, tree types.StateTree, from, to address.Address,
 	return nil
 }
 
-func UpgradeFaucetBurnRecovery(ctx context.Context, sm *StateManager, cb ExecCallback, root cid.Cid, ts *types.TipSet) (cid.Cid, error) {
+func UpgradeFaucetBurnRecovery(ctx context.Context, sm *StateManager, cb ExecCallback, root cid.Cid, epoch abi.ChainEpoch, ts *types.TipSet) (cid.Cid, error) {
 	// Some initial parameters
 	FundsForMiners := types.FromFil(1_000_000)
 	LookbackEpoch := abi.ChainEpoch(32000)
@@ -431,11 +438,9 @@ func UpgradeFaucetBurnRecovery(ctx context.Context, sm *StateManager, cb ExecCal
 	return tree.Flush(ctx)
 }
 
-func UpgradeIgnition(ctx context.Context, sm *StateManager, cb ExecCallback, root cid.Cid, ts *types.TipSet) (cid.Cid, error) {
+func UpgradeIgnition(ctx context.Context, sm *StateManager, cb ExecCallback, root cid.Cid, epoch abi.ChainEpoch, ts *types.TipSet) (cid.Cid, error) {
 	store := sm.cs.Store(ctx)
 
-	epoch := ts.Height() - 1
-
 	if build.UpgradeLiftoffHeight <= epoch {
 		return cid.Undef, xerrors.Errorf("liftoff height must be beyond ignition height")
 	}
@@ -488,11 +493,9 @@ func UpgradeIgnition(ctx context.Context, sm *StateManager, cb ExecCallback, roo
 	return tree.Flush(ctx)
 }
 
-func UpgradeActorsV2(ctx context.Context, sm *StateManager, cb ExecCallback, root cid.Cid, ts *types.TipSet) (cid.Cid, error) {
+func UpgradeActorsV2(ctx context.Context, sm *StateManager, cb ExecCallback, root cid.Cid, epoch abi.ChainEpoch, ts *types.TipSet) (cid.Cid, error) {
 	store := sm.cs.Store(ctx)
 
-	epoch := ts.Height() - 1
-
 	info, err := store.Put(ctx, new(types.StateInfo0))
 	if err != nil {
 		return cid.Undef, xerrors.Errorf("failed to create new state info for actors v2: %w", err)
@@ -539,7 +542,7 @@ func UpgradeActorsV2(ctx context.Context, sm *StateManager, cb ExecCallback, roo
 	return newRoot, nil
 }
 
-func UpgradeLiftoff(ctx context.Context, sm *StateManager, cb ExecCallback, root cid.Cid, ts *types.TipSet) (cid.Cid, error) {
+func UpgradeLiftoff(ctx context.Context, sm *StateManager, cb ExecCallback, root cid.Cid, epoch abi.ChainEpoch, ts *types.TipSet) (cid.Cid, error) {
 	tree, err := sm.StateTree(root)
 	if err != nil {
 		return cid.Undef, xerrors.Errorf("getting state tree: %w", err)
diff --git a/chain/stmgr/forks_test.go b/chain/stmgr/forks_test.go
index bb03f13b9..daa39a8d6 100644
--- a/chain/stmgr/forks_test.go
+++ b/chain/stmgr/forks_test.go
@@ -120,7 +120,7 @@ func TestForkHeightTriggers(t *testing.T) {
 			Network: 1,
 			Height:  testForkHeight,
 			Migration: func(ctx context.Context, sm *StateManager, cb ExecCallback,
-				root cid.Cid, ts *types.TipSet) (cid.Cid, error) {
+				root cid.Cid, height abi.ChainEpoch, ts *types.TipSet) (cid.Cid, error) {
 				cst := ipldcbor.NewCborStore(sm.ChainStore().Blockstore())
 
 				st, err := sm.StateTree(root)