optimize DHT

src/kademlia/routing_table.cpp

@@ -646,24 +646,35 @@ bool routing_table::add_node(node_entry e)
 			// nodes sharing a prefix. Find all nodes that do not
 			// have a unique prefix
 
-			std::sort(b.begin(), b.end(), boost::bind(&node_entry::id, _1) < boost::bind(&node_entry::id, _2));
+			// find node entries with duplicate prefixes in O(1)
+			std::vector<bucket_t::iterator> prefix(1 << (8 - mask_shift), b.end());
+			TORRENT_ASSERT(prefix.size() >= bucket_size_limit);
+
+			// the begin iterator from this object is used as a placeholder
+			// for an occupied slot whose node has already been added to the
+			// duplicate nodes list.
+			bucket_t placeholder;
 
-			int last_prefix = -1;
 			nodes.reserve(b.size());
 			for (j = b.begin(); j != b.end(); ++j)
 			{
 				node_id id = j->id;
 				id <<= bucket_index + 1;
-				int this_prefix = id[0] & mask;
+				int this_prefix = (id[0] & mask) >> mask_shift;
-				if (this_prefix != last_prefix)
+				TORRENT_ASSERT(this_prefix >= 0);
+				TORRENT_ASSERT(this_prefix < prefix.size());
+				if (prefix[this_prefix] != b.end())
 				{
-					last_prefix = this_prefix;
+					// there's already a node with this prefix. Remember both
-					continue;
+					// duplicates.
-				}
-
-				if (nodes.empty() || nodes.back() != j-1)
-					nodes.push_back(j-1);
 					nodes.push_back(j);
+
+					if (prefix[this_prefix] != placeholder.begin())
+					{
+						nodes.push_back(prefix[this_prefix]);
+						prefix[this_prefix] = placeholder.begin();
+					}
+				}
 			}
 
 			if (!nodes.empty())
@@ -959,6 +970,7 @@ void routing_table::find_node(node_id const& target
 
 	table_t::iterator j = i;
 
+	int unsorted_start_idx = 0;
 	for (; j != m_buckets.end() && int(l.size()) < count; ++j)
 	{
 		bucket_t& b = j->live_nodes;
@@ -974,32 +986,30 @@ void routing_table::find_node(node_id const& target
 				, !boost::bind(&node_entry::confirmed, _1));
 		}
 
-		if (int(l.size()) >= count)
+		if (int(l.size()) == count) return;
+
+		if (int(l.size()) > count)
 		{
 			// sort the nodes by how close they are to the target
-			std::sort(l.begin(), l.end(), boost::bind(&compare_ref
+			std::sort(l.begin() + unsorted_start_idx, l.end(), boost::bind(&compare_ref
 				, boost::bind(&node_entry::id, _1)
 				, boost::bind(&node_entry::id, _2), target));
 
 			l.resize(count);
 			return;
 		}
+		unsorted_start_idx = int(l.size());
 	}
 
 	// if we still don't have enough nodes, copy nodes
 	// further away from us
 
 	if (i == m_buckets.begin())
-	{
-		// sort the nodes by how close they are to the target
-		std::sort(l.begin(), l.end(), boost::bind(&compare_ref
-			, boost::bind(&node_entry::id, _1)
-			, boost::bind(&node_entry::id, _2), target));
 		return;
-	}
 
 	j = i;
 
+	unsorted_start_idx = int(l.size());
 	do
 	{
 		--j;
@@ -1015,26 +1025,23 @@ void routing_table::find_node(node_id const& target
 				, !boost::bind(&node_entry::confirmed, _1));
 		}
 
-		if (int(l.size()) >= count)
+		if (int(l.size()) == count) return;
+
+		if (int(l.size()) > count)
 		{
 			// sort the nodes by how close they are to the target
-			std::sort(l.begin(), l.end(), boost::bind(&compare_ref
+			std::sort(l.begin() + unsorted_start_idx, l.end(), boost::bind(&compare_ref
 				, boost::bind(&node_entry::id, _1)
 				, boost::bind(&node_entry::id, _2), target));
 
 			l.resize(count);
 			return;
 		}
+		unsorted_start_idx = int(l.size());
 	}
 	while (j != m_buckets.begin() && int(l.size()) < count);
 
-	// sort the nodes by how close they are to the target
+	TORRENT_ASSERT(int(l.size()) <= count);
-	std::sort(l.begin(), l.end(), boost::bind(&compare_ref
-		, boost::bind(&node_entry::id, _1)
-		, boost::bind(&node_entry::id, _2), target));
-
-	if (int(l.size()) >= count)
-		l.resize(count);
 }
 
 } } // namespace libtorrent::dht

test/test_dht.cpp

@@ -988,7 +988,8 @@ int test_main()
 		dht_settings s;
 		//	s.restrict_routing_ips = false;
 		node_id id = to_hash("3123456789abcdef01232456789abcdef0123456");
-		dht::routing_table table(id, 10, s);
+		const int bucket_size = 10;
+		dht::routing_table table(id, bucket_size, s);
 		std::vector<node_entry> nodes;
 		TEST_EQUAL(table.size().get<0>(), 0);
 
@@ -1107,9 +1108,9 @@ int test_main()
 		TEST_EQUAL(temp.size(), nodes.size());
 
 		std::generate(tmp.begin(), tmp.end(), &std::rand);
-		table.find_node(tmp, temp, 0, 7);
+		table.find_node(tmp, temp, 0, bucket_size);
 		std::cout << "returned: " << temp.size() << std::endl;
-		TEST_EQUAL(temp.size(), 7);
+		TEST_EQUAL(temp.size(), bucket_size);
 
 		std::sort(nodes.begin(), nodes.end(), boost::bind(&compare_ref
 				, boost::bind(&node_entry::id, _1)
@@ -1123,16 +1124,16 @@ int test_main()
 		{
 			int hit = std::find_if(nodes.begin(), nodes.end()
 				, boost::bind(&node_entry::id, _1) == i->id) - nodes.begin();
-			//		std::cerr << hit << std::endl;
+			std::cerr << hit << std::endl;
 			if (hit < int(temp.size())) ++hits;
 		}
 		std::cout << "hits: " << hits << std::endl;
-		TEST_CHECK(hits == int(temp.size()));
+		TEST_EQUAL(hits, int(temp.size()));
 
 		std::generate(tmp.begin(), tmp.end(), &std::rand);
-		table.find_node(tmp, temp, 0, 15);
+		table.find_node(tmp, temp, 0, bucket_size * 2);
 		std::cout << "returned: " << temp.size() << std::endl;
-		TEST_EQUAL(int(temp.size()), (std::min)(15, int(nodes.size())));
+		TEST_EQUAL(int(temp.size()), (std::min)(bucket_size * 2, int(nodes.size())));
 
 		std::sort(nodes.begin(), nodes.end(), boost::bind(&compare_ref
 				, boost::bind(&node_entry::id, _1)
@@ -1146,11 +1147,11 @@ int test_main()
 		{
 			int hit = std::find_if(nodes.begin(), nodes.end()
 				, boost::bind(&node_entry::id, _1) == i->id) - nodes.begin();
-			//		std::cerr << hit << std::endl;
+			std::cerr << hit << std::endl;
 			if (hit < int(temp.size())) ++hits;
 		}
 		std::cout << "hits: " << hits << std::endl;
-		TEST_CHECK(hits == int(temp.size()));
+		TEST_EQUAL(hits, int(temp.size()));
 
 		using namespace libtorrent::dht;