feat: cache the materialized wantlist

internal/decision/engine.go

@@ -328,8 +328,6 @@ func (e *Engine) WantlistForPeer(p peer.ID) []wl.Entry {
 	entries := partner.wantList.Entries()
 	partner.lk.Unlock()
 
-	wl.SortEntries(entries)
-
 	return entries
 }
 

internal/messagequeue/messagequeue.go

@@ -740,13 +740,6 @@ func (mq *MessageQueue) extractOutgoingMessage(supportsHave bool) (bsmsg.BitSwap
 
 	// 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)
-	}
 
 	for _, e := range peerEntries {
 		msgSize += mq.msg.AddEntry(e.Cid, e.Priority, e.WantType, true)
@@ -757,11 +750,6 @@ func (mq *MessageQueue) extractOutgoingMessage(supportsHave bool) (bsmsg.BitSwap
 		}
 	}
 
-	// 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 {
 		// Broadcast wants are sent as want-have

wantlist/wantlist.go

@@ -13,6 +13,9 @@ import (
 // Wantlist is a raw list of wanted blocks and their priorities
 type Wantlist struct {
 	set map[cid.Cid]Entry
+
+	// Re-computing this can get expensive so we memoize it.
+	cached []Entry
 }
 
 // Entry is an entry in a want list, consisting of a cid and its priority
@@ -58,11 +61,11 @@ func (w *Wantlist) Add(c cid.Cid, priority int32, wantType pb.Message_Wantlist_W
 		return false
 	}
 
-	w.set[c] = Entry{
+	w.put(c, Entry{
 		Cid:      c,
 		Priority: priority,
 		WantType: wantType,
-	}
+	})
 
 	return true
 }
@@ -74,7 +77,7 @@ func (w *Wantlist) Remove(c cid.Cid) bool {
 		return false
 	}
 
-	delete(w.set, c)
+	w.delete(c)
 	return true
 }
 
@@ -91,34 +94,49 @@ func (w *Wantlist) RemoveType(c cid.Cid, wantType pb.Message_Wantlist_WantType)
 		return false
 	}
 
-	delete(w.set, c)
+	w.delete(c)
 	return true
 }
 
+func (w *Wantlist) delete(c cid.Cid) {
+	delete(w.set, c)
+	w.cached = nil
+}
+
+func (w *Wantlist) put(c cid.Cid, e Entry) {
+	w.cached = nil
+	w.set[c] = e
+}
+
 // Contains returns the entry, if present, for the given CID, plus whether it
 // was present.
 func (w *Wantlist) Contains(c cid.Cid) (Entry, bool) {
 	e, ok := w.set[c]
 	return e, ok
 }
 
-// Entries returns all wantlist entries for a want list.
+// Entries returns all wantlist entries for a want list, sorted by priority.
+//
+// DO NOT MODIFY. The returned list is cached.
 func (w *Wantlist) Entries() []Entry {
+	if w.cached != nil {
+		return w.cached
+	}
 	es := make([]Entry, 0, len(w.set))
 	for _, e := range w.set {
 		es = append(es, e)
 	}
-	return es
+	sort.Sort(entrySlice(es))
+	w.cached = es
+	return es[0:len(es):len(es)]
 }
 
 // Absorb all the entries in other into this want list
 func (w *Wantlist) Absorb(other *Wantlist) {
+	// Invalidate the cache up-front to avoid doing any work trying to keep it up-to-date.
+	w.cached = nil
+
 	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))
-}

wantlist/wantlist_test.go

@@ -5,6 +5,7 @@ import (
 
 	pb "github.com/ipfs/go-bitswap/message/pb"
 	cid "github.com/ipfs/go-cid"
+	"github.com/stretchr/testify/require"
 )
 
 var testcids []cid.Cid
@@ -211,11 +212,24 @@ func TestSortEntries(t *testing.T) {
 	wl.Add(testcids[2], 4, pb.Message_Wantlist_Have)
 
 	entries := wl.Entries()
-	SortEntries(entries)
-
 	if !entries[0].Cid.Equals(testcids[1]) ||
 		!entries[1].Cid.Equals(testcids[2]) ||
 		!entries[2].Cid.Equals(testcids[0]) {
 		t.Fatal("wrong order")
 	}
+
+}
+
+// Test adding and removing interleaved with checking entries to make sure we clear the cache.
+func TestCache(t *testing.T) {
+	wl := New()
+
+	wl.Add(testcids[0], 3, pb.Message_Wantlist_Block)
+	require.Len(t, wl.Entries(), 1)
+
+	wl.Add(testcids[1], 3, pb.Message_Wantlist_Block)
+	require.Len(t, wl.Entries(), 2)
+
+	wl.Remove(testcids[1])
+	require.Len(t, wl.Entries(), 1)
 }