From 37d85ac2f7519809216e420766bb24f0cd572513 Mon Sep 17 00:00:00 2001
From: arvidn <arvid@libtorrent.org>
Date: Sun, 19 Aug 2018 11:08:09 +0200
Subject: [PATCH] improve piece picker performance in tracking pad-blocks

---
 include/libtorrent/piece_picker.hpp | 18 ++++++++++++++----
 src/piece_picker.cpp                | 27 +++++++++++++++++++++------
 2 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/include/libtorrent/piece_picker.hpp b/include/libtorrent/piece_picker.hpp
index faab92161..c9fa41c36 100644
--- a/include/libtorrent/piece_picker.hpp
+++ b/include/libtorrent/piece_picker.hpp
@@ -43,6 +43,7 @@ POSSIBILITY OF SUCH DAMAGE.
 #include <cstdint>
 #include <tuple>
 #include <set>
+#include <unordered_map>
 
 #include "libtorrent/peer_id.hpp"
 #include "libtorrent/assert.hpp"
@@ -464,7 +465,7 @@ namespace libtorrent {
 
 	private:
 
-		int num_pad_blocks() const { return int(m_pad_blocks.size()); }
+		int num_pad_blocks() const { return m_num_pad_blocks; }
 
 		aux::typed_span<block_info> mutable_blocks_for_piece(downloading_piece const& dp);
 
@@ -725,10 +726,19 @@ namespace libtorrent {
 		// TODO: should this be allocated lazily?
 		mutable aux::vector<piece_pos, piece_index_t> m_piece_map;
 
-		// this maps pieces to a range of blocks that are pad files and should not
-		// be picked
+		// this indicates whether a block has been marked as a pad
+		// block or not. It's indexed by block index, i.e. piece_index
+		// * blocks_per_piece + block. These blocks should not be
+		// picked and are considered to be had
 		// TODO: this could be a much more efficient data structure
-		std::set<piece_block> m_pad_blocks;
+		bitfield m_pad_blocks;
+
+		// tracks the number of blocks in a specific piece that are pad blocks
+		std::unordered_map<piece_index_t, int> m_pads_in_piece;
+
+		// the number of bits set in the m_pad_blocks bitfield, i.e.
+		// the number of blocks marked as pads
+		int m_num_pad_blocks = 0;
 
 		// the number of pad blocks that we already have
 		int m_have_pad_blocks = 0;
diff --git a/src/piece_picker.cpp b/src/piece_picker.cpp
index f7368b078..d9ba497ca 100644
--- a/src/piece_picker.cpp
+++ b/src/piece_picker.cpp
@@ -234,7 +234,7 @@ namespace libtorrent {
 		{
 			info.num_peers = 0;
 			info.state = block_info::state_none;
-			if (m_pad_blocks.count(piece_block(piece, block_idx)))
+			if (!m_pad_blocks.empty() && m_pad_blocks.get_bit(static_cast<int>(piece) * m_blocks_per_piece + block_idx))
 			{
 				info.state = block_info::state_finished;
 				++ret.finished;
@@ -3506,7 +3506,23 @@ get_out:
 
 	void piece_picker::mark_as_pad(piece_block block)
 	{
-		m_pad_blocks.insert(block);
+		// if this is the first block we mark as a pad, allocate the bitfield
+		if (m_pad_blocks.empty())
+		{
+			m_pad_blocks.resize(int(m_piece_map.size() * m_blocks_per_piece));
+		}
+
+		int const block_index = static_cast<int>(block.piece_index) * m_blocks_per_piece + block.block_index;
+		TORRENT_ASSERT(block_index < m_pad_blocks.size());
+		TORRENT_ASSERT(block_index >= 0);
+		TORRENT_ASSERT(m_pad_blocks.get_bit(block_index) == false);
+
+		m_pad_blocks.set_bit(block_index);
+		++m_num_pad_blocks;
+		TORRENT_ASSERT(m_pad_blocks.count() == m_num_pad_blocks);
+
+		++m_pads_in_piece[block.piece_index];
+
 		// if we mark and entire piece as a pad file, we need to also
 		// consder that piece as "had" and increment some counters
 		int const blocks = blocks_in_piece(block.piece_index);
@@ -3519,10 +3535,9 @@ get_out:
 
 	int piece_picker::pad_blocks_in_piece(piece_index_t const index) const
 	{
-		int const blocks = blocks_in_piece(index);
-		auto const begin = m_pad_blocks.lower_bound(piece_block(index, 0));
-		auto const end = m_pad_blocks.upper_bound(piece_block(index, blocks));
-		return int(std::distance(begin, end));
+		auto const it = m_pads_in_piece.find(index);
+		if (it == m_pads_in_piece.end()) return 0;
+		return it->second;
 	}
 
 	void piece_picker::get_downloaders(std::vector<torrent_peer*>& d