fix: update wishlist when files wanted changed (#7733)

* refactor: extract salt calculation to method

* fix: update wishlist when files wanted changed

libtransmission/peer-mgr-wishlist.cc

@@ -332,6 +332,31 @@ private:
             [block](auto const& c) { return c.block_span.begin <= block && block < c.block_span.end; });
     }
 
+    static constexpr tr_piece_index_t get_salt(
+        tr_piece_index_t const piece,
+        tr_piece_index_t const n_pieces,
+        tr_piece_index_t const random_salt,
+        bool const is_sequential)
+    {
+        if (!is_sequential)
+        {
+            return random_salt;
+        }
+
+        // Download first and last piece first
+        if (piece == 0U)
+        {
+            return 0U;
+        }
+
+        if (piece == n_pieces - 1U)
+        {
+            return 1U;
+        }
+
+        return piece + 1U;
+    }
+
     void maybe_rebuild_candidate_list()
     {
         if (!candidates_dirty_)
@@ -352,31 +377,58 @@ private:
                 continue;
             }
 
-            auto const salt = [&]()
+            auto const salt = get_salt(piece, n_pieces, salter(), is_sequential);
-            {
-                if (!is_sequential)
-                {
-                    return salter();
-                }
-
-                // Download first and last piece first
-                if (piece == 0U)
-                {
-                    return 0U;
-                }
-
-                if (piece == n_pieces - 1U)
-                {
-                    return 1U;
-                }
-
-                return piece + 1U;
-            }();
             candidates_.emplace_back(piece, salt, &mediator_);
         }
         std::sort(std::begin(candidates_), std::end(candidates_));
     }
 
