Fix relayMessageCount read lock

Signed-off-by: Yilun <[email protected]>
nknorg · Jan 18, 2019 · 0f8815e · 0f8815e
1 parent 312abe0
commit 0f8815e
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 6 deletions.
diff --git a/net/node/localnode.go b/net/node/localnode.go
@@ -37,8 +37,8 @@ type LocalNode struct {
 	*messageHandlerStore
 
 	sync.RWMutex
-	syncOnce *sync.Once
-	relayMessageCount int        // count node relayData
+	syncOnce          *sync.Once
+	relayMessageCount uint64 // count how many messages node has relayed since start
 }
 
 func (localNode *LocalNode) MarshalJSON() ([]byte, error) {
@@ -56,7 +56,7 @@ func (localNode *LocalNode) MarshalJSON() ([]byte, error) {
 
 	out["height"] = localNode.GetHeight()
 	out["version"] = config.Version
-	out["relayMessageCount"] = localNode.relayMessageCount
+	out["relayMessageCount"] = localNode.GetRelayMessageCount()
 
 	return json.Marshal(out)
 }
@@ -197,6 +197,18 @@ func (localNode *LocalNode) maybeAddRemoteNode(remoteNode *nnetnode.RemoteNode)
 	return nil
 }
 
+func (localNode *LocalNode) GetRelayMessageCount() uint64 {
+	localNode.RLock()
+	defer localNode.RUnlock()
+	return localNode.relayMessageCount
+}
+
+func (localNode *LocalNode) IncrementRelayMessageCount() {
+	localNode.Lock()
+	localNode.relayMessageCount++
+	localNode.Unlock()
+}
+
 func (localNode *LocalNode) GetTxnPool() *pool.TxnPool {
 	return localNode.TxnPool
 }

diff --git a/net/node/message.go b/net/node/message.go
@@ -202,9 +202,8 @@ func (localNode *LocalNode) remoteMessageRouted(remoteMessage *nnetnode.RemoteMe
 				log.Errorf("Marshal new relay msg body error: %v", err)
 				return nil, nil, nil, false
 			}
-			localNode.RWMutex.Lock()
-			localNode.relayMessageCount++
-			localNode.RWMutex.Unlock()
+
+			localNode.IncrementRelayMessageCount()
 		}
 
 		// msg send to local node