Address comments.

Author: arijitAD

dot/network/message_cache.go

@@ -1,6 +1,7 @@
 package network
 
 import (
+	"errors"
 	"time"
 
 	"github.com/ChainSafe/gossamer/lib/common"
@@ -31,8 +32,8 @@ func newMessageCache(config ristretto.Config, ttl time.Duration) (*messageCache,
 	return &messageCache{cache: cache, ttl: ttl}, nil
 }
 
-// Put appends peer ID and message data and stores it in cache with TTL.
-func (m *messageCache) Put(peer peer.ID, msg []byte) (bool, error) {
+// put appends peer ID and message data and stores it in cache with TTL.
+func (m *messageCache) put(peer peer.ID, msg NotificationsMessage) (bool, error) {
 	key, err := generateCacheKey(peer, msg)
 	if err != nil {
 		return false, err
@@ -47,8 +48,8 @@ func (m *messageCache) Put(peer peer.ID, msg []byte) (bool, error) {
 	return ok, nil
 }
 
-// Exists checks if <peer ID, message data> exist in cache.
-func (m *messageCache) Exists(peer peer.ID, msg []byte) bool {
+// exists checks if <peer ID, message> exist in cache.
+func (m *messageCache) exists(peer peer.ID, msg NotificationsMessage) bool {
 	key, err := generateCacheKey(peer, msg)
 	if err != nil {
 		return false
@@ -58,8 +59,12 @@ func (m *messageCache) Exists(peer peer.ID, msg []byte) bool {
 	return ok
 }
 
-func generateCacheKey(peer peer.ID, msg []byte) ([]byte, error) {
-	peerMsgHash, err := common.Blake2bHash(append([]byte(peer), msg...))
+func generateCacheKey(peer peer.ID, msg NotificationsMessage) ([]byte, error) {
+	if msg.IsHandshake() {
+		return nil, errors.New("cache does not support handshake messages")
+	}
+
+	peerMsgHash, err := common.Blake2bHash(append([]byte(peer), msg.Hash().ToBytes()...))
 	if err != nil {
 		return nil, err
 	}

dot/network/message_cache_test.go

@@ -1,9 +1,12 @@
 package network
 
 import (
+	"math/big"
 	"testing"
 	"time"
 
+	"github.com/ChainSafe/gossamer/dot/types"
+	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/dgraph-io/ristretto"
 	"github.com/libp2p/go-libp2p-core/peer"
 	"github.com/stretchr/testify/require"
@@ -19,22 +22,53 @@ func TestMessageCache(t *testing.T) {
 			return int64(1)
 		},
 	}, 800*time.Millisecond)
+	require.NoError(t, err)
 
 	peerID := peer.ID("gossamer")
-	msgData := []byte("testData")
-	require.NoError(t, err)
+	msg := &BlockAnnounceMessage{
+		ParentHash:     common.Hash{1},
+		Number:         big.NewInt(77),
+		StateRoot:      common.Hash{2},
+		ExtrinsicsRoot: common.Hash{3},
+		Digest:         types.Digest{},
+	}
 
-	ok, err := msgCache.Put(peerID, msgData)
+	ok, err := msgCache.put(peerID, msg)
 	require.NoError(t, err)
 	require.True(t, ok)
 
 	time.Sleep(750 * time.Millisecond)
 
-	ok = msgCache.Exists(peerID, msgData)
+	ok = msgCache.exists(peerID, msg)
 	require.True(t, ok)
 
 	time.Sleep(50 * time.Millisecond)
 
-	ok = msgCache.Exists(peerID, msgData)
+	ok = msgCache.exists(peerID, msg)
+	require.False(t, ok)
+}
+
+func TestMessageCacheError(t *testing.T) {
+	cacheSize := 64 << 20 // 64 MB
+	msgCache, err := newMessageCache(ristretto.Config{
+		NumCounters: int64(float64(cacheSize) * 0.05 * 2),
+		MaxCost:     int64(float64(cacheSize) * 0.95),
+		BufferItems: 64,
+		Cost: func(value interface{}) int64 {
+			return int64(1)
+		},
+	}, 800*time.Millisecond)
+	require.NoError(t, err)
+
+	peerID := peer.ID("gossamer")
+	msg := &BlockAnnounceHandshake{
+		Roles:           4,
+		BestBlockNumber: 77,
+		BestBlockHash:   common.Hash{1},
+		GenesisHash:     common.Hash{2},
+	}
+
+	ok, err := msgCache.put(peerID, msg)
+	require.Error(t, err, "cache does not support handshake messages")
 	require.False(t, ok)
 }

dot/network/notifications.go

@@ -243,17 +243,14 @@ func (s *Service) broadcastExcluding(info *notificationsProtocol, excluding peer
 			err = s.host.send(peer, info.protocolID, hs)
 		} else {
 			if s.host.messageCache != nil {
-				var encMsg []byte
-				encMsg, err = msg.Encode()
+				var added bool
+				added, err = s.host.messageCache.put(peer, msg)
 				if err != nil {
-					logger.Error("failed to encode message", "peer", peer, "error", err)
+					logger.Error("failed to add message to cache", "peer", peer, "error", err)
 					continue
 				}
 
-				var added bool
-				added, err = s.host.messageCache.Put(peer, encMsg)
-				if err != nil || !added {
-					logger.Error("failed to add message to cache", "peer", peer, "error", err)
+				if !added {
 					continue
 				}
 			}