local_collations.rs 6.56 KB
Newer Older
Shawn Tabrizi's avatar
Shawn Tabrizi committed
1
// Copyright 2018-2020 Parity Technologies (UK) Ltd.
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
// This file is part of Polkadot.

// Polkadot 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.

// Polkadot 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 Polkadot.  If not, see <http://www.gnu.org/licenses/>.

//! Local collations to be circulated to validators.
//!
//! Collations are attempted to be repropagated when a new validator connects,
//! a validator changes his session key, or when they are generated.

asynchronous rob's avatar
asynchronous rob committed
22
use polkadot_primitives::v0::{Hash, ValidatorId};
23
use crate::legacy::collator_pool::Role;
24
use std::collections::{HashMap, HashSet};
Gavin Wood's avatar
Gavin Wood committed
25
26
use std::time::Duration;
use wasm_timer::Instant;
27
use rand::seq::SliceRandom;
28
29
30
31

const LIVE_FOR: Duration = Duration::from_secs(60 * 5);

struct LocalCollation<C> {
32
	targets: HashSet<ValidatorId>,
33
34
35
36
37
38
	collation: C,
	live_since: Instant,
}

/// Tracker for locally collated values and which validators to send them to.
pub struct LocalCollations<C> {
39
	primary_for: HashSet<ValidatorId>,
40
41
42
	local_collations: HashMap<Hash, LocalCollation<C>>,
}

43
44
45
46
47
48
impl<C: Clone> Default for LocalCollations<C> {
	fn default() -> Self {
		Self::new()
	}
}

49
50
51
52
53
54
55
56
57
58
59
impl<C: Clone> LocalCollations<C> {
	/// Create a new `LocalCollations` tracker.
	pub fn new() -> Self {
		LocalCollations {
			primary_for: HashSet::new(),
			local_collations: HashMap::new(),
		}
	}

	/// Validator gave us a new role. If the new role is "primary", this function might return
	/// a set of collations to send to that validator.
60
	pub fn note_validator_role(&mut self, key: ValidatorId, role: Role) -> Vec<(Hash, C)> {
61
62
63
64
65
66
		match role {
			Role::Backup => {
				self.primary_for.remove(&key);
				Vec::new()
			}
			Role::Primary => {
Gav Wood's avatar
Gav Wood committed
67
				let new_primary = self.primary_for.insert(key.clone());
68
69
70
71
72
73
74
75
76
77
78
				if new_primary {
					self.collations_targeting(&key)
				} else {
					Vec::new()
				}
			}
		}
	}

	/// Fresh session key from a validator. Returns a vector of collations to send
	/// to the validator.
79
	pub fn fresh_key(&mut self, old_key: &ValidatorId, new_key: &ValidatorId) -> Vec<(Hash, C)> {
80
		if self.primary_for.remove(old_key) {
Gav Wood's avatar
Gav Wood committed
81
			self.primary_for.insert(new_key.clone());
82
83
84
85
86
87
88
89

			self.collations_targeting(new_key)
		} else {
			Vec::new()
		}
	}

	/// Validator disconnected.
90
	pub fn on_disconnect(&mut self, key: &ValidatorId) {
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
		self.primary_for.remove(key);
	}

	/// Mark collations relevant to the given parent hash as obsolete.
	pub fn collect_garbage(&mut self, relay_parent: Option<&Hash>) {
		if let Some(relay_parent) = relay_parent {
			self.local_collations.remove(relay_parent);
		}

		let now = Instant::now();
		self.local_collations.retain(|_, v| v.live_since + LIVE_FOR > now);
	}

	/// Add a collation. Returns an iterator of session keys to send to and lazy copies of the collation.
	pub fn add_collation<'a>(
		&'a mut self,
		relay_parent: Hash,
108
		targets: HashSet<ValidatorId>,
109
		collation: C
110
	) -> impl Iterator<Item=(ValidatorId, C)> + 'a {
111
112
113
114
115
116
117
118
119
120
		self.local_collations.insert(relay_parent, LocalCollation {
			targets,
			collation,
			live_since: Instant::now(),
		});

		let local = self.local_collations.get(&relay_parent)
			.expect("just inserted to this key; qed");

		let borrowed_collation = &local.collation;
