feat: cache the materialized wantlist

This can become a hot-spot.

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,10 +94,31 @@ 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) {
+	if w.cached != nil {
+		if _, had := w.set[c]; had {
+			// invalidate the cache.
+			w.cached = nil
+		} else {
+			// This is a small optimization that will keep the cache up-to-date instead
+			// of simply invalidating it.
+			//
+			// We intentionally _only_ do this if the cache is non-empty.
+			w.cached = append(w.cached, e)
+		}
+	}
+	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) {
@@ -104,15 +128,22 @@ func (w *Wantlist) Contains(c cid.Cid) (Entry, bool) {
 
 // Entries returns all wantlist entries for a want list.
 func (w *Wantlist) Entries() []Entry {
-	es := make([]Entry, 0, len(w.set))
+	if w.cached != nil {
+		return w.cached
+	}
+	es := make([]Entry, len(w.set))
 	for _, e := range w.set {
 		es = append(es, e)
 	}
-	return 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)
 	}