From 5bbaedb2b9e5813b00710239014b2b10c798e6e4 Mon Sep 17 00:00:00 2001
From: Jorropo <jorropo.pgm@gmail.com>
Date: Wed, 6 Apr 2022 04:21:01 +0200
Subject: [PATCH] fix: shard lock to fix contention issue in the notify handler

---
 internal/decision/engine.go    |  6 ++++--
 network/connecteventmanager.go | 10 +++++++---
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/internal/decision/engine.go b/internal/decision/engine.go
index c8c33097..297deac8 100644
--- a/internal/decision/engine.go
+++ b/internal/decision/engine.go
@@ -914,20 +914,20 @@ func (e *Engine) MessageSent(p peer.ID, m bsmsg.BitSwapMessage) {
 // sending blocks.
 func (e *Engine) PeerConnected(p peer.ID) {
 	e.lock.Lock()
-	defer e.lock.Unlock()
 
 	_, ok := e.ledgerMap[p]
 	if !ok {
 		e.ledgerMap[p] = newLedger(p)
 	}
 
+	e.lock.Unlock()
+
 	e.scoreLedger.PeerConnected(p)
 }
 
 // PeerDisconnected is called when a peer disconnects.
 func (e *Engine) PeerDisconnected(p peer.ID) {
 	e.lock.Lock()
-	defer e.lock.Unlock()
 
 	ledger, ok := e.ledgerMap[p]
 	if ok {
@@ -941,6 +941,8 @@ func (e *Engine) PeerDisconnected(p peer.ID) {
 	}
 	delete(e.ledgerMap, p)
 
+	e.lock.Unlock()
+
 	e.scoreLedger.PeerDisconnected(p)
 }
 
diff --git a/network/connecteventmanager.go b/network/connecteventmanager.go
index bbde7af2..01428196 100644
--- a/network/connecteventmanager.go
+++ b/network/connecteventmanager.go
@@ -31,7 +31,6 @@ func newConnectEventManager(connListener ConnectionListener) *connectEventManage
 
 func (c *connectEventManager) Connected(p peer.ID) {
 	c.lk.Lock()
-	defer c.lk.Unlock()
 
 	state, ok := c.conns[p]
 	if !ok {
@@ -41,13 +40,15 @@ func (c *connectEventManager) Connected(p peer.ID) {
 	state.refs++
 
 	if state.refs == 1 && state.responsive {
+		c.lk.Unlock()
 		c.connListener.PeerConnected(p)
+		return
 	}
+	c.lk.Unlock()
 }
 
 func (c *connectEventManager) Disconnected(p peer.ID) {
 	c.lk.Lock()
-	defer c.lk.Unlock()
 
 	state, ok := c.conns[p]
 	if !ok {
@@ -57,11 +58,14 @@ func (c *connectEventManager) Disconnected(p peer.ID) {
 	state.refs--
 
 	if state.refs == 0 {
+		delete(c.conns, p)
 		if state.responsive {
+			c.lk.Unlock()
 			c.connListener.PeerDisconnected(p)
+			return
 		}
-		delete(c.conns, p)
 	}
+	c.lk.Unlock()
 }
 
 func (c *connectEventManager) MarkUnresponsive(p peer.ID) {