bitswap: clean up ledgers when disconnecting

License: MIT
Signed-off-by: Jeromy <why@ipfs.io>

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

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 {

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
 }