Limit the set of known extrinsics and blocks. (paritytech#1842)

* Limit the set of known extrinsics and blocks. * Update Cargo.lock * Always insert and pop oldest entry if needed. Always inserting first ensures that the element's LRU position is updated.
freedomfromfiat · Feb 21, 2019 · 8888fc5 · 8888fc5
1 parent bd5f345
commit 8888fc5
Show file tree

Hide file tree

Showing 5 changed files with 97 additions and 9 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/core/network/Cargo.toml b/core/network/Cargo.toml
@@ -16,6 +16,7 @@ error-chain = "0.12"
 bitflags = "1.0"
 futures = "0.1.17"
 linked-hash-map = "0.5"
+linked_hash_set = "0.1.3"
 lru-cache = "0.1.1"
 rustc-hex = "2.0"
 rand = "0.6"

diff --git a/core/network/src/lib.rs b/core/network/src/lib.rs
@@ -27,6 +27,7 @@ mod protocol;
 mod chain;
 mod blocks;
 mod on_demand;
+mod util;
 pub mod config;
 pub mod consensus_gossip;
 pub mod error;

diff --git a/core/network/src/protocol.rs b/core/network/src/protocol.rs
@@ -29,15 +29,13 @@ use crate::sync::{ChainSync, Status as SyncStatus, SyncState};
 use crate::service::{NetworkChan, NetworkMsg, TransactionPool, ExHashT};
 use crate::config::{ProtocolConfig, Roles};
 use rustc_hex::ToHex;
-use std::collections::{BTreeMap, HashMap, HashSet};
+use std::collections::{BTreeMap, HashMap};
 use std::sync::Arc;
-use std::thread;
-use std::time;
-use std::cmp;
+use std::{cmp, num::NonZeroUsize, thread, time};
 use log::{trace, debug, warn};
 use crate::chain::Client;
 use client::light::fetcher::ChangesProof;
-use crate::error;
+use crate::{error, util::LruHashSet};
 
 const REQUEST_TIMEOUT_SEC: u64 = 40;
 const TICK_TIMEOUT: time::Duration = time::Duration::from_millis(1000);
@@ -90,9 +88,9 @@ struct Peer<B: BlockT, H: ExHashT> {
 	/// Requests we are no longer insterested in.
 	obsolete_requests: HashMap<message::RequestId, time::Instant>,
 	/// Holds a set of transactions known to this peer.
-	known_extrinsics: HashSet<H>,
+	known_extrinsics: LruHashSet<H>,
 	/// Holds a set of blocks known to this peer.
-	known_blocks: HashSet<B::Hash>,
+	known_blocks: LruHashSet<B::Hash>,
 	/// Request counter,
 	next_request_id: message::RequestId,
 }
@@ -685,6 +683,8 @@ impl<B: BlockT, S: NetworkSpecialization<B>, H: ExHashT> Protocol<B, S, H> {
 				}
 			}
 
+			let cache_limit = NonZeroUsize::new(1_000_000).expect("1_000_000 > 0; qed");
+
 			let peer = Peer {
 				info: PeerInfo {
 					protocol_version: status.version,
@@ -693,8 +693,8 @@ impl<B: BlockT, S: NetworkSpecialization<B>, H: ExHashT> Protocol<B, S, H> {
 					best_number: status.best_number
 				},
 				block_request: None,
-				known_extrinsics: HashSet::new(),
-				known_blocks: HashSet::new(),
+				known_extrinsics: LruHashSet::new(cache_limit),
+				known_blocks: LruHashSet::new(cache_limit),
 				next_request_id: 0,
 				obsolete_requests: HashMap::new(),
 			};

diff --git a/core/network/src/util.rs b/core/network/src/util.rs
@@ -0,0 +1,76 @@
+// Copyright 2019 Parity Technologies (UK) Ltd.
+// This file is part of Substrate.
+
+// Substrate is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Substrate is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
+
+use linked_hash_set::LinkedHashSet;
+use std::{hash::Hash, num::NonZeroUsize};
+
+/// Wrapper around `LinkedHashSet` which grows bounded.
+///
+/// In the limit, for each element inserted the oldest existing element will be removed.
+#[derive(Debug, Clone)]
+pub(crate) struct LruHashSet<T: Hash + Eq> {
+	set: LinkedHashSet<T>,
+	limit: NonZeroUsize
+}
+
+impl<T: Hash + Eq> LruHashSet<T> {
+	/// Create a new `LruHashSet` with the given (exclusive) limit.
+	pub(crate) fn new(limit: NonZeroUsize) -> Self {
+		Self { set: LinkedHashSet::new(), limit }
+	}
+
+	/// Insert element into the set.
+	///
+	/// Returns `true` if this is a new element to the set, `false` otherwise.
+	/// Maintains the limit of the set by removing the oldest entry if necessary.
+	/// Inserting the same element will update its LRU position.
+	pub(crate) fn insert(&mut self, e: T) -> bool {
+		if self.set.insert(e) {
+			if self.set.len() == usize::from(self.limit) {
+				self.set.pop_front(); // remove oldest entry
+			}
+			return true
+		}
+		false
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	#[test]
+	fn maintains_limit() {
+		let three = NonZeroUsize::new(3).unwrap();
+		let mut set = LruHashSet::<u8>::new(three);
+
+		// First element.
+		assert!(set.insert(1));
+		assert_eq!(vec![&1], set.set.iter().collect::<Vec<_>>());
+
+		// Second element.
+		assert!(set.insert(2));
+		assert_eq!(vec![&1, &2], set.set.iter().collect::<Vec<_>>());
+
+		// Inserting the same element updates its LRU position.
+		assert!(!set.insert(1));
+		assert_eq!(vec![&2, &1], set.set.iter().collect::<Vec<_>>());
+
+		// We reached the limit. The next element forces the oldest one out.
+		assert!(set.insert(3));
+		assert_eq!(vec![&1, &3], set.set.iter().collect::<Vec<_>>());
+	}
+}