core/crypto/dkg: concurrent access for gpk (#548)

* core/crypto/dkg: add benchmark
* core/crypto/dkg: concurrent access for gpk

14 files changed, 133 insertions, 44 deletions

diff --git a/core/crypto/dkg/dkg.go b/core/crypto/dkg/dkg.go
index 796609d..b89ad10 100644
--- a/core/crypto/dkg/dkg.go
+++ b/core/crypto/dkg/dkg.go
@@ -21,6 +21,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"sync"
+	"sync/atomic"
 
 	"github.com/dexon-foundation/bls/ffi/go/bls"
 	"github.com/dexon-foundation/dexon/rlp"
@@ -210,23 +212,30 @@ func (prvs *PrivateKeyShares) DecodeRLP(s *rlp.Stream) error {
 	return nil
 }
 
+type publicKeySharesCache struct {
+	share []PublicKey
+	index map[ID]int
+}
+
 // PublicKeyShares represents a public key shares for DKG protocol.
 type PublicKeyShares struct {
-	shareCaches     []PublicKey
-	shareCacheIndex map[ID]int
+	cache           atomic.Value
+	lock            sync.Mutex
 	masterPublicKey []bls.PublicKey
 }
 
 // Equal checks equality of two PublicKeyShares instance.
 func (pubs *PublicKeyShares) Equal(other *PublicKeyShares) bool {
+	cache := pubs.cache.Load().(*publicKeySharesCache)
+	cacheOther := other.cache.Load().(*publicKeySharesCache)
 	// Check shares.
-	for dID, idx := range pubs.shareCacheIndex {
-		otherIdx, exists := other.shareCacheIndex[dID]
+	for dID, idx := range cache.index {
+		otherIdx, exists := cacheOther.index[dID]
 		if !exists {
 			continue
 		}
-		if !pubs.shareCaches[idx].publicKey.IsEqual(
-			&other.shareCaches[otherIdx].publicKey) {
+		if !cache.share[idx].publicKey.IsEqual(
+			&cacheOther.share[otherIdx].publicKey) {
 			return false
 		}
 	}
@@ -267,7 +276,7 @@ func (pubs *PublicKeyShares) DecodeRLP(s *rlp.Stream) error {
 		ps.masterPublicKey = append(ps.masterPublicKey, key)
 	}
 
-	*pubs = *ps
+	*pubs = *ps.Move()
 	return nil
 }
 
@@ -349,13 +358,12 @@ func NewPrivateKeyShares(t int) (*PrivateKeyShares, *PublicKeyShares) {
 	prv.SetByCSPRNG()
 	msk := prv.GetMasterSecretKey(t)
 	mpk := bls.GetMasterPublicKey(msk)
+	pubShare := NewEmptyPublicKeyShares()
+	pubShare.masterPublicKey = mpk
 	return &PrivateKeyShares{
-			masterPrivateKey: msk,
-			shareIndex:       make(map[ID]int),
-		}, &PublicKeyShares{
-			shareCacheIndex: make(map[ID]int),
-			masterPublicKey: mpk,
-		}
+		masterPrivateKey: msk,
+		shareIndex:       make(map[ID]int),
+	}, pubShare
 }
 
 // NewEmptyPrivateKeyShares creates an empty private key shares.
@@ -442,16 +450,25 @@ func (prvs *PrivateKeyShares) Share(ID ID) (*PrivateKey, bool) {
 
 // NewEmptyPublicKeyShares creates an empty public key shares.
 func NewEmptyPublicKeyShares() *PublicKeyShares {
-	return &PublicKeyShares{
-		shareCacheIndex: make(map[ID]int),
+	cache := &publicKeySharesCache{
+		index: make(map[ID]int),
 	}
+	pubShares := &PublicKeyShares{}
+	pubShares.cache.Store(cache)
+	return pubShares
+}
+
+// Move will invalidate itself. Do not access to original reference.
+func (pubs *PublicKeyShares) Move() *PublicKeyShares {
+	return pubs
 }
 
 // Share returns the share for the ID.
 func (pubs *PublicKeyShares) Share(ID ID) (*PublicKey, error) {
-	idx, exist := pubs.shareCacheIndex[ID]
+	cache := pubs.cache.Load().(*publicKeySharesCache)
+	idx, exist := cache.index[ID]
 	if exist {
-		return &pubs.shareCaches[idx], nil
+		return &cache.share[idx], nil
 	}
 	var pk PublicKey
 	if err := pk.publicKey.Set(pubs.masterPublicKey, &ID); err != nil {
@@ -465,14 +482,19 @@ func (pubs *PublicKeyShares) Share(ID ID) (*PublicKey, error) {
 
 // AddShare adds a share.
 func (pubs *PublicKeyShares) AddShare(ID ID, share *PublicKey) error {
-	if idx, exist := pubs.shareCacheIndex[ID]; exist {
-		if !share.publicKey.IsEqual(&pubs.shareCaches[idx].publicKey) {
+	cache := pubs.cache.Load().(*publicKeySharesCache)
+	if idx, exist := cache.index[ID]; exist {
+		if !share.publicKey.IsEqual(&cache.share[idx].publicKey) {
 			return ErrDuplicatedShare
 		}
 		return nil
 	}
-	pubs.shareCacheIndex[ID] = len(pubs.shareCaches)
-	pubs.shareCaches = append(pubs.shareCaches, *share)
+	pubs.lock.Lock()
+	defer pubs.lock.Unlock()
+	cache = pubs.cache.Load().(*publicKeySharesCache)
+	cache.index[ID] = len(cache.share)
+	cache.share = append(cache.share, *share)
+	pubs.cache.Store(cache)
 	return nil
 }
 
diff --git a/core/crypto/dkg/dkg_test.go b/core/crypto/dkg/dkg_test.go
index f278932..a679cf4 100644
--- a/core/crypto/dkg/dkg_test.go
+++ b/core/crypto/dkg/dkg_test.go
@@ -307,11 +307,10 @@ func (s *DKGTestSuite) TestPrivateKeyRLPEncodeDecode() {
 
 func (s *DKGTestSuite) TestPublicKeySharesRLPEncodeDecode() {
 	p := NewEmptyPublicKeyShares()
-	for i, id := range s.genID(1) {
+	for _, id := range s.genID(1) {
 		privkey := NewPrivateKey()
 		pubkey := privkey.PublicKey().(PublicKey)
-		p.shareCaches = append(p.shareCaches, pubkey)
-		p.shareCacheIndex[id] = i
+		p.AddShare(id, &pubkey)
 		p.masterPublicKey = append(p.masterPublicKey, pubkey.publicKey)
 	}
 
@@ -371,6 +370,30 @@ func (s *DKGTestSuite) TestPublicKeySharesEquality() {
 	req.True(pubShares2.Equal(pubShares1))
 }
 
+func (s *DKGTestSuite) TestPublicKeySharesMove() {
+	var req = s.Require()
+	IDs := s.genID(2)
+	_, pubShares1 := NewPrivateKeyShares(4)
+	// Make a copy from an empty share.
+	pubShares2 := pubShares1.Clone()
+	req.True(pubShares1.Equal(pubShares2))
+	// Move from pubShare1.
+	pubShares3 := pubShares1.Move()
+	// Add two shares.
+	prvKey1 := NewPrivateKey()
+	pubKey1 := prvKey1.PublicKey().(PublicKey)
+	req.NoError(pubShares3.AddShare(IDs[0], &pubKey1))
+	prvKey2 := NewPrivateKey()
+	pubKey2 := prvKey2.PublicKey().(PublicKey)
+	req.True(pubShares3.Equal(pubShares2))
+	// Clone the shares.
+	req.NoError(pubShares2.AddShare(IDs[0], &pubKey1))
+	req.NoError(pubShares2.AddShare(IDs[1], &pubKey2))
+	// They should be equal now.
+	req.True(pubShares3.Equal(pubShares2))
+	req.True(pubShares2.Equal(pubShares3))
+}
+
 func (s *DKGTestSuite) TestPrivateKeySharesEquality() {
 	var req = s.Require()
 	IDs := s.genID(2)
@@ -578,3 +601,47 @@ func BenchmarkDKGProtocol(b *testing.B) {
 		}
 	})
 }
+
+func BenchmarkGPKShare81_121(b *testing.B) { benchmarkGPKShare(b, 81, 121) }
+
+func benchmarkGPKShare(b *testing.B, t, n int) {
+	_, pubShare := NewPrivateKeyShares(t)
+	IDs := make(IDs, n)
+	for i := range IDs {
+		id := common.NewRandomHash()
+		IDs[i] = NewID(id[:])
+	}
+
+	for _, id := range IDs {
+		_, err := pubShare.Share(id)
+		if err != nil {
+			panic(err)
+		}
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		for _, id := range IDs {
+			pubShare.Share(id)
+		}
+	}
+}
+
+func BenchmarkGPKAddShare81_121(b *testing.B) { benchmarkGPKAddShare(b, 81, 121) }
+
+func benchmarkGPKAddShare(b *testing.B, t, n int) {
+	IDs := make(IDs, n)
+	for i := range IDs {
+		id := common.NewRandomHash()
+		IDs[i] = NewID(id[:])
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		_, pubShare := NewPrivateKeyShares(t)
+		b.StartTimer()
+		for _, id := range IDs {
+			pubShare.Share(id)
+		}
+	}
+}
diff --git a/core/db/level-db_test.go b/core/db/level-db_test.go
index d7f70bc..78e3b77 100644
--- a/core/db/level-db_test.go
+++ b/core/db/level-db_test.go
@@ -214,8 +214,8 @@ func (s *LevelDBTestSuite) TestDKGProtocolInfoRLPEncodeDecode() {
 			types.NodeID{Hash: common.Hash{0x02}}: dkg.ID{},
 		},
 		MpkMap: NodeIDToPubShares{
-			types.NodeID{Hash: common.Hash{0x01}}: &dkg.PublicKeyShares{},
-			types.NodeID{Hash: common.Hash{0x02}}: &dkg.PublicKeyShares{},
+			types.NodeID{Hash: common.Hash{0x01}}: dkg.NewEmptyPublicKeyShares(),
+			types.NodeID{Hash: common.Hash{0x02}}: dkg.NewEmptyPublicKeyShares(),
 		},
 		AntiComplaintReceived: NodeIDToNodeIDs{
 			types.NodeID{Hash: common.Hash{0x01}}: map[types.NodeID]struct{}{
@@ -278,8 +278,8 @@ func (s *LevelDBTestSuite) TestNodeIDRLPEncodeDecode() {
 
 func (s *LevelDBTestSuite) TestNodeIDToPubSharesRLPEncodeDecode() {
 	m := NodeIDToPubShares{
-		types.NodeID{Hash: common.Hash{0x01}}: &dkg.PublicKeyShares{},
-		types.NodeID{Hash: common.Hash{0x02}}: &dkg.PublicKeyShares{},
+		types.NodeID{Hash: common.Hash{0x01}}: dkg.NewEmptyPublicKeyShares(),
+		types.NodeID{Hash: common.Hash{0x02}}: dkg.NewEmptyPublicKeyShares(),
 	}
 
 	b, err := rlp.EncodeToBytes(&m)
diff --git a/core/dkg-tsig-protocol.go b/core/dkg-tsig-protocol.go
index e4ae14c..0612bda 100644
--- a/core/dkg-tsig-protocol.go
+++ b/core/dkg-tsig-protocol.go
@@ -230,7 +230,7 @@ func newDKGProtocol(
 		Round:           round,
 		Reset:           reset,
 		DKGID:           typesDKG.NewID(ID),
-		PublicKeyShares: *pubShare,
+		PublicKeyShares: *pubShare.Move(),
 	})
 
 	return &dkgProtocol{
diff --git a/core/dkg-tsig-protocol_test.go b/core/dkg-tsig-protocol_test.go
index f6909e8..89fd105 100644
--- a/core/dkg-tsig-protocol_test.go
+++ b/core/dkg-tsig-protocol_test.go
@@ -294,7 +294,7 @@ func (s *DKGTSIGProtocolTestSuite) TestErrMPKRegistered() {
 				Round:           round,
 				Reset:           reset,
 				DKGID:           typesDKG.NewID(ID),
-				PublicKeyShares: *mpk,
+				PublicKeyShares: *mpk.Move(),
 			})
 		}
 		gov.AddDKGMasterPublicKey(receiver.mpk)
@@ -995,7 +995,7 @@ func (s *DKGTSIGProtocolTestSuite) TestUnexpectedDKGResetCount() {
 		Round:           round,
 		Reset:           reset + 1,
 		DKGID:           typesDKG.NewID(sourceID),
-		PublicKeyShares: *mpk,
+		PublicKeyShares: *mpk.Move(),
 	})
 	err = protocols[sourceID].processMasterPublicKeys(
 		[]*typesDKG.MasterPublicKey{receivers[sourceID].mpk})
@@ -1031,7 +1031,7 @@ func benchmarkDKGGroupPubliKey(k, n int, b *testing.B) {
 			Round:           round,
 			Reset:           reset,
 			DKGID:           typesDKG.NewID(types.NewNodeID(pk)),
-			PublicKeyShares: *pubShare,
+			PublicKeyShares: *pubShare.Move(),
 		})
 	}
 
@@ -1075,7 +1075,7 @@ func benchmarkDKGNodePubliKeys(k, n int, b *testing.B) {
 			Round:           round,
 			Reset:           reset,
 			DKGID:           typesDKG.NewID(types.NewNodeID(pk)),
-			PublicKeyShares: *pubShare,
+			PublicKeyShares: *pubShare.Move(),
 		})
 	}
 
@@ -1119,7 +1119,7 @@ func benchmarkCalcQualified(k, n int, b *testing.B) {
 			Round:           round,
 			Reset:           reset,
 			DKGID:           typesDKG.NewID(types.NewNodeID(pk)),
-			PublicKeyShares: *pubShare,
+			PublicKeyShares: *pubShare.Move(),
 		})
 	}
 
diff --git a/core/test/app_test.go b/core/test/app_test.go
index 574604c..828a3c3 100644
--- a/core/test/app_test.go
+++ b/core/test/app_test.go
@@ -82,7 +82,7 @@ func (s *AppTestSuite) proposeMPK(
 			Round:           round,
 			Reset:           reset,
 			DKGID:           typesDKG.NewID(types.NewNodeID(pubKey)),
-			PublicKeyShares: *pubShare,
+			PublicKeyShares: *pubShare.Move(),
 		}
 		s.Require().NoError(s.signers[idx].SignDKGMasterPublicKey(mpk))
 		gov.AddDKGMasterPublicKey(mpk)
diff --git a/core/test/governance_test.go b/core/test/governance_test.go
index 474ec80..e075638 100644
--- a/core/test/governance_test.go
+++ b/core/test/governance_test.go
@@ -130,7 +130,7 @@ func (s *GovernanceTestSuite) TestProhibit() {
 		mpk := &typesDKG.MasterPublicKey{
 			Round:           round,
 			DKGID:           typesDKG.NewID(types.NewNodeID(k.PublicKey())),
-			PublicKeyShares: *pubShare,
+			PublicKeyShares: *pubShare.Move(),
 		}
 		s.Require().NoError(signer.SignDKGMasterPublicKey(mpk))
 		gov.AddDKGMasterPublicKey(mpk)
diff --git a/core/test/state_test.go b/core/test/state_test.go
index 5ad5a3e..e3ed9bb 100644
--- a/core/test/state_test.go
+++ b/core/test/state_test.go
@@ -50,7 +50,7 @@ func (s *StateTestSuite) newDKGMasterPublicKey(
 		Round:           round,
 		Reset:           reset,
 		DKGID:           dID,
-		PublicKeyShares: *pubShare,
+		PublicKeyShares: *pubShare.Move(),
 	}
 }
 
diff --git a/core/types/dkg/dkg.go b/core/types/dkg/dkg.go
index e9b22bc..868f0da 100644
--- a/core/types/dkg/dkg.go
+++ b/core/types/dkg/dkg.go
@@ -130,7 +130,7 @@ func (d *MasterPublicKey) DecodeRLP(s *rlp.Stream) error {
 		Round:           dec.Round,
 		Reset:           dec.Reset,
 		DKGID:           id,
-		PublicKeyShares: *dec.PublicKeyShares,
+		PublicKeyShares: *dec.PublicKeyShares.Move(),
 		Signature:       dec.Signature,
 	}
 	return err
diff --git a/core/types/dkg/dkg_test.go b/core/types/dkg/dkg_test.go
index ec16a2b..9f50feb 100644
--- a/core/types/dkg/dkg_test.go
+++ b/core/types/dkg/dkg_test.go
@@ -63,7 +63,7 @@ func (s *DKGTestSuite) TestRLPEncodeDecode() {
 		Round:           10,
 		Reset:           11,
 		DKGID:           dID,
-		PublicKeyShares: *pubShare,
+		PublicKeyShares: *pubShare.Clone(),
 		Signature: crypto.Signature{
 			Type:      "123",
 			Signature: []byte{4, 5, 6},
@@ -189,7 +189,7 @@ func (s *DKGTestSuite) TestMasterPublicKeyEquality() {
 	pubKey := prvKey.PublicKey().(cryptoDKG.PublicKey)
 	_, pubShares := cryptoDKG.NewPrivateKeyShares(2)
 	req.NoError(pubShares.AddShare(s.genID(), &pubKey))
-	master1.PublicKeyShares = *pubShares
+	master1.PublicKeyShares = *pubShares.Move()
 	// Prepare another master public key by copying every field.
 	master2 := &MasterPublicKey{}
 	s.clone(master1, master2)
diff --git a/core/utils/crypto_test.go b/core/utils/crypto_test.go
index 061f250..5dfd82b 100644
--- a/core/utils/crypto_test.go
+++ b/core/utils/crypto_test.go
@@ -179,7 +179,7 @@ func (s *CryptoTestSuite) TestDKGSignature() {
 		Round:           5,
 		Reset:           6,
 		DKGID:           id,
-		PublicKeyShares: *pkShare,
+		PublicKeyShares: *pkShare.Move(),
 	}
 	mpk.Signature, err = prv.Sign(hashDKGMasterPublicKey(mpk))
 	s.Require().NoError(err)
diff --git a/core/utils/penalty-helper_test.go b/core/utils/penalty-helper_test.go
index 7b61781..3e4f8b5 100644
--- a/core/utils/penalty-helper_test.go
+++ b/core/utils/penalty-helper_test.go
@@ -52,7 +52,7 @@ func (s *PenaltyHelperTestSuite) TestDKGComplaint() {
 	mpk := &typesDKG.MasterPublicKey{
 		ProposerID:      nID1,
 		DKGID:           typesDKG.NewID(nID1),
-		PublicKeyShares: *pubShares,
+		PublicKeyShares: *pubShares.Move(),
 	}
 	mpk.Signature, err = prv1.Sign(hashDKGMasterPublicKey(mpk))
 	s.Require().NoError(err)
diff --git a/core/utils/utils_test.go b/core/utils/utils_test.go
index 70f6e98..c3afea9 100644
--- a/core/utils/utils_test.go
+++ b/core/utils/utils_test.go
@@ -53,7 +53,7 @@ func (s *UtilsTestSuite) TestVerifyDKGComplaint() {
 	mpk := &typesDKG.MasterPublicKey{
 		ProposerID:      nID1,
 		DKGID:           typesDKG.NewID(nID1),
-		PublicKeyShares: *pubShares,
+		PublicKeyShares: *pubShares.Move(),
 	}
 	mpk.Signature, err = prv1.Sign(hashDKGMasterPublicKey(mpk))
 	s.Require().NoError(err)
diff --git a/integration_test/round-event_test.go b/integration_test/round-event_test.go
index f83a437..dca0834 100644
--- a/integration_test/round-event_test.go
+++ b/integration_test/round-event_test.go
@@ -82,7 +82,7 @@ func (s *RoundEventTestSuite) proposeMPK(
 			Round:           round,
 			Reset:           reset,
 			DKGID:           typesDKG.NewID(types.NewNodeID(pubKey)),
-			PublicKeyShares: *pubShare,
+			PublicKeyShares: *pubShare.Move(),
 		}
 		s.Require().NoError(s.signers[idx].SignDKGMasterPublicKey(mpk))
 		gov.AddDKGMasterPublicKey(mpk)