Merge pull request #30 from libp2p/rework-nearest-peers

Remove a lot of allocations, and fix some ambiguous naming

sorting.go

@@ -13,48 +13,49 @@ type peerDistance struct {
 	distance ID
 }
 
-// peerSorterArr implements sort.Interface to sort peers by xor distance
-type peerSorterArr []*peerDistance
+// peerDistanceSorter implements sort.Interface to sort peers by xor distance
+type peerDistanceSorter struct {
+	peers  []peerDistance
+	target ID
+}
 
-func (p peerSorterArr) Len() int      { return len(p) }
-func (p peerSorterArr) Swap(a, b int) { p[a], p[b] = p[b], p[a] }
-func (p peerSorterArr) Less(a, b int) bool {
-	return p[a].distance.less(p[b].distance)
+func (pds *peerDistanceSorter) Len() int      { return len(pds.peers) }
+func (pds *peerDistanceSorter) Swap(a, b int) { pds.peers[a], pds.peers[b] = pds.peers[b], pds.peers[a] }
+func (pds *peerDistanceSorter) Less(a, b int) bool {
+	return pds.peers[a].distance.less(pds.peers[b].distance)
 }
 
-//
+// Append the peer.ID to the sorter's slice. It may no longer be sorted.
+func (pds *peerDistanceSorter) appendPeer(p peer.ID) {
+	pds.peers = append(pds.peers, peerDistance{
+		p:        p,
+		distance: xor(pds.target, ConvertPeerID(p)),
+	})
+}
 
-func copyPeersFromList(target ID, peerArr peerSorterArr, peerList *list.List) peerSorterArr {
-	if cap(peerArr) < len(peerArr)+peerList.Len() {
-		newArr := make(peerSorterArr, len(peerArr), len(peerArr)+peerList.Len())
-		copy(newArr, peerArr)
-		peerArr = newArr
-	}
-	for e := peerList.Front(); e != nil; e = e.Next() {
-		p := e.Value.(peer.ID)
-		pID := ConvertPeerID(p)
-		pd := peerDistance{
-			p:        p,
-			distance: xor(target, pID),
-		}
-		peerArr = append(peerArr, &pd)
+// Append the peer.ID values in the list to the sorter's slice. It may no longer be sorted.
+func (pds *peerDistanceSorter) appendPeersFromList(l *list.List) {
+	for e := l.Front(); e != nil; e = e.Next() {
+		pds.appendPeer(e.Value.(peer.ID))
 	}
-	return peerArr
 }
 
+func (pds *peerDistanceSorter) sort() {
+	sort.Sort(pds)
+}
+
+// Sort the given peers by their ascending distance from the target. A new slice is returned.
 func SortClosestPeers(peers []peer.ID, target ID) []peer.ID {
-	psarr := make(peerSorterArr, 0, len(peers))
+	sorter := peerDistanceSorter{
+		peers:  make([]peerDistance, 0, len(peers)),
+		target: target,
+	}
 	for _, p := range peers {
-		pID := ConvertPeerID(p)
-		pd := &peerDistance{
-			p:        p,
-			distance: xor(target, pID),
-		}
-		psarr = append(psarr, pd)
+		sorter.appendPeer(p)
 	}
-	sort.Sort(psarr)
-	out := make([]peer.ID, 0, len(psarr))
-	for _, p := range psarr {
+	sorter.sort()
+	out := make([]peer.ID, 0, sorter.Len())
+	for _, p := range sorter.peers {
 		out = append(out, p.p)
 	}
 	return out

table.go

@@ -4,7 +4,6 @@ package kbucket
 import (
 	"errors"
 	"fmt"
-	"sort"
 	"sync"
 	"time"
 
@@ -169,6 +168,7 @@ func (rt *RoutingTable) NearestPeer(id ID) peer.ID {
 func (rt *RoutingTable) NearestPeers(id ID, count int) []peer.ID {
 	cpl := CommonPrefixLen(id, rt.local)
 
+	// It's assumed that this also protects the buckets.
 	rt.tabLock.RLock()
 
 	// Get bucket at cpl index or last bucket
@@ -178,32 +178,32 @@ func (rt *RoutingTable) NearestPeers(id ID, count int) []peer.ID {
 	}
 	bucket = rt.Buckets[cpl]
 
-	peerArr := make(peerSorterArr, 0, count)
-	peerArr = copyPeersFromList(id, peerArr, bucket.list)
-	if len(peerArr) < count {
+	pds := peerDistanceSorter{
+		peers:  make([]peerDistance, 0, 3*rt.bucketsize),
+		target: id,
+	}
+	pds.appendPeersFromList(bucket.list)
+	if pds.Len() < count {
 		// In the case of an unusual split, one bucket may be short or empty.
 		// if this happens, search both surrounding buckets for nearby peers
 		if cpl > 0 {
-			plist := rt.Buckets[cpl-1].list
-			peerArr = copyPeersFromList(id, peerArr, plist)
+			pds.appendPeersFromList(rt.Buckets[cpl-1].list)
 		}
-
 		if cpl < len(rt.Buckets)-1 {
-			plist := rt.Buckets[cpl+1].list
-			peerArr = copyPeersFromList(id, peerArr, plist)
+			pds.appendPeersFromList(rt.Buckets[cpl+1].list)
 		}
 	}
 	rt.tabLock.RUnlock()
 
 	// Sort by distance to local peer
-	sort.Sort(peerArr)
+	pds.sort()
 
-	if count < len(peerArr) {
-		peerArr = peerArr[:count]
+	if count < pds.Len() {
+		pds.peers = pds.peers[:count]
 	}
 
-	out := make([]peer.ID, 0, len(peerArr))
-	for _, p := range peerArr {
+	out := make([]peer.ID, 0, pds.Len())
+	for _, p := range pds.peers {
 		out = append(out, p.p)
 	}