Files
transmission/libtransmission/peer-mgr-wishlist.h

373 lines
10 KiB
C++

// This file Copyright © Mnemosyne LLC.
// It may be used under GPLv2 (SPDX: GPL-2.0-only), GPLv3 (SPDX: GPL-3.0-only),
// or any future license endorsed by Mnemosyne LLC.
// License text can be found in the licenses/ folder.
#pragma once
#ifndef LIBTRANSMISSION_PEER_MODULE
#error only the libtransmission peer module should #include this header.
#endif
#include <algorithm>
#include <compare>
#include <cstddef> // size_t
#include <functional>
#include <memory>
#include <vector>
#include <small/set.hpp>
#include "libtransmission/bitfield.h"
#include "libtransmission/crypto-utils.h"
#include "libtransmission/tr-macros.h"
#include "libtransmission/types.h"
/**
* Figures out what blocks we want to request next.
*/
class Wishlist
{
public:
struct Mediator
{
[[nodiscard]] virtual bool client_has_block(tr_block_index_t block) const = 0;
[[nodiscard]] virtual bool client_has_piece(tr_piece_index_t piece) const = 0;
[[nodiscard]] virtual bool client_wants_piece(tr_piece_index_t piece) const = 0;
[[nodiscard]] virtual bool is_sequential_download() const = 0;
[[nodiscard]] virtual tr_piece_index_t sequential_download_from_piece() const = 0;
[[nodiscard]] virtual size_t count_piece_replication(tr_piece_index_t piece) const = 0;
[[nodiscard]] virtual tr_block_span_t block_span(tr_piece_index_t piece) const = 0;
[[nodiscard]] virtual tr_piece_index_t piece_count() const = 0;
[[nodiscard]] virtual tr_priority_t priority(tr_piece_index_t piece) const = 0;
virtual ~Mediator() = default;
};
private:
struct Candidate
{
Candidate(tr_piece_index_t piece_in, tr_piece_index_t salt_in, Mediator const* mediator);
[[nodiscard]] constexpr auto operator<=>(Candidate const& that) const noexcept
{
// prefer pieces closer to completion
if (auto const val = std::size(unrequested) <=> std::size(that.unrequested); val != 0)
{
return val;
}
// prefer higher priority
if (auto const val = that.priority <=> priority; val != 0)
{
return val;
}
// prefer rarer pieces
if (auto const val = replication <=> that.replication; val != 0)
{
return val;
}
return salt <=> that.salt;
}
[[nodiscard]] constexpr auto operator==(Candidate const& that) const noexcept
{
return (*this <=> that) == 0;
}
[[nodiscard]] constexpr auto block_belongs(tr_block_index_t const block) const
{
return block_span.begin <= block && block < block_span.end;
}
tr_piece_index_t piece;
tr_block_span_t block_span;
tr_block_span_t raw_block_span;
// This is sorted in descending order so that smaller blocks indices
// can be taken from the end of the list, avoiding a move operation.
small::set<tr_block_index_t, small::default_inline_storage_v<tr_block_index_t>, std::greater<>> unrequested;
// Caching the following 2 values are highly beneficial, because:
// - they are often used (mainly because resort_piece() is called
// every time we receive a block)
// - does not change as often compared to missing blocks
// - calculating their values involves sifting through bitfield(s),
// which is expensive.
size_t replication;
tr_priority_t priority;
tr_piece_index_t salt;
};
using CandidateVec = std::vector<Candidate>;
public:
explicit Wishlist(Mediator& mediator_in)
: mediator_{ mediator_in }
{
candidate_list_upkeep();
}
void on_files_wanted_changed()
{
candidate_list_upkeep();
}
void on_got_bad_piece(tr_piece_index_t piece);
constexpr void on_got_bitfield(tr_bitfield const& bitfield)
{
inc_replication_bitfield(bitfield);
}
constexpr void on_got_block(tr_block_index_t const block)
{
if (auto const iter = find_by_block(block); iter != std::end(candidates_))
{
iter->unrequested.erase(block);
resort_piece(iter);
}
}
TR_CONSTEXPR_VEC void on_got_choke(tr_bitfield const& requests)
{
reset_blocks_bitfield(requests);
}
constexpr void on_got_have(tr_piece_index_t const piece)
{
if (auto const iter = find_by_piece(piece); iter != std::end(candidates_))
{
++iter->replication;
resort_piece(iter);
}
}
constexpr void on_got_have_all() noexcept
{
inc_replication();
}
constexpr void on_got_reject(tr_block_index_t const block)
{
reset_block(block);
}
TR_CONSTEXPR_VEC void on_peer_disconnect(tr_bitfield const& have, tr_bitfield const& requests)
{
dec_replication_bitfield(have);
reset_blocks_bitfield(requests);
}
constexpr void on_piece_completed(tr_piece_index_t const piece)
{
remove_piece(piece);
}
void on_priority_changed();
constexpr void on_sent_cancel(tr_block_index_t const block)
{
reset_block(block);
}
constexpr void on_sent_request(tr_block_span_t const block_span)
{
requested_block_span(block_span);
}
void on_sequential_download_changed()
{
recalculate_salt();
}
void on_sequential_download_from_piece_changed()
{
recalculate_salt();
}
// the next blocks that we should request from a peer
[[nodiscard]] std::vector<tr_block_span_t> next(
size_t n_wanted_blocks,
std::function<bool(tr_piece_index_t)> const& peer_has_piece);
private:
constexpr void dec_replication() noexcept
{
std::ranges::for_each(candidates_, [](Candidate& candidate) { --candidate.replication; });
}
constexpr void dec_replication_bitfield(tr_bitfield const& bitfield)
{
if (bitfield.has_none())
{
return;
}
if (bitfield.has_all())
{
dec_replication();
return;
}
for (auto& candidate : candidates_)
{
if (bitfield.test(candidate.piece))
{
--candidate.replication;
}
}
std::ranges::sort(candidates_);
}
constexpr void inc_replication() noexcept
{
std::ranges::for_each(candidates_, [](Candidate& candidate) { ++candidate.replication; });
}
constexpr void inc_replication_bitfield(tr_bitfield const& bitfield)
{
if (bitfield.has_none())
{
return;
}
if (bitfield.has_all())
{
inc_replication();
return;
}
for (auto& candidate : candidates_)
{
if (bitfield.test(candidate.piece))
{
++candidate.replication;
}
}
std::ranges::sort(candidates_);
}
// ---
constexpr void requested_block_span(tr_block_span_t const block_span)
{
for (auto block = block_span.begin; block < block_span.end;)
{
auto const it_p = find_by_block(block);
if (it_p == std::end(candidates_))
{
// std::unreachable();
break;
}
auto& unreq = it_p->unrequested;
auto it_b_end = std::end(unreq);
it_b_end = *std::prev(it_b_end) >= block_span.begin ? it_b_end : unreq.upper_bound(block_span.begin);
auto it_b_begin = std::begin(unreq);
it_b_begin = *it_b_begin < block_span.end ? it_b_begin : unreq.upper_bound(block_span.end);
unreq.erase(it_b_begin, it_b_end);
block = it_p->block_span.end;
resort_piece(it_p);
}
}
constexpr void reset_block(tr_block_index_t block)
{
if (auto const it_p = find_by_block(block); it_p != std::end(candidates_))
{
it_p->unrequested.insert(block);
resort_piece(it_p);
}
}
TR_CONSTEXPR_VEC void reset_blocks_bitfield(tr_bitfield const& requests)
{
for (auto& candidate : candidates_)
{
auto const [begin, end] = candidate.block_span;
if (requests.count(begin, end) == 0U)
{
continue;
}
for (auto i = end; i > begin; --i)
{
if (auto const block = i - 1U; requests.test(block))
{
candidate.unrequested.insert(block);
}
}
}
std::ranges::sort(candidates_);
}
// ---
[[nodiscard]] constexpr CandidateVec::iterator find_by_piece(tr_piece_index_t const piece)
{
return std::ranges::find_if(candidates_, [piece](auto const& c) { return c.piece == piece; });
}
[[nodiscard]] constexpr CandidateVec::iterator find_by_block(tr_block_index_t const block)
{
return std::ranges::find_if(candidates_, [block](auto const& c) { return c.block_belongs(block); });
}
tr_piece_index_t get_salt(tr_piece_index_t piece);
// ---
void candidate_list_upkeep();
// ---
constexpr void remove_piece(tr_piece_index_t const piece)
{
if (auto const iter = find_by_piece(piece); iter != std::end(candidates_))
{
candidates_.erase(iter);
}
}
// ---
void recalculate_salt();
// ---
constexpr void resort_piece(CandidateVec::iterator const& pos_old)
{
auto const pos_begin = std::begin(candidates_);
// Candidate needs to be moved towards the front of the list
if (auto const pos_next = std::next(pos_old); pos_old > pos_begin && *pos_old < *std::prev(pos_old))
{
auto const pos_new = std::lower_bound(pos_begin, pos_old, *pos_old);
std::rotate(pos_new, pos_old, pos_next);
}
// Candidate needs to be moved towards the end of the list
else if (auto const pos_end = std::end(candidates_); pos_next < pos_end && *pos_next < *pos_old)
{
auto const pos_new = std::lower_bound(pos_next, pos_end, *pos_old);
std::rotate(pos_old, pos_next, pos_new);
}
}
CandidateVec candidates_;
tr_salt_shaker<tr_piece_index_t> salter_ = {};
Mediator& mediator_;
};