From e8971ea0f1bfb26e2e05eea6e98a547dfb231bb9 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Sun, 19 Feb 2023 13:24:25 +1100
Subject: [PATCH] Add unit test for relative availability after HaveNone

Could help with https://github.com/anacrolix/torrent/issues/813.
---
 peerconn.go          |  2 +-
 prioritized_peers.go |  3 +++
 torrent.go           | 14 +++++++++-----
 torrent_test.go      | 36 ++++++++++++++++++++++++++++++++++++
 4 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/peerconn.go b/peerconn.go
index af9c0844..5b580bd6 100644
--- a/peerconn.go
+++ b/peerconn.go
@@ -899,7 +899,7 @@ func (cn *PeerConn) onPeerSentHaveAll() error {
 }
 
 func (cn *PeerConn) peerSentHaveNone() error {
-	if cn.peerSentHaveAll {
+	if !cn.peerSentHaveAll {
 		cn.t.decPeerPieceAvailability(&cn.Peer)
 	}
 	cn._peerPieces.Clear()
diff --git a/prioritized_peers.go b/prioritized_peers.go
index d0ef43ed..443d7207 100644
--- a/prioritized_peers.go
+++ b/prioritized_peers.go
@@ -45,6 +45,9 @@ func (me *prioritizedPeers) Each(f func(PeerInfo)) {
 }
 
 func (me *prioritizedPeers) Len() int {
+	if me == nil || me.om == nil {
+		return 0
+	}
 	return me.om.Len()
 }
 
diff --git a/torrent.go b/torrent.go
index f19deb75..637a9897 100644
--- a/torrent.go
+++ b/torrent.go
@@ -894,17 +894,21 @@ func (t *Torrent) close(wg *sync.WaitGroup) (err error) {
 	if t.storage != nil {
 		t.deletePieceRequestOrder()
 	}
+	t.assertAllPiecesRelativeAvailabilityZero()
+	t.pex.Reset()
+	t.cl.event.Broadcast()
+	t.pieceStateChanges.Close()
+	t.updateWantPeersEvent()
+	return
+}
+
+func (t *Torrent) assertAllPiecesRelativeAvailabilityZero() {
 	for i := range t.pieces {
 		p := t.piece(i)
 		if p.relativeAvailability != 0 {
 			panic(fmt.Sprintf("piece %v has relative availability %v", i, p.relativeAvailability))
 		}
 	}
-	t.pex.Reset()
-	t.cl.event.Broadcast()
-	t.pieceStateChanges.Close()
-	t.updateWantPeersEvent()
-	return
 }
 
 func (t *Torrent) requestOffset(r Request) int64 {
diff --git a/torrent_test.go b/torrent_test.go
index ce0af9f4..d2c7ba72 100644
--- a/torrent_test.go
+++ b/torrent_test.go
@@ -6,10 +6,14 @@ import (
 	"net"
 	"os"
 	"path/filepath"
+	"sync"
 	"testing"
 
+	"github.com/anacrolix/generics"
+	"github.com/anacrolix/log"
 	"github.com/anacrolix/missinggo/v2"
 	"github.com/anacrolix/missinggo/v2/bitmap"
+	qt "github.com/frankban/quicktest"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -215,3 +219,35 @@ func TestTorrentMetainfoIncompleteMetadata(t *testing.T) {
 	assert.False(t, tt.haveAllMetadataPieces())
 	assert.Nil(t, tt.Metainfo().InfoBytes)
 }
+
+func TestRelativeAvailabilityHaveNone(t *testing.T) {
+	c := qt.New(t)
+	var err error
+	cl := Client{
+		config: TestingConfig(t),
+	}
+	tt := Torrent{
+		cl:           &cl,
+		logger:       log.Default,
+		gotMetainfoC: make(chan struct{}),
+	}
+	tt.setChunkSize(2)
+	generics.MakeMapIfNil(&tt.conns)
+	pc := PeerConn{}
+	pc.t = &tt
+	pc.peerImpl = &pc
+	pc.initRequestState()
+	generics.InitNew(&pc.callbacks)
+	tt.conns[&pc] = struct{}{}
+	err = pc.peerSentHave(0)
+	c.Assert(err, qt.IsNil)
+	info := testutil.Greeting.Info(5)
+	err = tt.setInfo(&info)
+	c.Assert(err, qt.IsNil)
+	tt.onSetInfo()
+	err = pc.peerSentHaveNone()
+	c.Assert(err, qt.IsNil)
+	var wg sync.WaitGroup
+	tt.close(&wg)
+	tt.assertAllPiecesRelativeAvailabilityZero()
+}