make the upgrade schedule fully configurable

This should help with testing.

chain/stmgr/forks.go

@@ -9,6 +9,7 @@ import (
 	"github.com/filecoin-project/go-address"
 	"github.com/filecoin-project/go-state-types/abi"
 	"github.com/filecoin-project/go-state-types/big"
+	"github.com/filecoin-project/go-state-types/network"
 	"github.com/ipfs/go-cid"
 	cbor "github.com/ipfs/go-ipld-cbor"
 	"golang.org/x/xerrors"
@@ -33,17 +34,27 @@ import (
 	"github.com/filecoin-project/lotus/chain/vm"
 )
 
-var ForksAtHeight = map[abi.ChainEpoch]func(context.Context, *StateManager, ExecCallback, cid.Cid, *types.TipSet) (cid.Cid, error){
-	build.UpgradeBreezeHeight:   UpgradeFaucetBurnRecovery,
-	build.UpgradeIgnitionHeight: UpgradeIgnition,
-	build.UpgradeActorsV2Height: UpgradeActorsV2,
-	build.UpgradeLiftoffHeight:  UpgradeLiftoff,
+type UpgradeFunc func(context.Context, *StateManager, ExecCallback, cid.Cid, *types.TipSet) (cid.Cid, error)
+
+type Upgrade struct {
+	Height    abi.ChainEpoch
+	Migration UpgradeFunc
+}
+
+type UpgradeSchedule map[network.Version]Upgrade
+
+var DefaultUpgradeSchedule = UpgradeSchedule{
+	network.Version1: {build.UpgradeBreezeHeight, UpgradeFaucetBurnRecovery},
+	network.Version2: {build.UpgradeSmokeHeight, nil},
+	network.Version3: {build.UpgradeIgnitionHeight, UpgradeIgnition},
+	network.Version4: {build.UpgradeActorsV2Height, UpgradeActorsV2},
+	network.Version5: {build.UpgradeLiftoffHeight, UpgradeLiftoff},
 }
 
 func (sm *StateManager) handleStateForks(ctx context.Context, root cid.Cid, height abi.ChainEpoch, cb ExecCallback, ts *types.TipSet) (cid.Cid, error) {
 	retCid := root
 	var err error
-	f, ok := ForksAtHeight[height]
+	f, ok := sm.stateMigrations[height]
 	if ok {
 		retCid, err = f(ctx, sm, cb, root, ts)
 		if err != nil {

chain/stmgr/forks_test.go

@@ -19,7 +19,6 @@ import (
 	lotusinit "github.com/filecoin-project/lotus/chain/actors/builtin/init"
 	"github.com/filecoin-project/lotus/chain/actors/policy"
 	"github.com/filecoin-project/lotus/chain/gen"
-	"github.com/filecoin-project/lotus/chain/stmgr"
 	. "github.com/filecoin-project/lotus/chain/stmgr"
 	"github.com/filecoin-project/lotus/chain/types"
 	"github.com/filecoin-project/lotus/chain/vm"
@@ -110,51 +109,54 @@ func TestForkHeightTriggers(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	sm := NewStateManager(cg.ChainStore())
-
-	inv := vm.NewActorRegistry()
-
 	// predicting the address here... may break if other assumptions change
 	taddr, err := address.NewIDAddress(1002)
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	stmgr.ForksAtHeight[testForkHeight] = func(ctx context.Context, sm *StateManager, cb ExecCallback, root cid.Cid, ts *types.TipSet) (cid.Cid, error) {
-		cst := ipldcbor.NewCborStore(sm.ChainStore().Blockstore())
+	sm := NewStateManager(cg.ChainStore(),
+		WithUpgradeSchedule(UpgradeSchedule{
+			1: {Height: testForkHeight, Migration: func(ctx context.Context, sm *StateManager, cb ExecCallback,
+				root cid.Cid, ts *types.TipSet) (cid.Cid, error) {
+				cst := ipldcbor.NewCborStore(sm.ChainStore().Blockstore())
 
-		st, err := sm.StateTree(root)
-		if err != nil {
-			return cid.Undef, xerrors.Errorf("getting state tree: %w", err)
-		}
+				st, err := sm.StateTree(root)
+				if err != nil {
+					return cid.Undef, xerrors.Errorf("getting state tree: %w", err)
+				}
 
-		act, err := st.GetActor(taddr)
-		if err != nil {
-			return cid.Undef, err
-		}
+				act, err := st.GetActor(taddr)
+				if err != nil {
+					return cid.Undef, 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)
-		}
+				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)
+				}
 
-		tas.HasUpgraded = 55
+				tas.HasUpgraded = 55
 
-		ns, err := cst.Put(ctx, &tas)
-		if err != nil {
-			return cid.Undef, err
-		}
+				ns, err := cst.Put(ctx, &tas)
+				if err != nil {
+					return cid.Undef, err
+				}
 
-		act.Head = ns
+				act.Head = ns
 
-		if err := st.SetActor(taddr, act); err != nil {
-			return cid.Undef, err
-		}
+				if err := st.SetActor(taddr, act); err != nil {
+					return cid.Undef, err
+				}
 
-		return st.Flush(ctx)
-	}
+				return st.Flush(ctx)
+			}},
+		}),
+	)
 
+	inv := vm.NewActorRegistry()
 	inv.Register(nil, testActor{})
+
 	sm.SetVMConstructor(func(ctx context.Context, vmopt *vm.VMOpts) (*vm.VM, error) {
 		nvm, err := vm.NewVM(ctx, vmopt)
 		if err != nil {

chain/stmgr/options.go

@@ -0,0 +1,23 @@
+package stmgr
+
+type Option func(*config)
+
+type config struct {
+	upgradeSchedule UpgradeSchedule
+}
+
+func parseOptions(opts ...Option) *config {
+	cfg := &config{
+		upgradeSchedule: DefaultUpgradeSchedule,
+	}
+	for _, opt := range opts {
+		opt(cfg)
+	}
+	return cfg
+}
+
+func WithUpgradeSchedule(schedule UpgradeSchedule) Option {
+	return func(cfg *config) {
+		cfg.upgradeSchedule = schedule
+	}
+}

chain/stmgr/stmgr.go

@@ -42,6 +42,12 @@ var log = logging.Logger("statemgr")
 type StateManager struct {
 	cs *store.ChainStore
 
+	// Sorted network upgrade epochs (starting at version 1).
+	// -1 means the version is skipped.
+	networkVersions []abi.ChainEpoch
+	// Maps chain epochs to upgrade functions.
+	stateMigrations map[abi.ChainEpoch]UpgradeFunc
+
 	stCache              map[string][]cid.Cid
 	compWait             map[string]chan struct{}
 	stlk                 sync.Mutex
@@ -51,12 +57,36 @@ type StateManager struct {
 	postIgnitionGenInfos *genesisInfo
 }
 
-func NewStateManager(cs *store.ChainStore) *StateManager {
+func NewStateManager(cs *store.ChainStore, options ...Option) *StateManager {
+	cfg := parseOptions(options...)
+	stateMigrations := make(map[abi.ChainEpoch]UpgradeFunc, len(cfg.upgradeSchedule))
+	networkVersions := make([]abi.ChainEpoch, 0, len(stateMigrations))
+
+	// Iterate version by version, to make sure we handle skipped version numbers.
+	// Always skip version 0.
+	for i, version := 0, network.Version(1); i < len(cfg.upgradeSchedule); version++ {
+		upgrade, ok := cfg.upgradeSchedule[version]
+		if ok {
+			// We've processed an upgrade.
+			i++
+		}
+
+		epoch := abi.ChainEpoch(-1)
+		if ok && upgrade.Height >= 0 {
+			epoch = upgrade.Height
+			stateMigrations[epoch] = upgrade.Migration
+		}
+
+		networkVersions = append(networkVersions, epoch)
+	}
+
 	return &StateManager{
-		newVM:    vm.NewVM,
-		cs:       cs,
-		stCache:  make(map[string][]cid.Cid),
-		compWait: make(map[string]chan struct{}),
+		networkVersions: networkVersions,
+		stateMigrations: stateMigrations,
+		newVM:           vm.NewVM,
+		cs:              cs,
+		stCache:         make(map[string][]cid.Cid),
+		compWait:        make(map[string]chan struct{}),
 	}
 }
 
@@ -1259,29 +1289,12 @@ func (sm *StateManager) GetCirculatingSupply(ctx context.Context, height abi.Cha
 }
 
 func (sm *StateManager) GetNtwkVersion(ctx context.Context, height abi.ChainEpoch) network.Version {
-	// TODO: move hard fork epoch checks to a schedule defined in build/
-
-	if build.UseNewestNetwork() {
-		return build.NewestNetworkVersion
+	for v, epoch := range sm.networkVersions {
+		if epoch >= 0 && height <= epoch {
+			return network.Version(v + 1) // we've skipped version 0
+		}
 	}
-
-	if height <= build.UpgradeBreezeHeight {
-		return network.Version0
-	}
-
-	if height <= build.UpgradeSmokeHeight {
-		return network.Version1
-	}
-
-	if height <= build.UpgradeIgnitionHeight {
-		return network.Version2
-	}
-
-	if height <= build.UpgradeActorsV2Height {
-		return network.Version3
-	}
-
-	return build.NewestNetworkVersion
+	return network.Version(len(sm.networkVersions))
 }
 
 func (sm *StateManager) GetPaychState(ctx context.Context, addr address.Address, ts *types.TipSet) (*types.Actor, paych.State, error) {

go.mod

@@ -33,7 +33,7 @@ require (
 	github.com/filecoin-project/go-multistore v0.0.3
 	github.com/filecoin-project/go-padreader v0.0.0-20200903213702-ed5fae088b20
 	github.com/filecoin-project/go-paramfetch v0.0.2-0.20200701152213-3e0f0afdc261
-	github.com/filecoin-project/go-state-types v0.0.0-20200928172055-2df22083d8ab
+	github.com/filecoin-project/go-state-types v0.0.0-20201001162932-93a412d8e9fa
 	github.com/filecoin-project/go-statemachine v0.0.0-20200925024713-05bd7c71fbfe
 	github.com/filecoin-project/go-statestore v0.1.0
 	github.com/filecoin-project/go-storedcounter v0.0.0-20200421200003-1c99c62e8a5b

go.sum

@@ -263,6 +263,8 @@ github.com/filecoin-project/go-state-types v0.0.0-20200904021452-1883f36ca2f4/go
 github.com/filecoin-project/go-state-types v0.0.0-20200905071437-95828685f9df/go.mod h1:IQ0MBPnonv35CJHtWSN3YY1Hz2gkPru1Q9qoaYLxx9I=
 github.com/filecoin-project/go-state-types v0.0.0-20200928172055-2df22083d8ab h1:cEDC5Ei8UuT99hPWhCjA72SM9AuRtnpvdSTIYbnzN8I=
 github.com/filecoin-project/go-state-types v0.0.0-20200928172055-2df22083d8ab/go.mod h1:ezYnPf0bNkTsDibL/psSz5dy4B5awOJ/E7P2Saeep8g=
+github.com/filecoin-project/go-state-types v0.0.0-20201001162932-93a412d8e9fa h1:WVxLrmc6w5naeX3QikfRvc1G7EFfBZbhXuzM8Lj2J+o=
+github.com/filecoin-project/go-state-types v0.0.0-20201001162932-93a412d8e9fa/go.mod h1:ezYnPf0bNkTsDibL/psSz5dy4B5awOJ/E7P2Saeep8g=
 github.com/filecoin-project/go-statemachine v0.0.0-20200714194326-a77c3ae20989/go.mod h1:FGwQgZAt2Gh5mjlwJUlVB62JeYdo+if0xWxSEfBD9ig=
 github.com/filecoin-project/go-statemachine v0.0.0-20200925024713-05bd7c71fbfe h1:dF8u+LEWeIcTcfUcCf3WFVlc81Fr2JKg8zPzIbBDKDw=
 github.com/filecoin-project/go-statemachine v0.0.0-20200925024713-05bd7c71fbfe/go.mod h1:FGwQgZAt2Gh5mjlwJUlVB62JeYdo+if0xWxSEfBD9ig=