Merge pull request #3437 from ipfs/feat/bitswap-cleanup-ledger

bitswap: clean up ledgers when disconnecting

Author: whyrusleeping

exchange/bitswap/bitswap.go

@@ -414,6 +414,7 @@ func (bs *Bitswap) updateReceiveCounters(b blocks.Block) {
 // Connected/Disconnected warns bitswap about peer connections
 func (bs *Bitswap) PeerConnected(p peer.ID) {
 	bs.wm.Connected(p)
+	bs.engine.PeerConnected(p)
 }
 
 // Connected/Disconnected warns bitswap about peer connections

exchange/bitswap/decision/engine.go

@@ -298,8 +298,32 @@ func (e *Engine) MessageSent(p peer.ID, m bsmsg.BitSwapMessage) error {
 	return nil
 }
 
+func (e *Engine) PeerConnected(p peer.ID) {
+	e.lock.Lock()
+	l, ok := e.ledgerMap[p]
+	if !ok {
+		l = newLedger(p)
+		e.ledgerMap[p] = l
+	}
+	l.lk.Lock()
+	l.ref++
+	l.lk.Unlock()
+	e.lock.Unlock()
+}
+
 func (e *Engine) PeerDisconnected(p peer.ID) {
-	// TODO: release ledger
+	e.lock.Lock()
+	defer e.lock.Unlock()
+	l, ok := e.ledgerMap[p]
+	if !ok {
+		return
+	}
+	l.lk.Lock()
+	l.ref--
+	if l.ref <= 0 {
+		delete(e.ledgerMap, p)
+	}
+	l.lk.Unlock()
 }
 
 func (e *Engine) numBytesSentTo(p peer.ID) uint64 {

exchange/bitswap/decision/engine_test.go

@@ -89,6 +89,11 @@ func TestPeerIsAddedToPeersWhenMessageReceivedOrSent(t *testing.T) {
 	if !peerIsPartner(sanfrancisco.Peer, seattle.Engine) {
 		t.Fatal("Peer wasn't added as a Partner")
 	}
+
+	seattle.Engine.PeerDisconnected(sanfrancisco.Peer)
+	if peerIsPartner(sanfrancisco.Peer, seattle.Engine) {
+		t.Fatal("expected peer to be removed")
+	}
 }
 
 func peerIsPartner(p peer.ID, e *Engine) bool {

exchange/bitswap/decision/ledger.go

@@ -43,6 +43,10 @@ type ledger struct {
 	// to a given peer
 	sentToPeer map[string]time.Time
 
+	// ref is the reference count for this ledger, its used to ensure we
+	// don't drop the reference to this ledger in multi-connection scenarios
+	ref int
+
 	lk sync.Mutex
 }