+    // ---
+
+    void recalculate_wanted_pieces()
+    {
+        if (candidates_dirty_)
+        {
+            return;
+        }
+
+        auto n_old_c = std::size(candidates_);
+        auto salter = tr_salt_shaker<tr_piece_index_t>{};
+        auto const is_sequential = mediator_.is_sequential_download();
+        auto const n_pieces = mediator_.piece_count();
+
+        std::sort(
+            std::begin(candidates_),
+            std::end(candidates_),
+            [](auto const& lhs, auto const& rhs) { return lhs.piece < rhs.piece; });
+
+        for (tr_piece_index_t piece = 0U, idx_c = 0U; piece < n_pieces; ++piece)
+        {
+            auto const existing_candidate = idx_c < n_old_c && piece == candidates_[idx_c].piece;
+            if (mediator_.client_wants_piece(piece))
+            {
+                if (existing_candidate)
+                {
+                    ++idx_c;
+                }
+                else
+                {
+                    auto const salt = get_salt(piece, n_pieces, salter(), is_sequential);
+                    candidates_.emplace_back(piece, salt, &mediator_);
+                }
+            }
+            else if (existing_candidate)
+            {
+                candidates_.erase(std::next(std::begin(candidates_), idx_c));
+                --n_old_c;
+            }
+        }
+
+        std::sort(std::begin(candidates_), std::end(candidates_));
+    }
+
+    // ---
+
     TR_CONSTEXPR20 void remove_piece(tr_piece_index_t const piece)
     {
         if (candidates_dirty_)
@@ -417,13 +469,15 @@ private:
     bool candidates_dirty_ = true;
     bool is_endgame_ = false;
 
-    std::array<libtransmission::ObserverTag, 13U> const tags_;
+    std::array<libtransmission::ObserverTag, 14U> const tags_;
 
     Mediator& mediator_;
 };
 
 Wishlist::Impl::Impl(Mediator& mediator_in)
     : tags_{ {
+          mediator_in.observe_files_wanted_changed([this](tr_torrent*, tr_file_index_t const*, tr_file_index_t, bool)
+                                                   { recalculate_wanted_pieces(); }),
           mediator_in.observe_peer_disconnect([this](tr_torrent*, tr_bitfield const& b, tr_bitfield const& ar)
                                               { peer_disconnect(b, ar); }),
           mediator_in.observe_got_bad_piece([this](tr_torrent*, tr_piece_index_t) { set_candidates_dirty(); }),

libtransmission/peer-mgr-wishlist.h

@@ -43,6 +43,8 @@ public:
         [[nodiscard]] virtual tr_piece_index_t piece_count() const = 0;
         [[nodiscard]] virtual tr_priority_t priority(tr_piece_index_t piece) const = 0;
 
+        [[nodiscard]] virtual libtransmission::ObserverTag observe_files_wanted_changed(
+            libtransmission::SimpleObservable<tr_torrent*, tr_file_index_t const*, tr_file_index_t, bool>::Observer) = 0;
         [[nodiscard]] virtual libtransmission::ObserverTag observe_peer_disconnect(
             libtransmission::SimpleObservable<tr_torrent*, tr_bitfield const&, tr_bitfield const&>::Observer observer) = 0;
         [[nodiscard]] virtual libtransmission::ObserverTag observe_got_bad_piece(

libtransmission/peer-mgr.cc

@@ -321,6 +321,9 @@ public:
         [[nodiscard]] tr_piece_index_t piece_count() const override;
         [[nodiscard]] tr_priority_t priority(tr_piece_index_t piece) const override;
 
+        [[nodiscard]] libtransmission::ObserverTag observe_files_wanted_changed(
+            libtransmission::SimpleObservable<tr_torrent*, tr_file_index_t const*, tr_file_index_t, bool>::Observer observer)
+            override;
         [[nodiscard]] libtransmission::ObserverTag observe_peer_disconnect(
             libtransmission::SimpleObservable<tr_torrent*, tr_bitfield const&, tr_bitfield const&>::Observer observer) override;
         [[nodiscard]] libtransmission::ObserverTag observe_got_bad_piece(
@@ -986,6 +989,12 @@ tr_priority_t tr_swarm::WishlistMediator::priority(tr_piece_index_t piece) const
     return tor_.piece_priority(piece);
 }
 
+libtransmission::ObserverTag tr_swarm::WishlistMediator::observe_files_wanted_changed(
+    libtransmission::SimpleObservable<tr_torrent*, tr_file_index_t const*, tr_file_index_t, bool>::Observer observer)
+{
+    return tor_.files_wanted_changed_.observe(std::move(observer));
+}
+
 libtransmission::ObserverTag tr_swarm::WishlistMediator::observe_peer_disconnect(
     libtransmission::SimpleObservable<tr_torrent*, tr_bitfield const&, tr_bitfield const&>::Observer observer)
 {

libtransmission/torrent.h

@@ -992,6 +992,7 @@ struct tr_torrent
     libtransmission::SimpleObservable<tr_torrent*> started_;
     libtransmission::SimpleObservable<tr_torrent*> stopped_;
     libtransmission::SimpleObservable<tr_torrent*> swarm_is_all_upload_only_;
+    libtransmission::SimpleObservable<tr_torrent*, tr_file_index_t const*, tr_file_index_t, bool> files_wanted_changed_;
     libtransmission::SimpleObservable<tr_torrent*, tr_file_index_t const*, tr_file_index_t, tr_priority_t> priority_changed_;
     libtransmission::SimpleObservable<tr_torrent*, bool> sequential_download_changed_;
 
@@ -1223,6 +1224,7 @@ private:
 
         files_wanted_.set(files, n_files, wanted);
         completion_.invalidate_size_when_done();
+        files_wanted_changed_.emit(this, files, n_files, wanted);
 
         if (!is_bootstrapping)
         {

tests/libtransmission/peer-mgr-wishlist-test.cc

@@ -85,6 +85,13 @@ protected:
             return piece_priority_[piece];
         }
 
+        [[nodiscard]] libtransmission::ObserverTag observe_files_wanted_changed(
+            libtransmission::SimpleObservable<tr_torrent*, tr_file_index_t const*, tr_file_index_t, bool>::Observer observer)
+            override
+        {
+            return parent_.files_wanted_changed_.observe(std::move(observer));
+        }
+
         [[nodiscard]] libtransmission::ObserverTag observe_peer_disconnect(
             libtransmission::SimpleObservable<tr_torrent*, tr_bitfield const&, tr_bitfield const&>::Observer observer) override
         {
@@ -165,6 +172,7 @@ protected:
         }
     };
 
+    libtransmission::SimpleObservable<tr_torrent*, tr_file_index_t const*, tr_file_index_t, bool> files_wanted_changed_;
     libtransmission::SimpleObservable<tr_torrent*, tr_bitfield const&, tr_bitfield const&> peer_disconnect_;
     libtransmission::SimpleObservable<tr_torrent*, tr_piece_index_t> got_bad_piece_;
     libtransmission::SimpleObservable<tr_torrent*, tr_bitfield const&> got_bitfield_;
@@ -1546,3 +1554,120 @@ TEST_F(PeerMgrWishlistTest, settingSequentialDownloadRebuildsWishlist)
         EXPECT_EQ(100U, requested.count(200, 300));
     }
 }
+
+TEST_F(PeerMgrWishlistTest, setFileWantedUpdatesCandidateListAdd)
+{
+    auto const get_spans = [this](size_t n_wanted)
+    {
+        auto mediator = MockMediator{ *this };
+
+        // setup: four pieces, all missing
+        mediator.piece_count_ = 4;
+        mediator.block_span_[0] = { 0, 100 };
+        mediator.block_span_[1] = { 100, 200 };
+        mediator.block_span_[2] = { 200, 300 };
+        mediator.block_span_[3] = { 300, 400 };
+
+        // peer has all pieces
+        mediator.piece_replication_[0] = 1;
+        mediator.piece_replication_[1] = 1;
+        mediator.piece_replication_[2] = 1;
+        mediator.piece_replication_[3] = 1;
+
+        // we initially only want the first 2 pieces
+        mediator.client_wants_piece_.insert(0);
+        mediator.client_wants_piece_.insert(1);
+
+        // allow the wishlist to build its cache
+        auto wishlist = Wishlist{ mediator };
+        (void)wishlist.next(1, PeerHasAllPieces, ClientHasNoActiveRequests);
+
+        // now we want the file that consists of piece 2 and piece 3 also
+        mediator.client_wants_piece_.insert(2);
+        mediator.client_wants_piece_.insert(3);
+        files_wanted_changed_.emit(nullptr, nullptr, 0, true);
+
+        // a candidate should be inserted into the wishlist for
+        // piece 2 and piece 3
+        return wishlist.next(n_wanted, PeerHasAllPieces, ClientHasNoActiveRequests);
+    };
+
+    // We should request all 4 pieces here.
+    // NB: when all other things are equal in the wishlist, pieces are
+    // picked at random so this test -could- pass even if there's a bug.
+    // So test several times to shake out any randomness
+    static auto constexpr NumRuns = 1000;
+    for (int run = 0; run < NumRuns; ++run)
+    {
+        auto requested = tr_bitfield{ 400 };
+        auto const spans = get_spans(350);
+        for (auto const& [begin, end] : spans)
+        {
+            requested.set_span(begin, end);
+        }
+        EXPECT_EQ(350U, requested.count());
+        EXPECT_NE(0U, requested.count(0, 100));
+        EXPECT_NE(0U, requested.count(100, 200));
+        EXPECT_NE(0U, requested.count(200, 300));
+        EXPECT_NE(0U, requested.count(300, 400));
+    }
+}
+
+TEST_F(PeerMgrWishlistTest, setFileWantedUpdatesCandidateListRemove)
+{
+    auto const get_spans = [this](size_t n_wanted)
+    {
+        auto mediator = MockMediator{ *this };
+
+        // setup: four pieces, all missing
+        mediator.piece_count_ = 4;
+        mediator.block_span_[0] = { 0, 100 };
+        mediator.block_span_[1] = { 100, 200 };
+        mediator.block_span_[2] = { 200, 300 };
+        mediator.block_span_[3] = { 300, 400 };
+
+        // peer has all pieces
+        mediator.piece_replication_[0] = 1;
+        mediator.piece_replication_[1] = 1;
+        mediator.piece_replication_[2] = 1;
+        mediator.piece_replication_[3] = 1;
+
+        // we initially only want the all 4 pieces
+        mediator.client_wants_piece_.insert(0);
+        mediator.client_wants_piece_.insert(1);
+        mediator.client_wants_piece_.insert(2);
+        mediator.client_wants_piece_.insert(3);
+
+        // allow the wishlist to build its cache
+        auto wishlist = Wishlist{ mediator };
+        (void)wishlist.next(1, PeerHasAllPieces, ClientHasNoActiveRequests);
+
+        // we no longer want the file that consists of piece 2 and piece 3
+        mediator.client_wants_piece_.erase(2);
+        mediator.client_wants_piece_.erase(3);
+        files_wanted_changed_.emit(nullptr, nullptr, 0, true);
+
+        // the candidate objects for piece 2 and piece 3 should be removed
+        return wishlist.next(n_wanted, PeerHasAllPieces, ClientHasNoActiveRequests);
+    };
+
+    // We should request only the first 2 pieces here.
+    // NB: when all other things are equal in the wishlist, pieces are
+    // picked at random so this test -could- pass even if there's a bug.
+    // So test several times to shake out any randomness
+    static auto constexpr NumRuns = 1000;
+    for (int run = 0; run < NumRuns; ++run)
+    {
+        auto requested = tr_bitfield{ 400 };
+        auto const spans = get_spans(350);
+        for (auto const& [begin, end] : spans)
+        {
+            requested.set_span(begin, end);
+        }
+        EXPECT_EQ(200U, requested.count());
+        EXPECT_NE(0U, requested.count(0, 100));
+        EXPECT_NE(0U, requested.count(100, 200));
+        EXPECT_EQ(0U, requested.count(200, 300));
+        EXPECT_EQ(0U, requested.count(300, 400));
+    }
+}