diff --git a/exchange/bitswap/bitswap.go b/exchange/bitswap/bitswap.go index d76dbb32086..e7a20008ba7 100644 --- a/exchange/bitswap/bitswap.go +++ b/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 diff --git a/exchange/bitswap/decision/engine.go b/exchange/bitswap/decision/engine.go index 6c1a9e93653..37e370db071 100644 --- a/exchange/bitswap/decision/engine.go +++ b/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 { diff --git a/exchange/bitswap/decision/engine_test.go b/exchange/bitswap/decision/engine_test.go index 851e1469da9..fdac4eba106 100644 --- a/exchange/bitswap/decision/engine_test.go +++ b/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 { diff --git a/exchange/bitswap/decision/ledger.go b/exchange/bitswap/decision/ledger.go index ac83624674c..cb93f0e95fe 100644 --- a/exchange/bitswap/decision/ledger.go +++ b/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 }