core: use notarySet for BA module. (#153)

8 files changed, 117 insertions, 82 deletions

diff --git a/core/consensus.go b/core/consensus.go
index 08b6ceb..842d9aa 100644
--- a/core/consensus.go
+++ b/core/consensus.go
@@ -58,16 +58,16 @@ var (
 		"position of block is incorrect")
 )
 
-// consensusReceiver implements agreementReceiver.
-type consensusReceiver struct {
+// consensusBAReceiver implements agreementReceiver.
+type consensusBAReceiver struct {
 	// TODO(mission): consensus would be replaced by shard and network.
 	consensus       *Consensus
 	agreementModule *agreement
 	chainID         uint32
-	restart         chan struct{}
+	restartNotary   chan bool
 }
 
-func (recv *consensusReceiver) ProposeVote(vote *types.Vote) {
+func (recv *consensusBAReceiver) ProposeVote(vote *types.Vote) {
 	if err := recv.agreementModule.prepareVote(vote); err != nil {
 		log.Println(err)
 		return
@@ -81,7 +81,7 @@ func (recv *consensusReceiver) ProposeVote(vote *types.Vote) {
 	}()
 }
 
-func (recv *consensusReceiver) ProposeBlock() {
+func (recv *consensusBAReceiver) ProposeBlock() {
 	block := recv.consensus.proposeBlock(recv.chainID)
 	recv.consensus.baModules[recv.chainID].addCandidateBlock(block)
 	if err := recv.consensus.preProcessBlock(block); err != nil {
@@ -91,7 +91,7 @@ func (recv *consensusReceiver) ProposeBlock() {
 	recv.consensus.network.BroadcastBlock(block)
 }
 
-func (recv *consensusReceiver) ConfirmBlock(hash common.Hash) {
+func (recv *consensusBAReceiver) ConfirmBlock(hash common.Hash) {
 	block, exist := recv.consensus.baModules[recv.chainID].findCandidateBlock(hash)
 	if !exist {
 		log.Println(ErrUnknownBlockConfirmed, hash)
@@ -101,7 +101,7 @@ func (recv *consensusReceiver) ConfirmBlock(hash common.Hash) {
 		log.Println(err)
 		return
 	}
-	recv.restart <- struct{}{}
+	recv.restartNotary <- false
 }
 
 // consensusDKGReceiver implements dkgReceiver.
@@ -172,7 +172,7 @@ type Consensus struct {
 
 	// BA.
 	baModules []*agreement
-	receivers []*consensusReceiver
+	receivers []*consensusBAReceiver
 
 	// DKG.
 	dkgRunning int32
@@ -193,6 +193,7 @@ type Consensus struct {
 
 	// Misc.
 	nodeSetCache *NodeSetCache
+	round        uint64
 	lock         sync.RWMutex
 	ctx          context.Context
 	ctxCancel    context.CancelFunc
@@ -215,13 +216,13 @@ func NewConsensus(
 	nodeSetCache := NewNodeSetCache(gov)
 	crs := gov.GetCRS(round)
 	// Setup acking by information returned from Governace.
-	nodes, err := nodeSetCache.GetNodeIDs(0)
+	nodes, err := nodeSetCache.GetNodeSet(0)
 	if err != nil {
 		panic(err)
 	}
 	rb := newReliableBroadcast()
 	rb.setChainNum(config.NumChains)
-	for nID := range nodes {
+	for nID := range nodes.IDs {
 		rb.addNode(nID)
 	}
 	// Setup context.
@@ -230,7 +231,7 @@ func NewConsensus(
 	// Setup sequencer by information returned from Governace.
 	to := newTotalOrdering(
 		uint64(config.K),
-		uint64(float32(len(nodes)-1)*config.PhiRatio+1),
+		uint64(float32(len(nodes.IDs)-1)*config.PhiRatio+1),
 		config.NumChains)
 
 	ID := types.NewNodeID(prv.PublicKey())
@@ -247,7 +248,7 @@ func NewConsensus(
 		gov)
 	// Register DKG for the initial round. This is a temporary function call for
 	// simulation.
-	cfgModule.registerDKG(0, len(nodes)/3)
+	cfgModule.registerDKG(0, len(nodes.IDs)/3)
 
 	// Check if the application implement Debug interface.
 	debug, _ := app.(Debug)
@@ -272,18 +273,18 @@ func NewConsensus(
 	}
 
 	con.baModules = make([]*agreement, config.NumChains)
-	con.receivers = make([]*consensusReceiver, config.NumChains)
+	con.receivers = make([]*consensusBAReceiver, config.NumChains)
 	for i := uint32(0); i < config.NumChains; i++ {
 		chainID := i
-		recv := &consensusReceiver{
-			consensus: con,
-			chainID:   chainID,
-			restart:   make(chan struct{}, 1),
+		recv := &consensusBAReceiver{
+			consensus:     con,
+			chainID:       chainID,
+			restartNotary: make(chan bool, 1),
 		}
 		agreementModule := newAgreement(
 			con.ID,
-			con.receivers[chainID],
-			nodes,
+			recv,
+			nodes.IDs,
 			newGenesisLeaderSelector(crs),
 			con.authModule,
 		)
@@ -325,13 +326,10 @@ func (con *Consensus) Run() {
 
 func (con *Consensus) runBA(chainID uint32, tick <-chan struct{}) {
 	// TODO(jimmy-dexon): move this function inside agreement.
-	nodes, err := con.nodeSetCache.GetNodeIDs(0)
-	if err != nil {
-		panic(err)
-	}
 	agreement := con.baModules[chainID]
 	recv := con.receivers[chainID]
-	recv.restart <- struct{}{}
+	recv.restartNotary <- true
+	nIDs := make(map[types.NodeID]struct{})
 	// Reset ticker
 	<-tick
 BALoop:
@@ -345,14 +343,21 @@ BALoop:
 			<-tick
 		}
 		select {
-		case <-recv.restart:
-			// TODO(jimmy-dexon): handling change of notary set.
+		case newNotary := <-recv.restartNotary:
+			if newNotary {
+				nodes, err := con.nodeSetCache.GetNodeSet(con.round)
+				if err != nil {
+					panic(err)
+				}
+				nIDs = nodes.GetSubSet(con.gov.GetConfiguration(con.round).NumNotarySet,
+					types.NewNotarySetTarget(con.gov.GetCRS(con.round), 0, chainID))
+			}
 			aID := types.Position{
 				ShardID: 0,
 				ChainID: chainID,
 				Height:  con.rbModule.nextHeight(chainID),
 			}
-			agreement.restart(nodes, aID)
+			agreement.restart(nIDs, aID)
 		default:
 		}
 		err := agreement.nextState()
@@ -378,18 +383,17 @@ func (con *Consensus) runDKGTSIG() {
 			con.dkgReady.Broadcast()
 			con.dkgRunning = 2
 		}()
-		round := con.cfgModule.dkg.round
+		round := con.round
 		if err := con.cfgModule.runDKG(round); err != nil {
 			panic(err)
 		}
-		nodes, err := con.nodeSetCache.GetNodeIDs(0)
+		nodes, err := con.nodeSetCache.GetNodeSet(round)
 		if err != nil {
-			// TODO(mission): should be done in some bootstrap routine.
 			panic(err)
 		}
 		hash := HashConfigurationBlock(
-			nodes,
-			con.gov.GetConfiguration(0),
+			nodes.IDs,
+			con.gov.GetConfiguration(round),
 			common.Hash{},
 			con.cfgModule.prevHash)
 		psig, err := con.cfgModule.preparePartialSignature(round, hash)
diff --git a/core/interfaces.go b/core/interfaces.go
index 0319e8d..37adf36 100644
--- a/core/interfaces.go
+++ b/core/interfaces.go
@@ -97,6 +97,9 @@ type Governance interface {
 	// Return the genesis CRS if round == 0.
 	GetCRS(round uint64) []byte
 
+	// Propose a CRS of round.
+	ProposeCRS(round uint64, crs []byte)
+
 	// GetNodeSet returns the node set at a given round.
 	// Return the genesis node set if round == 0.
 	GetNodeSet(round uint64) []crypto.PublicKey
diff --git a/core/nodeset-cache.go b/core/nodeset-cache.go
index d574817..610131b 100644
--- a/core/nodeset-cache.go
+++ b/core/nodeset-cache.go
@@ -34,7 +34,7 @@ var (
 type NodeSetCache struct {
 	lock    sync.RWMutex
 	gov     Governance
-	rounds  map[uint64]map[types.NodeID]struct{}
+	rounds  map[uint64]*types.NodeSet
 	keyPool map[types.NodeID]*struct {
 		pubKey crypto.PublicKey
 		refCnt int
@@ -45,7 +45,7 @@ type NodeSetCache struct {
 func NewNodeSetCache(gov Governance) *NodeSetCache {
 	return &NodeSetCache{
 		gov:    gov,
-		rounds: make(map[uint64]map[types.NodeID]struct{}),
+		rounds: make(map[uint64]*types.NodeSet),
 		keyPool: make(map[types.NodeID]*struct {
 			pubKey crypto.PublicKey
 			refCnt int
@@ -63,7 +63,7 @@ func (cache *NodeSetCache) Exists(
 			return
 		}
 	}
-	_, exists = nIDs[nodeID]
+	_, exists = nIDs.IDs[nodeID]
 	return
 }
 
@@ -81,9 +81,9 @@ func (cache *NodeSetCache) GetPublicKey(
 	return
 }
 
-// GetNodeIDs returns IDs of nodes set of this round as map.
-func (cache *NodeSetCache) GetNodeIDs(
-	round uint64) (nIDs map[types.NodeID]struct{}, err error) {
+// GetNodeSet returns IDs of nodes set of this round as map.
+func (cache *NodeSetCache) GetNodeSet(
+	round uint64) (nIDs *types.NodeSet, err error) {
 
 	IDs, exists := cache.get(round)
 	if !exists {
@@ -91,11 +91,7 @@ func (cache *NodeSetCache) GetNodeIDs(
 			return
 		}
 	}
-	// Clone the map.
-	nIDs = make(map[types.NodeID]struct{})
-	for ID := range IDs {
-		nIDs[ID] = struct{}{}
-	}
+	nIDs = IDs.Clone()
 	return
 }
 
@@ -104,7 +100,7 @@ func (cache *NodeSetCache) GetNodeIDs(
 // This cache would maintain 10 rounds before the updated round and purge
 // rounds not in this range.
 func (cache *NodeSetCache) update(
-	round uint64) (nIDs map[types.NodeID]struct{}, err error) {
+	round uint64) (nIDs *types.NodeSet, err error) {
 
 	cache.lock.Lock()
 	defer cache.lock.Unlock()
@@ -117,10 +113,10 @@ func (cache *NodeSetCache) update(
 		return
 	}
 	// Cache new round.
-	nIDs = make(map[types.NodeID]struct{})
+	nIDs = types.NewNodeSet()
 	for _, key := range keySet {
 		nID := types.NewNodeID(key)
-		nIDs[nID] = struct{}{}
+		nIDs.Add(nID)
 		if rec, exists := cache.keyPool[nID]; exists {
 			rec.refCnt++
 		} else {
@@ -136,7 +132,7 @@ func (cache *NodeSetCache) update(
 		if round-rID <= 5 {
 			continue
 		}
-		for nID := range nIDs {
+		for nID := range nIDs.IDs {
 			rec := cache.keyPool[nID]
 			if rec.refCnt--; rec.refCnt == 0 {
 				delete(cache.keyPool, nID)
@@ -148,7 +144,7 @@ func (cache *NodeSetCache) update(
 }
 
 func (cache *NodeSetCache) get(
-	round uint64) (nIDs map[types.NodeID]struct{}, exists bool) {
+	round uint64) (nIDs *types.NodeSet, exists bool) {
 
 	cache.lock.RLock()
 	defer cache.lock.RUnlock()
diff --git a/core/nodeset-cache_test.go b/core/nodeset-cache_test.go
index e9d8867..947e239 100644
--- a/core/nodeset-cache_test.go
+++ b/core/nodeset-cache_test.go
@@ -33,6 +33,7 @@ type testGov struct {
 
 func (g *testGov) GetConfiguration(round uint64) (cfg *types.Config) { return }
 func (g *testGov) GetCRS(round uint64) (b []byte)                    { return }
+func (g *testGov) ProposeCRS(uint64, []byte)                         {}
 func (g *testGov) GetNodeSet(round uint64) []crypto.PublicKey {
 	// Randomly generating keys, and check them for verification.
 	g.curKeys = []crypto.PublicKey{}
@@ -90,20 +91,20 @@ func (s *NodeSetCacheTestSuite) TestBasicUsage() {
 	}
 
 	// Try to get round 0.
-	nodeSet0, err := cache.GetNodeIDs(0)
+	nodeSet0, err := cache.GetNodeSet(0)
 	req.NoError(err)
-	chk(cache, 0, nodeSet0)
+	chk(cache, 0, nodeSet0.IDs)
 	// Try to get round 1.
-	nodeSet1, err := cache.GetNodeIDs(1)
+	nodeSet1, err := cache.GetNodeSet(1)
 	req.NoError(err)
-	chk(cache, 0, nodeSet0)
-	chk(cache, 1, nodeSet1)
+	chk(cache, 0, nodeSet0.IDs)
+	chk(cache, 1, nodeSet1.IDs)
 	// Try to get round 6, round 0 should be purged.
-	nodeSet6, err := cache.GetNodeIDs(6)
+	nodeSet6, err := cache.GetNodeSet(6)
 	req.NoError(err)
-	chk(cache, 1, nodeSet1)
-	chk(cache, 6, nodeSet6)
-	for nID := range nodeSet0 {
+	chk(cache, 1, nodeSet1.IDs)
+	chk(cache, 6, nodeSet6.IDs)
+	for nID := range nodeSet0.IDs {
 		_, exists := cache.GetPublicKey(nID)
 		req.False(exists)
 	}
diff --git a/core/test/governance.go b/core/test/governance.go
index 69d8bdd..0fc962b 100644
--- a/core/test/governance.go
+++ b/core/test/governance.go
@@ -38,6 +38,7 @@ type Governance struct {
 	lambdaBA           time.Duration
 	lambdaDKG          time.Duration
 	privateKeys        map[types.NodeID]crypto.PrivateKey
+	crs                map[uint64][]byte
 	tsig               map[uint64]crypto.Signature
 	DKGComplaint       map[uint64][]*types.DKGComplaint
 	DKGMasterPublicKey map[uint64][]*types.DKGMasterPublicKey
@@ -54,6 +55,7 @@ func NewGovernance(nodeCount int, lambda time.Duration) (
 		lambdaBA:           lambda,
 		lambdaDKG:          lambda * 10,
 		privateKeys:        make(map[types.NodeID]crypto.PrivateKey),
+		crs:                map[uint64][]byte{0: []byte("__ DEXON")},
 		tsig:               make(map[uint64]crypto.Signature),
 		DKGComplaint:       make(map[uint64][]*types.DKGComplaint),
 		DKGMasterPublicKey: make(map[uint64][]*types.DKGMasterPublicKey),
@@ -102,7 +104,12 @@ func (g *Governance) GetConfiguration(_ uint64) *types.Config {
 
 // GetCRS returns the CRS for a given round.
 func (g *Governance) GetCRS(round uint64) []byte {
-	return []byte("__ DEXON")
+	return g.crs[round]
+}
+
+// ProposeCRS propose a CRS.
+func (g *Governance) ProposeCRS(round uint64, crs []byte) {
+	g.crs[round] = crs
 }
 
 // GetPrivateKeys return the private key for that node, this function
diff --git a/core/types/nodeset.go b/core/types/nodeset.go
index 9025cd7..e6efd76 100644
--- a/core/types/nodeset.go
+++ b/core/types/nodeset.go
@@ -27,7 +27,7 @@ import (
 
 // NodeSet is the node set structure as defined in DEXON consensus core.
 type NodeSet struct {
-	Nodes map[NodeID]struct{}
+	IDs map[NodeID]struct{}
 }
 
 // SubSetTarget is the sub set target for GetSubSet().
@@ -66,7 +66,7 @@ func (h *rankHeap) Pop() interface{} {
 // NewNodeSet creates a new NodeSet instance.
 func NewNodeSet() *NodeSet {
 	return &NodeSet{
-		Nodes: make(map[NodeID]struct{}),
+		IDs: make(map[NodeID]struct{}),
 	}
 }
 
@@ -97,11 +97,26 @@ func NewDKGSetTarget(crs []byte, round uint64) SubSetTarget {
 	return newTarget(targetDKGSet, crs, binaryRound)
 }
 
+// Add a NodeID to the set.
+func (ns *NodeSet) Add(ID NodeID) {
+	ns.IDs[ID] = struct{}{}
+}
+
+// Clone the NodeSet.
+func (ns *NodeSet) Clone() *NodeSet {
+	nsCopy := NewNodeSet()
+	for ID := range ns.IDs {
+		nsCopy.Add(ID)
+	}
+	return nsCopy
+}
+
 // GetSubSet returns the subset of given target.
-func (ns *NodeSet) GetSubSet(size int, target SubSetTarget) NodeIDs {
+func (ns *NodeSet) GetSubSet(
+	size int, target SubSetTarget) map[NodeID]struct{} {
 	h := rankHeap{}
 	idx := 0
-	for nID := range ns.Nodes {
+	for nID := range ns.IDs {
 		if idx < size {
 			h = append(h, newNodeRank(nID, target))
 		} else if idx == size {
@@ -117,9 +132,9 @@ func (ns *NodeSet) GetSubSet(size int, target SubSetTarget) NodeIDs {
 		idx++
 	}
 
-	nIDs := make(NodeIDs, 0, size)
+	nIDs := make(map[NodeID]struct{}, size)
 	for _, rank := range h {
-		nIDs = append(nIDs, rank.ID)
+		nIDs[rank.ID] = struct{}{}
 	}
 
 	return nIDs
diff --git a/core/types/nodeset_test.go b/core/types/nodeset_test.go
index 6fe9092..021a0a4 100644
--- a/core/types/nodeset_test.go
+++ b/core/types/nodeset_test.go
@@ -32,20 +32,20 @@ func (s *NodeSetTestSuite) TestGetSubSet() {
 	total := 10
 	crs := common.NewRandomHash()
 	nodes := NewNodeSet()
-	for len(nodes.Nodes) < total {
-		nodes.Nodes[NodeID{common.NewRandomHash()}] = struct{}{}
+	for len(nodes.IDs) < total {
+		nodes.IDs[NodeID{common.NewRandomHash()}] = struct{}{}
 	}
 	target := NewNotarySetTarget(crs[:], 0, 0)
-	ranks := make(map[NodeID]*nodeRank, len(nodes.Nodes))
-	for nID := range nodes.Nodes {
+	ranks := make(map[NodeID]*nodeRank, len(nodes.IDs))
+	for nID := range nodes.IDs {
 		ranks[nID] = newNodeRank(nID, target)
 	}
 	size := 4
 	notarySet := nodes.GetSubSet(size, target)
-	for _, notary := range notarySet {
+	for notary := range notarySet {
 		win := 0
 		rank := ranks[notary].rank
-		for node := range nodes.Nodes {
+		for node := range nodes.IDs {
 			if rank.Cmp(ranks[node].rank) < 0 {
 				win++
 			}
diff --git a/simulation/governance.go b/simulation/governance.go
index 3351202..9e7cbaf 100644
--- a/simulation/governance.go
+++ b/simulation/governance.go
@@ -37,7 +37,7 @@ type simGovernance struct {
 	k                  int
 	phiRatio           float32
 	chainNum           uint32
-	crs                []byte
+	crs                map[uint64][]byte
 	tsig               map[uint64]crypto.Signature
 	dkgComplaint       map[uint64][]*types.DKGComplaint
 	dkgMasterPublicKey map[uint64][]*types.DKGMasterPublicKey
@@ -52,19 +52,23 @@ func newSimGovernance(
 	id types.NodeID,
 	numNodes int, consensusConfig config.Consensus) *simGovernance {
 	return &simGovernance{
-		id:                 id,
-		nodeSet:            make(map[types.NodeID]crypto.PublicKey),
-		expectedNumNodes:   numNodes,
-		k:                  consensusConfig.K,
-		phiRatio:           consensusConfig.PhiRatio,
-		chainNum:           consensusConfig.ChainNum,
-		crs:                []byte(consensusConfig.GenesisCRS),
+		id:               id,
+		nodeSet:          make(map[types.NodeID]crypto.PublicKey),
+		expectedNumNodes: numNodes,
+		k:                consensusConfig.K,
+		phiRatio:         consensusConfig.PhiRatio,
+		chainNum:         consensusConfig.ChainNum,
+		crs: map[uint64][]byte{
+			0: []byte(consensusConfig.GenesisCRS)},
 		tsig:               make(map[uint64]crypto.Signature),
 		dkgComplaint:       make(map[uint64][]*types.DKGComplaint),
 		dkgMasterPublicKey: make(map[uint64][]*types.DKGMasterPublicKey),
-		lambdaBA:           time.Duration(consensusConfig.LambdaBA) * time.Millisecond,
-		lambdaDKG:          time.Duration(consensusConfig.LambdaDKG) * time.Millisecond,
-		roundInterval:      time.Duration(consensusConfig.RoundInterval) * time.Millisecond,
+		lambdaBA: time.Duration(consensusConfig.LambdaBA) *
+			time.Millisecond,
+		lambdaDKG: time.Duration(consensusConfig.LambdaDKG) *
+			time.Millisecond,
+		roundInterval: time.Duration(consensusConfig.RoundInterval) *
+			time.Millisecond,
 	}
 }
 
@@ -102,7 +106,12 @@ func (g *simGovernance) GetConfiguration(round uint64) *types.Config {
 
 // GetCRS returns the CRS for a given round.
 func (g *simGovernance) GetCRS(round uint64) []byte {
-	return g.crs
+	return g.crs[round]
+}
+
+// ProposeCRS proposes a CRS of round.
+func (g *simGovernance) ProposeCRS(round uint64, crs []byte) {
+	g.crs[round] = crs
 }
 
 // addNode add a new node into the simulated governance contract.