refactor torrent._rechoke()

1 files changed, 58 insertions, 50 deletions

diff --git a/lib/torrent.js b/lib/torrent.js
index 6e98725..3be12af 100644
--- a/lib/torrent.js
+++ b/lib/torrent.js
@@ -1348,71 +1348,79 @@ class Torrent extends EventEmitter {
   _rechoke () {
     if (!this.ready) return
 
-    if (this._rechokeOptimisticTime > 0) this._rechokeOptimisticTime -= 1
-    else this._rechokeOptimisticWire = null
+    // wires in increasing order of quality (pop() gives next best peer)
+    const wireStack =
+      this.wires
+        .map(wire => ({ wire, random: Math.random() })) // insert a random seed for randomizing the sort
+        .sort((objA, objB) => {
+          const wireA = objA.wire
+          const wireB = objB.wire
+
+          // prefer peers that send us data faster
+          if (wireA.downloadSpeed() !== wireB.downloadSpeed()) {
+            return wireA.downloadSpeed() - wireB.downloadSpeed()
+          }
 
-    const peers = []
+          // then prefer peers that can download data from us faster
+          if (wireA.uploadSpeed() !== wireB.uploadSpeed()) {
+            return wireA.uploadSpeed() - wireB.uploadSpeed()
+          }
 
-    this.wires.forEach(wire => {
-      if (!wire.isSeeder && wire !== this._rechokeOptimisticWire) {
-        peers.push({
-          wire,
-          downloadSpeed: wire.downloadSpeed(),
-          uploadSpeed: wire.uploadSpeed(),
-          salt: Math.random(),
-          isChoked: true
-        })
-      }
-    })
+          // then prefer already unchoked peers (to minimize fibrillation)
+          if (wireA.amChoking !== wireB.amChoking) {
+            return wireA.amChoking ? -1 : 1 // choking < unchoked
+          }
 
-    peers.sort(rechokeSort)
+          // otherwise random order
+          return objA.random - objB.random
+        })
+        .map(obj => obj.wire) // return array of wires (remove random seed)
 
-    let unchokeInterested = 0
-    let i = 0
-    for (; i < peers.length && unchokeInterested < this._rechokeNumSlots; ++i) {
-      peers[i].isChoked = false
-      if (peers[i].wire.peerInterested) unchokeInterested += 1
+    if (this._rechokeOptimisticTime <= 0) {
+      // clear old optimistic peer, so it can be rechoked normally and then replaced
+      this._rechokeOptimisticWire = null
+    } else {
+      this._rechokeOptimisticTime -= 1
     }
 
-    // Optimistically unchoke a peer
-    if (!this._rechokeOptimisticWire && i < peers.length && this._rechokeNumSlots) {
-      const candidates = peers.slice(i).filter(peer => peer.wire.peerInterested)
-      const optimistic = candidates[randomInt(candidates.length)]
+    let numInterestedUnchoked = 0
+    // leave one rechoke slot open for optimistic unchoking
+    while (wireStack.length > 0 && numInterestedUnchoked < this._rechokeNumSlots - 1) {
+      const wire = wireStack.pop() // next best quality peer
 
-      if (optimistic) {
-        optimistic.isChoked = false
-        this._rechokeOptimisticWire = optimistic.wire
-        this._rechokeOptimisticTime = RECHOKE_OPTIMISTIC_DURATION
+      if (wire.isSeeder || wire === this._rechokeOptimisticWire) {
+        continue
       }
-    }
 
-    // Unchoke best peers
-    peers.forEach(peer => {
-      if (peer.wire.amChoking !== peer.isChoked) {
-        if (peer.isChoked) peer.wire.choke()
-        else peer.wire.unchoke()
-      }
-    })
+      wire.unchoke()
 
-    function rechokeSort (peerA, peerB) {
-      // Prefer higher download speed
-      if (peerA.downloadSpeed !== peerB.downloadSpeed) {
-        return peerB.downloadSpeed - peerA.downloadSpeed
+      // only stop unchoking once we fill the slots with interested peers that will actually download
+      if (wire.peerInterested) {
+        numInterestedUnchoked++
       }
+    }
 
-      // Prefer higher upload speed
-      if (peerA.uploadSpeed !== peerB.uploadSpeed) {
-        return peerB.uploadSpeed - peerA.uploadSpeed
-      }
+    // fill optimistic unchoke slot if empty
+    if (this._rechokeOptimisticWire === null && this._rechokeNumSlots > 0) {
+      // don't optimistically unchoke uninterested peers
+      const remaining = wireStack.filter(wire => wire.peerInterested)
 
-      // Prefer unchoked
-      if (peerA.wire.amChoking !== peerB.wire.amChoking) {
-        return peerA.wire.amChoking ? 1 : -1
-      }
+      if (remaining.length > 0) {
+        // select random remaining (not yet unchoked) peer
+        const newOptimisticPeer = remaining[randomInt(remaining.length)]
 
-      // Random order
-      return peerA.salt - peerB.salt
+        newOptimisticPeer.unchoke()
+
+        this._rechokeOptimisticWire = newOptimisticPeer
+
+        this._rechokeOptimisticTime = RECHOKE_OPTIMISTIC_DURATION
+      }
     }
+
+    // choke the rest
+    wireStack
+      .filter(wire => wire !== this._rechokeOptimisticWire) // except the optimistically unchoked peer
+      .forEach(wire => wire.choke())
   }
 
   /**