core: optimize core (#428)

* core: Use a channel to process ba confirmed block

* change the implementation of done() to react faster

* Fix restart

* Wait tipRound to change

* fix corner case

* Check for context

1 files changed, 94 insertions, 43 deletions

diff --git a/core/agreement-mgr.go b/core/agreement-mgr.go
index 9e86369..7410977 100644
--- a/core/agreement-mgr.go
+++ b/core/agreement-mgr.go
@@ -419,7 +419,10 @@ Loop:
 		// Run BA for this round.
 		recv.roundValue.Store(currentRound)
 		recv.changeNotaryTime = roundEndTime
-		recv.restartNotary <- types.Position{ChainID: math.MaxUint32}
+		recv.restartNotary <- types.Position{
+			Round:   setting.recv.round(),
+			ChainID: math.MaxUint32,
+		}
 		if err := mgr.baRoutineForOneRound(&setting); err != nil {
 			mgr.logger.Error("BA routine failed",
 				"error", err,
@@ -435,6 +438,79 @@ func (mgr *agreementMgr) baRoutineForOneRound(
 	agr := setting.agr
 	recv := setting.recv
 	oldPos := agr.agreementID()
+	restart := func(restartPos types.Position) (breakLoop bool, err error) {
+		if !isStop(restartPos) {
+			if restartPos.Round > oldPos.Round {
+				for {
+					select {
+					case <-mgr.ctx.Done():
+						break
+					default:
+					}
+					tipRound := mgr.lattice.TipRound(setting.chainID)
+					if tipRound > restartPos.Round {
+						// It's a vary rare that this go routine sleeps for entire round.
+						break
+					} else if tipRound != restartPos.Round {
+						mgr.logger.Debug("Waiting lattice to change round...",
+							"pos", &restartPos)
+					} else {
+						break
+					}
+					time.Sleep(100 * time.Millisecond)
+				}
+				// This round is finished.
+				breakLoop = true
+				return
+			}
+			if restartPos.Older(&oldPos) {
+				// The restartNotary event is triggered by 'BlockConfirmed'
+				// of some older block.
+				return
+			}
+		}
+		var nextHeight uint64
+		var nextTime time.Time
+		for {
+			nextHeight, nextTime, err =
+				mgr.lattice.NextBlock(recv.round(), setting.chainID)
+			if err != nil {
+				mgr.logger.Debug("Error getting next height",
+					"error", err,
+					"round", recv.round(),
+					"chainID", setting.chainID)
+				err = nil
+				nextHeight = restartPos.Height
+			}
+			if isStop(oldPos) && nextHeight == 0 {
+				break
+			}
+			if isStop(restartPos) && nextHeight == 0 {
+				break
+			}
+			if nextHeight > restartPos.Height {
+				break
+			}
+			mgr.logger.Debug("Lattice not ready!!!",
+				"old", &oldPos, "restart", &restartPos, "next", nextHeight)
+			time.Sleep(100 * time.Millisecond)
+		}
+		nextPos := types.Position{
+			Round:   recv.round(),
+			ChainID: setting.chainID,
+			Height:  nextHeight,
+		}
+		oldPos = nextPos
+		var leader types.NodeID
+		leader, err = mgr.cache.GetLeaderNode(nextPos)
+		if err != nil {
+			return
+		}
+		time.Sleep(nextTime.Sub(time.Now()))
+		setting.ticker.Restart()
+		agr.restart(setting.notarySet, nextPos, leader, setting.crs)
+		return
+	}
 Loop:
 	for {
 		select {
@@ -442,55 +518,30 @@ Loop:
 			break Loop
 		default:
 		}
-		select {
-		case restartPos := <-recv.restartNotary:
-			if !isStop(restartPos) {
-				if restartPos.Round > oldPos.Round {
-					// This round is finished.
-					break Loop
-				}
-				if restartPos.Older(&oldPos) {
-					// The restartNotary event is triggered by 'BlockConfirmed'
-					// of some older block.
-					break
-				}
-			}
-			var nextHeight uint64
-			var nextTime time.Time
-			for {
-				nextHeight, nextTime, err =
-					mgr.lattice.NextBlock(recv.round(), setting.chainID)
+		if agr.confirmed() {
+			// Block until receive restartPos
+			select {
+			case restartPos := <-recv.restartNotary:
+				breakLoop, err := restart(restartPos)
 				if err != nil {
-					mgr.logger.Debug("Error getting next height",
-						"error", err,
-						"round", recv.round(),
-						"chainID", setting.chainID)
-					err = nil
-					nextHeight = restartPos.Height
-				}
-				if isStop(restartPos) || nextHeight == 0 {
-					break
+					return err
 				}
-				if nextHeight > restartPos.Height {
-					break
+				if breakLoop {
+					break Loop
 				}
-				mgr.logger.Debug("Lattice not ready!!!",
-					"old", &restartPos, "next", nextHeight)
-				time.Sleep(100 * time.Millisecond)
-			}
-			nextPos := types.Position{
-				Round:   recv.round(),
-				ChainID: setting.chainID,
-				Height:  nextHeight,
+			case <-mgr.ctx.Done():
+				break Loop
 			}
-			oldPos = nextPos
-			leader, err := mgr.cache.GetLeaderNode(nextPos)
+		}
+		select {
+		case restartPos := <-recv.restartNotary:
+			breakLoop, err := restart(restartPos)
 			if err != nil {
 				return err
 			}
-			time.Sleep(nextTime.Sub(time.Now()))
-			setting.ticker.Restart()
-			agr.restart(setting.notarySet, nextPos, leader, setting.crs)
+			if breakLoop {
+				break Loop
+			}
 		default:
 		}
 		if agr.pullVotes() {