2 files changed, 32 insertions, 13 deletions

diff --git a/src/cryptonote_protocol/block_queue.cpp b/src/cryptonote_protocol/block_queue.cpp
index c39d67ceb..05f4189fb 100644
--- a/src/cryptonote_protocol/block_queue.cpp
+++ b/src/cryptonote_protocol/block_queue.cpp
@@ -57,7 +57,11 @@ void block_queue::add_blocks(uint64_t height, std::vector<cryptonote::block_comp
   bool has_hashes = remove_span(height, &hashes);
   blocks.insert(span(height, std::move(bcel), connection_id, rate, size));
   if (has_hashes)
+  {
+    for (const crypto::hash &h: hashes)
+      requested_hashes.insert(h);
     set_span_hashes(height, connection_id, hashes);
+  }
 }
 
 void block_queue::add_blocks(uint64_t height, uint64_t nblocks, const boost::uuids::uuid &connection_id, boost::posix_time::ptime time)
@@ -76,11 +80,19 @@ void block_queue::flush_spans(const boost::uuids::uuid &connection_id, bool all)
     block_map::iterator j = i++;
     if (j->connection_id == connection_id && (all || j->blocks.size() == 0))
     {
-      blocks.erase(j);
+      erase_block(j);
     }
   }
 }
 
+void block_queue::erase_block(block_map::iterator j)
+{
+  CHECK_AND_ASSERT_THROW_MES(j != blocks.end(), "Invalid iterator");
+  for (const crypto::hash &h: j->hashes)
+    requested_hashes.erase(h);
+  blocks.erase(j);
+}
+
 void block_queue::flush_stale_spans(const std::set<boost::uuids::uuid> &live_connections)
 {
   boost::unique_lock<boost::recursive_mutex> lock(mutex);
@@ -92,7 +104,7 @@ void block_queue::flush_stale_spans(const std::set<boost::uuids::uuid> &live_con
     block_map::iterator j = i++;
     if (live_connections.find(j->connection_id) == live_connections.end() && j->blocks.size() == 0)
     {
-      blocks.erase(j);
+      erase_block(j);
     }
   }
 }
@@ -106,7 +118,7 @@ bool block_queue::remove_span(uint64_t start_block_height, std::vector<crypto::h
     {
       if (hashes)
         *hashes = std::move(i->hashes);
-      blocks.erase(i);
+      erase_block(i);
       return true;
     }
   }
@@ -121,7 +133,7 @@ void block_queue::remove_spans(const boost::uuids::uuid &connection_id, uint64_t
     block_map::iterator j = i++;
     if (j->connection_id == connection_id && j->start_block_height <= start_block_height)
     {
-      blocks.erase(j);
+      erase_block(j);
     }
   }
 }
@@ -160,16 +172,15 @@ std::string block_queue::get_overview() const
   return s;
 }
 
+inline bool block_queue::requested_internal(const crypto::hash &hash) const
+{
+  return requested_hashes.find(hash) != requested_hashes.end();
+}
+
 bool block_queue::requested(const crypto::hash &hash) const
 {
   boost::unique_lock<boost::recursive_mutex> lock(mutex);
-  for (const auto &span: blocks)
-  {
-    for (const auto &h: span.hashes)
-      if (h == hash)
-        return true;
-  }
-  return false;
+  return requested_internal(hash);
 }
 
 std::pair<uint64_t, uint64_t> block_queue::reserve_span(uint64_t first_block_height, uint64_t last_block_height, uint64_t max_blocks, const boost::uuids::uuid &connection_id, const std::vector<crypto::hash> &block_hashes, boost::posix_time::ptime time)
@@ -184,7 +195,7 @@ std::pair<uint64_t, uint64_t> block_queue::reserve_span(uint64_t first_block_hei
 
   uint64_t span_start_height = last_block_height - block_hashes.size() + 1;
   std::vector<crypto::hash>::const_iterator i = block_hashes.begin();
-  while (i != block_hashes.end() && requested(*i))
+  while (i != block_hashes.end() && requested_internal(*i))
   {
     ++i;
     ++span_start_height;
@@ -256,8 +267,10 @@ void block_queue::set_span_hashes(uint64_t start_height, const boost::uuids::uui
     if (i->start_block_height == start_height && i->connection_id == connection_id)
     {
       span s = *i;
-      blocks.erase(i);
+      erase_block(i);
       s.hashes = std::move(hashes);
+      for (const crypto::hash &h: s.hashes)
+        requested_hashes.insert(h);
       blocks.insert(s);
       return;
     }
diff --git a/src/cryptonote_protocol/block_queue.h b/src/cryptonote_protocol/block_queue.h
index 9059e89ac..9cce95075 100644
--- a/src/cryptonote_protocol/block_queue.h
+++ b/src/cryptonote_protocol/block_queue.h
@@ -33,6 +33,7 @@
 #include <string>
 #include <vector>
 #include <set>
+#include <unordered_set>
 #include <boost/thread/recursive_mutex.hpp>
 #include <boost/uuid/uuid.hpp>
 
@@ -93,7 +94,12 @@ namespace cryptonote
     bool requested(const crypto::hash &hash) const;
 
   private:
+    void erase_block(block_map::iterator j);
+    inline bool requested_internal(const crypto::hash &hash) const;
+
+  private:
     block_map blocks;
     mutable boost::recursive_mutex mutex;
+    std::unordered_set<crypto::hash> requested_hashes;
   };
 }