From 509abd7f8eec95920b4c5f711728573c4e59e7ef Mon Sep 17 00:00:00 2001
From: rickyyangz <yangzhongzhong@ruiquewenhua.com>
Date: Thu, 4 Apr 2019 11:05:11 +0800
Subject: [PATCH] upgrade logic for splitting val address

---
 app/fee_distribution.go   |  8 +++++++-
 common/upgrade/upgrade.go | 20 ++------------------
 2 files changed, 9 insertions(+), 19 deletions(-)

diff --git a/app/fee_distribution.go b/app/fee_distribution.go
index b6d221e33..b0c76dd70 100644
--- a/app/fee_distribution.go
+++ b/app/fee_distribution.go
@@ -8,6 +8,7 @@ import (
 	"github.com/binance-chain/node/common/fees"
 	"github.com/binance-chain/node/common/log"
 	"github.com/binance-chain/node/common/types"
+	"github.com/binance-chain/node/common/upgrade"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/x/auth"
@@ -15,10 +16,15 @@ import (
 )
 
 func NewValAddrCache(stakeKeeper stake.Keeper) *ValAddrCache {
-	return &ValAddrCache{
+	cache := &ValAddrCache{
 		cache:       make(map[string]sdk.AccAddress),
 		stakeKeeper: stakeKeeper,
 	}
+
+	upgrade.Mgr.RegisterBeginBlocker(sdk.UpgradeSeparateValAddrName, func(ctx sdk.Context) {
+		cache.ClearCache()
+	})
+	return cache
 }
 
 type ValAddrCache struct {
diff --git a/common/upgrade/upgrade.go b/common/upgrade/upgrade.go
index 6c2c0bd77..44a892831 100644
--- a/common/upgrade/upgrade.go
+++ b/common/upgrade/upgrade.go
@@ -10,26 +10,10 @@ var Mgr = sdk.UpgradeMgr
 const FixOrderSeqInPriceLevelName = "fixOrderSeqInPriceLevel"
 const FixDropFilledOrderSeqName = "fixDropFilledOrderSeq"
 
-func Upgrade(name string, before func(), in func(), after func()) {
-	if sdk.IsUpgradeHeight(name) {
-		if in != nil {
-			in()
-		}
-	} else if sdk.IsUpgrade(name) {
-		if after != nil {
-			after()
-		}
-	} else {
-		if before != nil {
-			before()
-		}
-	}
-}
-
 func FixOrderSeqInPriceLevel(before func(), in func(), after func()) {
-	Upgrade(FixOrderSeqInPriceLevelName, before, in, after)
+	sdk.Upgrade(FixOrderSeqInPriceLevelName, before, in, after)
 }
 
 func FixDropFilledOrderSeq(before func(), after func()) {
-	Upgrade(FixDropFilledOrderSeqName, before, nil, after)
+	sdk.Upgrade(FixDropFilledOrderSeqName, before, nil, after)
 }