121
122
123
124
125
126
127
128

		// If we are conected to multiple validators,
		// make sure we always send the collation to one of the validators
		// we are registered as backup. This ensures that one collator that
		// is primary at multiple validators, desn't block the Parachain from progressing.
		let mut rng = rand::thread_rng();
		let diff = local.targets.difference(&self.primary_for).collect::<Vec<_>>();

129
130
		local.targets
			.intersection(&self.primary_for)
131
			.chain(diff.choose(&mut rng).map(|r| r.clone()))
Gav Wood's avatar
Gav Wood committed
132
			.map(move |k| (k.clone(), borrowed_collation.clone()))
133
134
	}

135
	fn collations_targeting(&self, key: &ValidatorId) -> Vec<(Hash, C)> {
136
137
138
139
140
141
142
143
144
145
		self.local_collations.iter()
			.filter(|&(_, ref v)| v.targets.contains(key))
			.map(|(h, v)| (*h, v.collation.clone()))
			.collect()
	}
}

#[cfg(test)]
mod tests {
	use super::*;
146
	use sp_core::crypto::UncheckedInto;
asynchronous rob's avatar
asynchronous rob committed
147
	use polkadot_primitives::v0::ValidatorId;
148
149
150

	#[test]
	fn add_validator_with_ready_collation() {
Gav Wood's avatar
Gav Wood committed
151
		let key: ValidatorId = [1; 32].unchecked_into();
152
153
154
		let relay_parent = [2; 32].into();
		let targets = {
			let mut set = HashSet::new();
Gav Wood's avatar
Gav Wood committed
155
			set.insert(key.clone());
156
157
158
159
			set
		};

		let mut tracker = LocalCollations::new();
160
		assert!(tracker.add_collation(relay_parent, targets, 5).next().is_some());
161
162
163
164
165
		assert_eq!(tracker.note_validator_role(key, Role::Primary), vec![(relay_parent, 5)]);
	}

	#[test]
	fn rename_with_ready() {
Gav Wood's avatar
Gav Wood committed
166
167
		let orig_key: ValidatorId = [1; 32].unchecked_into();
		let new_key: ValidatorId  = [2; 32].unchecked_into();
168
169
170
		let relay_parent = [255; 32].into();
		let targets = {
			let mut set = HashSet::new();
Gav Wood's avatar
Gav Wood committed
171
			set.insert(new_key.clone());
172
173
174
175
			set
		};

		let mut tracker: LocalCollations<u8> = LocalCollations::new();
176
		assert!(tracker.add_collation(relay_parent, targets, 5).next().is_some());
Gav Wood's avatar
Gav Wood committed
177
		assert!(tracker.note_validator_role(orig_key.clone(), Role::Primary).is_empty());
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
		assert_eq!(tracker.fresh_key(&orig_key, &new_key), vec![(relay_parent, 5u8)]);
	}

	#[test]
	fn collecting_garbage() {
		let relay_parent_a = [255; 32].into();
		let relay_parent_b = [222; 32].into();

		let mut tracker: LocalCollations<u8> = LocalCollations::new();
		assert!(tracker.add_collation(relay_parent_a, HashSet::new(), 5).next().is_none());
		assert!(tracker.add_collation(relay_parent_b, HashSet::new(), 69).next().is_none());

		let live_since = Instant::now() - LIVE_FOR - Duration::from_secs(10);
		tracker.local_collations.get_mut(&relay_parent_b).unwrap().live_since = live_since;

		tracker.collect_garbage(Some(&relay_parent_a));

		// first one pruned because of relay parent, other because of time.
		assert!(tracker.local_collations.is_empty());
	}

	#[test]
	fn add_collation_with_connected_target() {
Gav Wood's avatar
Gav Wood committed
201
		let key: ValidatorId = [1; 32].unchecked_into();
202
203
204
		let relay_parent = [2; 32].into();
		let targets = {
			let mut set = HashSet::new();
Gav Wood's avatar
Gav Wood committed
205
			set.insert(key.clone());
206
207
208
209
			set
		};

		let mut tracker = LocalCollations::new();
Gav Wood's avatar
Gav Wood committed
210
		assert!(tracker.note_validator_role(key.clone(), Role::Primary).is_empty());
211
212
213
214
		assert_eq!(tracker.add_collation(relay_parent, targets, 5).next(), Some((key, 5)));

	}
}