Repository: ipfs/go-bitswap
Pull request #361 (2020-04-21 00:02:39)

fix: minimize time holding wantlist lock
Instead of holding the lock the entire time we prepare a message, hold the lock while we retrieve the wantlist entries, process the entries without the lock, retake the lock, then mark entries as sent.

This means:

1. We never sort entries while holding the lock.
2. We allocate exactly three times while holding the lock (once per entry list).

This change also tries to be as lazy as possible when sorting.

Stebalien: Note: this doesn't appear to be showing up as a significant source of lock contention at this point anyways. I believe the issue was mainly due to collecting _very_ large wantlists due to unresponsive peers.

Patch:
diff --git a/go.mod b/go.mod
index 29e392c..6b3460f 100644
--- a/go.mod
+++ b/go.mod
@@ -27,6 +27,7 @@ require (
 	github.com/libp2p/go-libp2p-testing v0.1.1
 	github.com/libp2p/go-msgio v0.0.4
 	github.com/multiformats/go-multiaddr v0.2.1
+	github.com/multiformats/go-multistream v0.1.1
 	go.uber.org/zap v1.14.1
 )
 
diff --git a/internal/decision/engine.go b/internal/decision/engine.go
index 620bb86..81ef9b9 100644
--- a/internal/decision/engine.go
+++ b/internal/decision/engine.go
@@ -337,9 +337,13 @@ func (e *Engine) onPeerRemoved(p peer.ID) {
 // WantlistForPeer returns the currently understood want list for a given peer
 func (e *Engine) WantlistForPeer(p peer.ID) (out []wl.Entry) {
 	partner := e.findOrCreate(p)
+
 	partner.lk.Lock()
-	defer partner.lk.Unlock()
-	return partner.wantList.SortedEntries()
+	entries := partner.wantList.Entries()
+	partner.lk.Unlock()
+
+	wl.SortEntries(entries)
+	return
 }
 
 // LedgerForPeer returns aggregated data about blocks swapped and communication
diff --git a/internal/messagequeue/messagequeue.go b/internal/messagequeue/messagequeue.go
index b8323a7..7bcc087 100644
--- a/internal/messagequeue/messagequeue.go
+++ b/internal/messagequeue/messagequeue.go
@@ -115,9 +115,15 @@ func (r *recallWantlist) RemoveType(c cid.Cid, wtype pb.Message_Wantlist_WantTyp
 }
 
 // MarkSent moves the want from the pending to the sent list
-func (r *recallWantlist) MarkSent(e wantlist.Entry) {
-	r.pending.RemoveType(e.Cid, e.WantType)
+//
+// Returns true if the want was marked as sent. Returns false if the want wasn't
+// pending.
+func (r *recallWantlist) MarkSent(e wantlist.Entry) bool {
+	if !r.pending.RemoveType(e.Cid, e.WantType) {
+		return false
+	}
 	r.sent.Add(e.Cid, e.Priority, e.WantType)
+	return true
 }
 
 type peerConn struct {
@@ -539,74 +545,77 @@ func (mq *MessageQueue) pendingWorkCount() int {
 
 // Convert the lists of wants into a Bitswap message
 func (mq *MessageQueue) extractOutgoingMessage(supportsHave bool) bsmsg.BitSwapMessage {
-	mq.wllock.Lock()
-	defer mq.wllock.Unlock()
-
 	// Get broadcast and regular wantlist entries.
-	// SortedEntries() slows down the MessageQueue a lot, and entries only need
-	// to be sorted if the number of wants will overflow the size of the
-	// message (to make sure that the highest priority wants are sent in the
-	// first message).
-	// We prioritize cancels, then regular wants, then broadcast wants.
-	var peerEntries []bswl.Entry
-	var bcstEntries []bswl.Entry
-	maxCancelsSize := mq.cancels.Len() * bsmsg.MaxEntrySize
-	maxPeerSize := mq.peerWants.pending.Len() * bsmsg.MaxEntrySize
-	maxBcstSize := mq.bcstWants.pending.Len() * bsmsg.MaxEntrySize
-
-	if maxCancelsSize+maxPeerSize < mq.maxMessageSize {
-		peerEntries = mq.peerWants.pending.Entries()
-	} else {
-		peerEntries = mq.peerWants.pending.SortedEntries()
-	}
-	if maxCancelsSize+maxPeerSize+maxBcstSize < mq.maxMessageSize {
-		bcstEntries = mq.bcstWants.pending.Entries()
-	} else {
-		bcstEntries = mq.bcstWants.pending.SortedEntries()
+	mq.wllock.Lock()
+	peerEntries := mq.peerWants.pending.Entries()
+	bcstEntries := mq.bcstWants.pending.Entries()
+	cancels := mq.cancels.Keys()
+	if !supportsHave {
+		filteredPeerEntries := peerEntries[:0]
+		// If the remote peer doesn't support HAVE / DONT_HAVE messages,
+		// don't send want-haves (only send want-blocks)
+		//
+		// Doing this here under the lock makes everything else in this
+		// function simpler.
+		//
+		// TODO: We should _try_ to avoid recording these in the first
+		// place if possible.
+		for _, e := range peerEntries {
+			if e.WantType == pb.Message_Wantlist_Have {
+				mq.peerWants.RemoveType(e.Cid, pb.Message_Wantlist_Have)
+			} else {
+				filteredPeerEntries = append(filteredPeerEntries, e)
+			}
+		}
+		peerEntries = filteredPeerEntries
 	}
+	mq.wllock.Unlock()
 
-	// Size of the message so far
-	msgSize := 0
+	// We prioritize cancels, then regular wants, then broadcast wants.
 
-	// Always prioritize cancels, then targeted, then broadcast.
+	var (
+		msgSize         = 0 // size of message so far
+		sentCancels     = 0 // number of cancels in message
+		sentPeerEntries = 0 // number of peer entries in message
+		sentBcstEntries = 0 // number of broadcast entries in message
+	)
 
 	// Add each cancel to the message
-	cancels := mq.cancels.Keys()
 	for _, c := range cancels {
+		msgSize += mq.msg.Cancel(c)
+		sentCancels++
+
 		if msgSize >= mq.maxMessageSize {
-			break
+			goto FINISH
 		}
-		msgSize += mq.msg.Cancel(c)
+	}
 
-		// Clear the cancel - we make a best effort to let peers know about
-		// cancels but won't save them to resend if there's a failure.
-		mq.cancels.Remove(c)
+	// Next, add the wants. If we have too many entries to fit into a single
+	// message, sort by priority and include the high priority ones first.
+	// However, avoid sorting till we really need to as this code is a
+	// called frequently.
+
+	// Add each regular want-have / want-block to the message.
+	if msgSize+(len(peerEntries)*bsmsg.MaxEntrySize) > mq.maxMessageSize {
+		bswl.SortEntries(peerEntries)
 	}
 
-	// Add each regular want-have / want-block to the message
 	for _, e := range peerEntries {
+		msgSize += mq.msg.AddEntry(e.Cid, e.Priority, e.WantType, true)
+		sentPeerEntries++
+
 		if msgSize >= mq.maxMessageSize {
-			break
+			goto FINISH
 		}
+	}
 
-		// If the remote peer doesn't support HAVE / DONT_HAVE messages,
-		// don't send want-haves (only send want-blocks)
-		if !supportsHave && e.WantType == pb.Message_Wantlist_Have {
-			mq.peerWants.RemoveType(e.Cid, pb.Message_Wantlist_Have)
-		} else {
-			msgSize += mq.msg.AddEntry(e.Cid, e.Priority, e.WantType, true)
-
-			// Move the key from pending to sent
-			mq.peerWants.MarkSent(e)
-		}
+	// Add each broadcast want-have to the message.
+	if msgSize+(len(bcstEntries)*bsmsg.MaxEntrySize) > mq.maxMessageSize {
+		bswl.SortEntries(bcstEntries)
 	}
 
 	// Add each broadcast want-have to the message
 	for _, e := range bcstEntries {
-		if msgSize >= mq.maxMessageSize {
-			break
-		}
-
 		// Broadcast wants are sent as want-have
 		wantType := pb.Message_Wantlist_Have
 
@@ -617,11 +626,40 @@ func (mq *MessageQueue) extractOutgoingMessage(supportsHave bool) bsmsg.BitSwapM
 		}
 
 		msgSize += mq.msg.AddEntry(e.Cid, e.Priority, wantType, false)
+		sentBcstEntries++
 
-		// Move the key from pending to sent
-		mq.bcstWants.MarkSent(e)
+		if msgSize >= mq.maxMessageSize {
+			goto FINISH
+		}
 	}
 
+FINISH:
+
+	// Finally, re-take the lock, mark sent and remove any entries from our
+	// message that we've decided to cancel at the last minute.
+	mq.wllock.Lock()
+	for _, e := range peerEntries[:sentPeerEntries] {
+		if !mq.peerWants.MarkSent(e) {
+			// It changed.
+			mq.msg.Remove(e.Cid)
+		}
+	}
+
+	for _, e := range bcstEntries[:sentBcstEntries] {
+		if !mq.bcstWants.MarkSent(e) {
+			mq.msg.Remove(e.Cid)
+		}
+	}
+
+	for _, c := range cancels[:sentCancels] {
+		if !mq.cancels.Has(c) {
+			mq.msg.Remove(c)
+		} else {
+			mq.cancels.Remove(c)
+		}
+	}
+	mq.wllock.Unlock()
+
 	return mq.msg
 }
 
diff --git a/message/message.go b/message/message.go
index f820c9d..88c3f7d 100644
--- a/message/message.go
+++ b/message/message.go
@@ -44,6 +44,10 @@ type BitSwapMessage interface {
 	// Returns the size of the CANCEL entry in the protobuf
 	Cancel(key cid.Cid) int
 
+	// Remove removes any entries for the given CID. Useful when the want
+	// status for the CID changes when preparing a message.
+	Remove(key cid.Cid)
+
 	// Empty indicates whether the message has any information
 	Empty() bool
 	// Size returns the size of the message in bytes
@@ -298,6 +302,10 @@ func (m *impl) SetPendingBytes(pendingBytes int32) {
 	m.pendingBytes = pendingBytes
 }
 
+func (m *impl) Remove(k cid.Cid) {
+	delete(m.wantlist, k)
+}
+
 func (m *impl) Cancel(k cid.Cid) int {
 	return m.addEntry(k, 0, true, pb.Message_Wantlist_Block, false)
 }
diff --git a/wantlist/wantlist.go b/wantlist/wantlist.go
index e18567d..555c293 100644
--- a/wantlist/wantlist.go
+++ b/wantlist/wantlist.go
@@ -111,16 +111,14 @@ func (w *Wantlist) Entries() []Entry {
 	return es
 }
 
-// SortedEntries returns wantlist entries ordered by priority.
-func (w *Wantlist) SortedEntries() []Entry {
-	es := w.Entries()
-	sort.Sort(entrySlice(es))
-	return es
-}
-
 // Absorb all the entries in other into this want list
 func (w *Wantlist) Absorb(other *Wantlist) {
 	for _, e := range other.Entries() {
 		w.Add(e.Cid, e.Priority, e.WantType)
 	}
 }
+
+// SortEntries sorts the list of entries by priority.
+func SortEntries(es []Entry) {
+	sort.Sort(entrySlice(es))
+}


Test patch:
diff --git a/wantlist/wantlist_test.go b/wantlist/wantlist_test.go
index 1139e87..49dc559 100644
--- a/wantlist/wantlist_test.go
+++ b/wantlist/wantlist_test.go
@@ -203,14 +203,16 @@ func TestAbsort(t *testing.T) {
 	}
 }
 
-func TestSortedEntries(t *testing.T) {
+func TestSortEntries(t *testing.T) {
 	wl := New()
 
 	wl.Add(testcids[0], 3, pb.Message_Wantlist_Block)
 	wl.Add(testcids[1], 5, pb.Message_Wantlist_Have)
 	wl.Add(testcids[2], 4, pb.Message_Wantlist_Have)
 
-	entries := wl.SortedEntries()
+	entries := wl.Entries()
+	SortEntries(entries)
+
 	if !entries[0].Cid.Equals(testcids[1]) ||
 		!entries[1].Cid.Equals(testcids[2]) ||
 		!entries[2].Cid.Equals(testcids[0]) {