From f07a1e6bafc3d1ab7e2e5aef6756bfafcd913d68 Mon Sep 17 00:00:00 2001
From: Arkadiy Paronyan <arkady.paronyan@gmail.com>
Date: Wed, 24 Aug 2016 17:00:14 +0200
Subject: [PATCH] Fixed neighbours collection (#1996)

---
 util/network/src/discovery.rs | 25 ++++++++++++++++++++-----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/util/network/src/discovery.rs b/util/network/src/discovery.rs
index 2b98f3230..a78f05804 100644
--- a/util/network/src/discovery.rs
+++ b/util/network/src/discovery.rs
@@ -59,7 +59,7 @@ pub struct BucketEntry {
 	pub timeout: Option<u64>,
 }
 
-struct NodeBucket {
+pub struct NodeBucket {
 	nodes: VecDeque<BucketEntry>, //sorted by last active
 }
 
@@ -281,12 +281,12 @@ impl Discovery {
 				if count == BUCKET_SIZE {
 					// delete the most distant element
 					let remove = {
-						let (_, last) = found.iter_mut().next_back().unwrap();
+						let (key, last) = found.iter_mut().next_back().unwrap();
 						last.pop();
-						last.is_empty()
+						if last.is_empty() { Some(key.clone()) } else { None }
 					};
-					if remove {
-						found.remove(&distance);
+					if let Some(remove) = remove {
+						found.remove(&remove);
 					}
 				}
 				else {
@@ -605,6 +605,21 @@ mod tests {
 		assert!(removed > 0);
 	}
 
+	#[test]
+	fn find_nearest_saturated() {
+		use super::*;
+		let mut buckets: Vec<_> = (0..256).map(|_| NodeBucket::new()).collect();
+		let ep = NodeEndpoint { address: SocketAddr::from_str("127.0.0.1:40447").unwrap(), udp_port: 40447 };
+		for _ in 0..(16 + 10) {
+			buckets[0].nodes.push_back(BucketEntry {
+				address: NodeEntry { id: NodeId::new(), endpoint: ep.clone() },
+				timeout: None
+			});
+		}
+		let nearest = Discovery::nearest_node_entries(&NodeId::new(), &buckets);
+		assert_eq!(nearest.len(), 16)
+	}
+
 	#[test]
 	fn packets() {
 		let key = KeyPair::create().unwrap();