From e865ba0c645ada6b026cdc2263b24461befc32bd Mon Sep 17 00:00:00 2001
From: Steven Allen <steven@stebalien.com>
Date: Thu, 1 Oct 2020 15:18:59 -0700
Subject: [PATCH] improve upgrade schedule configurable

Unfortunately, it looks like functional ops cannot be passed in via lotus's
di system.
---
 chain/stmgr/forks_test.go |  8 ++++----
 chain/stmgr/options.go    | 23 -----------------------
 chain/stmgr/stmgr.go      | 13 ++++++++-----
 node/builder.go           |  3 ++-
 4 files changed, 14 insertions(+), 33 deletions(-)
 delete mode 100644 chain/stmgr/options.go

diff --git a/chain/stmgr/forks_test.go b/chain/stmgr/forks_test.go
index 1b794f31c..db65d2485 100644
--- a/chain/stmgr/forks_test.go
+++ b/chain/stmgr/forks_test.go
@@ -115,8 +115,9 @@ func TestForkHeightTriggers(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	sm := NewStateManager(cg.ChainStore(),
-		WithUpgradeSchedule(UpgradeSchedule{
+	sm := NewStateManagerWithUpgradeSchedule(
+		cg.ChainStore(),
+		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())
@@ -151,8 +152,7 @@ func TestForkHeightTriggers(t *testing.T) {
 
 				return st.Flush(ctx)
 			}},
-		}),
-	)
+		})
 
 	inv := vm.NewActorRegistry()
 	inv.Register(nil, testActor{})
diff --git a/chain/stmgr/options.go b/chain/stmgr/options.go
deleted file mode 100644
index 8c63149a1..000000000
--- a/chain/stmgr/options.go
+++ /dev/null
@@ -1,23 +0,0 @@
-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
-	}
-}
diff --git a/chain/stmgr/stmgr.go b/chain/stmgr/stmgr.go
index 8560690ea..966049693 100644
--- a/chain/stmgr/stmgr.go
+++ b/chain/stmgr/stmgr.go
@@ -57,15 +57,18 @@ type StateManager struct {
 	postIgnitionGenInfos *genesisInfo
 }
 
-func NewStateManager(cs *store.ChainStore, options ...Option) *StateManager {
-	cfg := parseOptions(options...)
-	stateMigrations := make(map[abi.ChainEpoch]UpgradeFunc, len(cfg.upgradeSchedule))
+func NewStateManager(cs *store.ChainStore) *StateManager {
+	return NewStateManagerWithUpgradeSchedule(cs, DefaultUpgradeSchedule)
+}
+
+func NewStateManagerWithUpgradeSchedule(cs *store.ChainStore, us UpgradeSchedule) *StateManager {
+	stateMigrations := make(map[abi.ChainEpoch]UpgradeFunc, len(us))
 	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]
+	for i, version := 0, network.Version(1); i < len(us); version++ {
+		upgrade, ok := us[version]
 		if ok {
 			// We've processed an upgrade.
 			i++
diff --git a/node/builder.go b/node/builder.go
index 8512d0a0c..c738740ec 100644
--- a/node/builder.go
+++ b/node/builder.go
@@ -259,7 +259,8 @@ func Online() Option {
 			Override(new(ffiwrapper.Verifier), ffiwrapper.ProofVerifier),
 			Override(new(vm.SyscallBuilder), vm.Syscalls),
 			Override(new(*store.ChainStore), modules.ChainStore),
-			Override(new(*stmgr.StateManager), stmgr.NewStateManager),
+			Override(new(stmgr.UpgradeSchedule), stmgr.DefaultUpgradeSchedule),
+			Override(new(*stmgr.StateManager), stmgr.NewStateManagerWithUpgradeSchedule),
 			Override(new(*wallet.Wallet), wallet.NewWallet),
 			Override(new(*messagesigner.MessageSigner), messagesigner.NewMessageSigner),