fix optimistic unchoke

internal/unchoker/unchoker.go

@@ -1,13 +1,19 @@
 package unchoker
 
-import "sort"
+import (
+	"math/rand"
+	"sort"
+)
 
 type Unchoker struct {
 	getPeers              func() []Peer
 	torrentCompleted      *bool
 	numUnchoked           int
 	numOptimisticUnchoked int
 
+	// Every 3rd round an optimistic unchoke logic is applied.
+	round uint8
+
 	peersUnchoked           map[Peer]struct{}
 	peersUnchokedOptimistic map[Peer]struct{}
 }
@@ -59,17 +65,6 @@ func (u *Unchoker) candidatesUnchoke() []Peer {
 	return peers
 }
 
-func (u *Unchoker) candidatesUnchokeOptimistic() []Peer {
-	allPeers := u.getPeers()
-	peers := make([]Peer, 0, len(allPeers))
-	for _, pe := range allPeers {
-		if pe.Interested() && (pe.Choking() || pe.Optimistic()) {
-			peers = append(peers, pe)
-		}
-	}
-	return peers
-}
-
 func (u *Unchoker) sortPeers(peers []Peer) {
 	byUploadSpeed := func(i, j int) bool { return peers[i].UploadSpeed() > peers[j].UploadSpeed() }
 	byDownloadSpeed := func(i, j int) bool { return peers[i].DownloadSpeed() > peers[j].DownloadSpeed() }
@@ -80,31 +75,32 @@ func (u *Unchoker) sortPeers(peers []Peer) {
 	}
 }
 
+// TickUnchoke must be called at every 10 seconds.
 func (u *Unchoker) TickUnchoke() {
+	optimistic := u.round == 0
 	peers := u.candidatesUnchoke()
 	u.sortPeers(peers)
-	var unchoked int
-	for _, pe := range peers {
-		if unchoked < u.numUnchoked {
-			u.unchokePeer(pe)
-			unchoked++
-		} else {
-			u.chokePeer(pe)
+	var i, unchoked int
+	for ; i < len(peers) && unchoked < u.numUnchoked; i++ {
+		if !optimistic && peers[i].Optimistic() {
+			continue
 		}
+		u.unchokePeer(peers[i])
+		unchoked++
 	}
-}
-
-func (u *Unchoker) TickOptimisticUnchoke() {
-	peers := u.candidatesUnchokeOptimistic()
-	var unchoked int
-	for _, pe := range peers {
-		if unchoked < u.numOptimisticUnchoked {
+	peers = peers[i:]
+	if optimistic {
+		for i = 0; i < u.numOptimisticUnchoked && len(peers) > 0; i++ {
+			n := rand.Intn(len(peers))
+			pe := peers[n]
 			u.optimisticUnchokePeer(pe)
-			unchoked++
-		} else {
-			u.chokePeer(pe)
+			peers[n], peers = peers[len(peers)-1], peers[:len(peers)-1]
 		}
 	}
+	for _, pe := range peers {
+		u.chokePeer(pe)
+	}
+	u.round = (u.round + 1) % 3
 }
 
 func (u *Unchoker) chokePeer(pe Peer) {
@@ -134,7 +130,10 @@ func (u *Unchoker) unchokePeer(pe Peer) {
 func (u *Unchoker) optimisticUnchokePeer(pe Peer) {
 	if !pe.Choking() {
 		if !pe.Optimistic() {
-			panic("must not optimistic unchoke already unchoked peer")
+			// Move into optimistic unchoked peers
+			pe.SetOptimistic(true)
+			delete(u.peersUnchoked, pe)
+			u.peersUnchokedOptimistic[pe] = struct{}{}
 		}
 		return
 	}

internal/unchoker/unchoker_test.go

@@ -36,6 +36,7 @@ func TestTickUnchoke(t *testing.T) {
 	u := New(getPeers, new(bool), 2, 1)
 
 	// Must unchoke fastest downloading 2 peers
+	u.round = 1
 	u.TickUnchoke()
 	assert.Equal(t, []*TestPeer{
 		{
@@ -56,6 +57,7 @@ func TestTickUnchoke(t *testing.T) {
 	}, testPeers)
 
 	// Nothing has changed. Same peers stays unchoked.
+	u.round = 1
 	u.TickUnchoke()
 	assert.Equal(t, []*TestPeer{
 		{
@@ -76,7 +78,8 @@ func TestTickUnchoke(t *testing.T) {
 	}, testPeers)
 
 	// First choked peer must be unchoked optimistically.
-	u.TickOptimisticUnchoke()
+	u.round = 0
+	u.TickUnchoke()
 	assert.Equal(t, []*TestPeer{
 		{
 			interested: true,
@@ -96,18 +99,39 @@ func TestTickUnchoke(t *testing.T) {
 	}, testPeers)
 
 	// Optimistically unchoked peer has started downloading,
-	// regular unchoke must be applied on that peer.
 	testPeers[0].downloadSpeed = 3
+	u.round = 1
 	u.TickUnchoke()
 	assert.Equal(t, []*TestPeer{
 		{
+			interested:    true,
+			optimistic:    true,
 			downloadSpeed: 3,
+		},
+		{
 			interested:    true,
+			downloadSpeed: 2,
+		},
+		{
+			interested:    true,
+			downloadSpeed: 4,
+		},
+		{
+			choking: true,
+		},
+	}, testPeers)
+
+	u.round = 0
+	u.TickUnchoke()
+	assert.Equal(t, []*TestPeer{
+		{
+			interested:    true,
+			downloadSpeed: 3,
 		},
 		{
 			interested:    true,
 			downloadSpeed: 2,
-			choking:       true,
+			optimistic:    true,
 		},
 		{
 			interested:    true,

torrent/run.go

@@ -42,9 +42,6 @@ func (t *torrent) run() {
 	t.unchokeTicker = time.NewTicker(10 * time.Second)
 	defer t.unchokeTicker.Stop()
 
-	t.optimisticUnchokeTicker = time.NewTicker(30 * time.Second)
-	defer t.optimisticUnchokeTicker.Stop()
-
 	for {
 		select {
 		case <-t.closeC:
@@ -209,8 +206,6 @@ func (t *torrent) run() {
 			}
 		case <-t.unchokeTicker.C:
 			t.unchoker.TickUnchoke()
-		case <-t.optimisticUnchokeTicker.C:
-			t.unchoker.TickOptimisticUnchoke()
 		case ih := <-t.incomingHandshakerResultC:
 			delete(t.incomingHandshakers, ih)
 			if ih.Error != nil {

torrent/torrent.go

@@ -196,9 +196,6 @@ type torrent struct {
 	// A ticker that ticks periodically to keep a certain number of peers unchoked.
 	unchokeTicker *time.Ticker
 
-	// A ticker that ticks periodically to keep a random peer unchoked regardless of its upload rate.
-	optimisticUnchokeTicker *time.Ticker
-
 	// A worker that opens and allocates files on the disk.
 	allocator          *allocator.Allocator
 	allocatorProgressC chan allocator.Progress