test: add test.Stopper to stop by round (#255)

2 files changed, 162 insertions, 32 deletions

diff --git a/core/test/stopper.go b/core/test/stopper.go
index 9fe5592..71b215d 100644
--- a/core/test/stopper.go
+++ b/core/test/stopper.go
@@ -40,7 +40,6 @@ func NewStopByConfirmedBlocks(
 	blockCount int,
 	apps map[types.NodeID]*App,
 	dbs map[types.NodeID]blockdb.BlockDatabase) *StopByConfirmedBlocks {
-
 	confirmedBlocks := make(map[types.NodeID]int)
 	for nID := range apps {
 		confirmedBlocks[nID] = 0
@@ -58,7 +57,6 @@ func NewStopByConfirmedBlocks(
 func (s *StopByConfirmedBlocks) ShouldStop(nID types.NodeID) bool {
 	s.lock.Lock()
 	defer s.lock.Unlock()
-
 	// Accumulate confirmed blocks proposed by this node in this round.
 	lastChecked := s.lastCheckDelivered[nID]
 	currentConfirmedBlocks := s.confirmedBlocks[nID]
@@ -84,3 +82,59 @@ func (s *StopByConfirmedBlocks) ShouldStop(nID types.NodeID) bool {
 	}
 	return true
 }
+
+// StopByRound would make sure at least one block at round R is delivered
+// at each node.
+type StopByRound struct {
+	untilRound         uint64
+	currentRounds      map[types.NodeID]uint64
+	lastCheckDelivered map[types.NodeID]int
+	apps               map[types.NodeID]*App
+	dbs                map[types.NodeID]blockdb.BlockDatabase
+	lock               sync.Mutex
+}
+
+// NewStopByRound constructs an StopByRound instance.
+func NewStopByRound(
+	round uint64,
+	apps map[types.NodeID]*App,
+	dbs map[types.NodeID]blockdb.BlockDatabase) *StopByRound {
+	return &StopByRound{
+		untilRound:         round,
+		currentRounds:      make(map[types.NodeID]uint64),
+		lastCheckDelivered: make(map[types.NodeID]int),
+		apps:               apps,
+		dbs:                dbs,
+	}
+}
+
+// ShouldStop implements Stopper interface.
+func (s *StopByRound) ShouldStop(nID types.NodeID) bool {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	// Cache latest round of this node.
+	if curRound := s.currentRounds[nID]; curRound < s.untilRound {
+		lastChecked := s.lastCheckDelivered[nID]
+		db := s.dbs[nID]
+		s.apps[nID].Check(func(app *App) {
+			for _, h := range app.DeliverSequence[lastChecked:] {
+				b, err := db.Get(h)
+				if err != nil {
+					panic(err)
+				}
+				if b.Position.Round > curRound {
+					curRound = b.Position.Round
+				}
+			}
+			s.lastCheckDelivered[nID] = len(app.DeliverSequence)
+			s.currentRounds[nID] = curRound
+		})
+	}
+	// Check if latest round on each node is later than untilRound.
+	for _, round := range s.currentRounds {
+		if round < s.untilRound {
+			return false
+		}
+	}
+	return true
+}
diff --git a/core/test/stopper_test.go b/core/test/stopper_test.go
index 7678f99..23c0137 100644
--- a/core/test/stopper_test.go
+++ b/core/test/stopper_test.go
@@ -32,48 +32,55 @@ type StopperTestSuite struct {
 	suite.Suite
 }
 
+func (s *StopperTestSuite) deliver(
+	blocks []*types.Block, app *App, db blockdb.BlockDatabase) {
+	hashes := common.Hashes{}
+	for _, b := range blocks {
+		hashes = append(hashes, b.Hash)
+		s.Require().NoError(db.Put(*b))
+	}
+	for _, h := range hashes {
+		app.StronglyAcked(h)
+	}
+	app.TotalOrderingDelivered(hashes, core.TotalOrderingModeNormal)
+	for _, h := range hashes {
+		app.BlockDelivered(h, types.FinalizationResult{
+			Timestamp: time.Time{},
+		})
+	}
+}
+
+func (s *StopperTestSuite) deliverToAllNodes(
+	blocks []*types.Block,
+	apps map[types.NodeID]*App,
+	dbs map[types.NodeID]blockdb.BlockDatabase) {
+	for nID := range apps {
+		s.deliver(blocks, apps[nID], dbs[nID])
+	}
+}
+
 func (s *StopperTestSuite) TestStopByConfirmedBlocks() {
 	// This test case makes sure this stopper would stop when
 	// all nodes confirmed at least 'x' count of blocks produced
 	// by themselves.
 	var (
-		req = s.Require()
+		req   = s.Require()
+		apps  = make(map[types.NodeID]*App)
+		dbs   = make(map[types.NodeID]blockdb.BlockDatabase)
+		nodes = GenerateRandomNodeIDs(2)
 	)
-
-	apps := make(map[types.NodeID]*App)
-	dbs := make(map[types.NodeID]blockdb.BlockDatabase)
-	nodes := GenerateRandomNodeIDs(2)
-	db, err := blockdb.NewMemBackedBlockDB()
-	req.Nil(err)
 	for _, nID := range nodes {
 		apps[nID] = NewApp()
+		db, err := blockdb.NewMemBackedBlockDB()
+		req.NoError(err)
 		dbs[nID] = db
 	}
-	deliver := func(blocks []*types.Block) {
-		hashes := common.Hashes{}
-		for _, b := range blocks {
-			hashes = append(hashes, b.Hash)
-			req.Nil(db.Put(*b))
-		}
-		for _, nID := range nodes {
-			app := apps[nID]
-			for _, h := range hashes {
-				app.StronglyAcked(h)
-			}
-			app.TotalOrderingDelivered(hashes, core.TotalOrderingModeNormal)
-			for _, h := range hashes {
-				app.BlockDelivered(h, types.FinalizationResult{
-					Timestamp: time.Time{},
-				})
-			}
-		}
-	}
 	stopper := NewStopByConfirmedBlocks(2, apps, dbs)
 	b00 := &types.Block{
 		ProposerID: nodes[0],
 		Hash:       common.NewRandomHash(),
 	}
-	deliver([]*types.Block{b00})
+	s.deliverToAllNodes([]*types.Block{b00}, apps, dbs)
 	b10 := &types.Block{
 		ProposerID: nodes[1],
 		Hash:       common.NewRandomHash(),
@@ -83,21 +90,90 @@ func (s *StopperTestSuite) TestStopByConfirmedBlocks() {
 		ParentHash: b10.Hash,
 		Hash:       common.NewRandomHash(),
 	}
-	deliver([]*types.Block{b10, b11})
+	s.deliverToAllNodes([]*types.Block{b10, b11}, apps, dbs)
 	req.False(stopper.ShouldStop(nodes[1]))
 	b12 := &types.Block{
 		ProposerID: nodes[1],
 		ParentHash: b11.Hash,
 		Hash:       common.NewRandomHash(),
 	}
-	deliver([]*types.Block{b12})
+	s.deliverToAllNodes([]*types.Block{b12}, apps, dbs)
 	req.False(stopper.ShouldStop(nodes[1]))
 	b01 := &types.Block{
 		ProposerID: nodes[0],
 		ParentHash: b00.Hash,
 		Hash:       common.NewRandomHash(),
 	}
-	deliver([]*types.Block{b01})
+	s.deliverToAllNodes([]*types.Block{b01}, apps, dbs)
+	req.True(stopper.ShouldStop(nodes[0]))
+}
+
+func (s *StopperTestSuite) TestStopByRound() {
+	// This test case make sure at least one block from round R
+	// is delivered by each node.
+	var (
+		req   = s.Require()
+		apps  = make(map[types.NodeID]*App)
+		dbs   = make(map[types.NodeID]blockdb.BlockDatabase)
+		nodes = GenerateRandomNodeIDs(2)
+	)
+	for _, nID := range nodes {
+		apps[nID] = NewApp()
+		db, err := blockdb.NewMemBackedBlockDB()
+		req.NoError(err)
+		dbs[nID] = db
+	}
+	stopper := NewStopByRound(10, apps, dbs)
+	b00 := &types.Block{
+		ProposerID: nodes[0],
+		Position: types.Position{
+			Round:   0,
+			ChainID: 0,
+			Height:  0,
+		},
+		Hash: common.NewRandomHash(),
+	}
+	s.deliverToAllNodes([]*types.Block{b00}, apps, dbs)
+	b10 := &types.Block{
+		ProposerID: nodes[1],
+		Position: types.Position{
+			Round:   0,
+			ChainID: 1,
+			Height:  0,
+		},
+		Hash: common.NewRandomHash(),
+	}
+	b11 := &types.Block{
+		ProposerID: nodes[1],
+		ParentHash: b10.Hash,
+		Position: types.Position{
+			Round:   0,
+			ChainID: 1,
+			Height:  1,
+		},
+		Hash: common.NewRandomHash(),
+	}
+	s.deliverToAllNodes([]*types.Block{b10, b11}, apps, dbs)
+	req.False(stopper.ShouldStop(nodes[0]))
+	req.False(stopper.ShouldStop(nodes[1]))
+	// Deliver one block at round 10 to node0
+	b12 := &types.Block{
+		ProposerID: nodes[1],
+		ParentHash: b11.Hash,
+		Position: types.Position{
+			Round:   10,
+			ChainID: 1,
+			Height:  2,
+		},
+		Hash: common.NewRandomHash(),
+	}
+	// None should stop when only one node reach that round.
+	s.deliver([]*types.Block{b12}, apps[nodes[0]], dbs[nodes[0]])
+	req.False(stopper.ShouldStop(nodes[0]))
+	req.False(stopper.ShouldStop(nodes[1]))
+	// Everyone should stop now.
+	s.deliver([]*types.Block{b12}, apps[nodes[1]], dbs[nodes[1]])
+	req.True(stopper.ShouldStop(nodes[1]))
 	req.True(stopper.ShouldStop(nodes[0]))
 }