1
// Copyright (C) Moondance Labs Ltd.
2
// This file is part of Tanssi.
3

            
4
// Tanssi is free software: you can redistribute it and/or modify
5
// it under the terms of the GNU General Public License as published by
6
// the Free Software Foundation, either version 3 of the License, or
7
// (at your option) any later version.
8

            
9
// Tanssi is distributed in the hope that it will be useful,
10
// but WITHOUT ANY WARRANTY; without even the implied warranty of
11
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12
// GNU General Public License for more details.
13

            
14
// You should have received a copy of the GNU General Public License
15
// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>
16

            
17
use {
18
    dp_collator_assignment::AssignedCollators,
19
    frame_support::traits::Get,
20
    sp_std::{
21
        cmp,
22
        collections::{btree_map::BTreeMap, btree_set::BTreeSet},
23
        marker::PhantomData,
24
        mem,
25
        vec::Vec,
26
    },
27
    tp_traits::{ParaId, RemoveInvulnerables as RemoveInvulnerablesT},
28
};
29

            
30
// Separate import of `sp_std::vec!` macro, which cause issues with rustfmt if grouped
31
// with `sp_std::vec::Vec`.
32
use sp_std::vec;
33

            
34
/// Helper methods to implement collator assignment algorithm
35
pub struct Assignment<T>(PhantomData<T>);
36

            
37
impl<T> Assignment<T>
38
where
39
    T: crate::Config,
40
{
41
    /// Recompute collator assignment from scratch. If the list of collators and the list of
42
    /// container chains are shuffled, this returns a random assignment.
43
349
    pub fn assign_collators_rotate_all<TShuffle>(
44
349
        collators: Vec<T::AccountId>,
45
349
        orchestrator_chain: ChainNumCollators,
46
349
        chains: Vec<ChainNumCollators>,
47
349
        shuffle: Option<TShuffle>,
48
349
    ) -> Result<AssignedCollators<T::AccountId>, AssignmentError>
49
349
    where
50
349
        TShuffle: FnOnce(&mut Vec<T::AccountId>),
51
349
    {
52
349
        // This is just the "always_keep_old" algorithm but with an empty "old"
53
349
        let old_assigned = Default::default();
54
349

            
55
349
        Self::assign_collators_always_keep_old(
56
349
            collators,
57
349
            orchestrator_chain,
58
349
            chains,
59
349
            old_assigned,
60
349
            shuffle,
61
349
        )
62
349
    }
63

            
64
    /// Assign new collators to missing container_chains.
65
    /// Old collators always have preference to remain on the same chain.
66
    /// If there are no missing collators, nothing is changed.
67
    ///
68
    /// `chains` should be shuffled or at least rotated on every session to ensure
69
    /// a fair distribution, because the order of that list affects container chain priority:
70
    /// the first chain on that list will be the first one to get new collators.
71
    ///
72
    /// Similarly, in the `collators` list order means priority, the first collators will be more
73
    /// likely to get assigned. Unlike the list of `chains` which should already be shuffled,
74
    /// collators will be shuffled using the `shuffle` callback when needed. This allows the
75
    /// algorithm to truncate the list of collators and only shuffle the first N. This ensures that
76
    /// shuffling doesn't cause a collator with low priority to be assigned instead of a collator
77
    /// with higher priority.
78
3087
    pub fn assign_collators_always_keep_old<TShuffle>(
79
3087
        collators: Vec<T::AccountId>,
80
3087
        orchestrator_chain: ChainNumCollators,
81
3087
        mut chains: Vec<ChainNumCollators>,
82
3087
        mut old_assigned: AssignedCollators<T::AccountId>,
83
3087
        shuffle: Option<TShuffle>,
84
3087
    ) -> Result<AssignedCollators<T::AccountId>, AssignmentError>
85
3087
    where
86
3087
        TShuffle: FnOnce(&mut Vec<T::AccountId>),
87
3087
    {
88
3087
        if collators.is_empty() && !T::ForceEmptyOrchestrator::get() {
89
4
            return Err(AssignmentError::ZeroCollators);
90
3083
        }
91
3083
        // The rest of this function mostly treats orchestrator chain as another container chain, so move it into
92
3083
        // `old_assigned.container_chains`
93
3083
        let old_orchestrator_assigned = mem::take(&mut old_assigned.orchestrator_chain);
94
3083
        old_assigned
95
3083
            .container_chains
96
3083
            .insert(orchestrator_chain.para_id, old_orchestrator_assigned);
97
3083
        let mut old_assigned = old_assigned.container_chains;
98
3083
        // Orchestrator chain must be the first one in the list because it always has priority
99
3083
        chains.insert(0, orchestrator_chain);
100
7887
        let all_para_ids: Vec<ParaId> = chains.iter().map(|cc| cc.para_id).collect();
101
3083
        let collators_set = BTreeSet::from_iter(collators.iter().cloned());
102
3083
        let chains_with_collators =
103
3083
            Self::select_chains_with_collators(collators.len() as u32, &chains);
104
3083
        let chains_with_collators_set: BTreeSet<ParaId> = chains_with_collators
105
3083
            .iter()
106
5592
            .map(|(para_id, _num_collators)| *para_id)
107
3083
            .collect();
108
3083
        Self::retain_valid_old_assigned(
109
3083
            &mut old_assigned,
110
3083
            &chains_with_collators_set,
111
3083
            &collators_set,
112
3083
        );
113
3083

            
114
3083
        // Ensure the first `min_orchestrator_collators` of orchestrator chain are invulnerables
115
3083
        Self::prioritize_invulnerables(&collators, orchestrator_chain, &mut old_assigned);
116

            
117
3083
        let new_assigned_chains =
118
3083
            Self::assign_full(collators, chains_with_collators, old_assigned, shuffle)?;
119

            
120
3083
        let mut new_assigned = AssignedCollators {
121
3083
            container_chains: new_assigned_chains,
122
3083
            ..Default::default()
123
3083
        };
124

            
125
        // Add container chains with 0 collators so that they are shown in UI
126
10970
        for para_id in all_para_ids {
127
7887
            new_assigned.container_chains.entry(para_id).or_default();
128
7887
        }
129

            
130
        // The rest of this function mostly treats orchestrator chain as another container chain, remove it from
131
        // container chains before returning the final assignment.
132
3083
        let orchestrator_assigned = new_assigned
133
3083
            .container_chains
134
3083
            .remove(&orchestrator_chain.para_id)
135
3083
            .unwrap();
136
3083
        // Sanity check to avoid bricking orchestrator chain
137
3083
        if orchestrator_assigned.is_empty() && !T::ForceEmptyOrchestrator::get() {
138
            return Err(AssignmentError::EmptyOrchestrator);
139
3083
        }
140
3083
        new_assigned.orchestrator_chain = orchestrator_assigned;
141
3083

            
142
3083
        Ok(new_assigned)
143
3087
    }
144

            
145
    /// Select which container chains will be assigned collators and how many collators, but do not specify which
146
    /// collator goes to which chain.
147
    ///
148
    /// Each chain has a min and max number of collators. If the number of collators is not enough to reach the min,
149
    /// no collators are assigned to that chain.
150
    ///
151
    /// If the available number of collators is:
152
    /// * lower than the min of the first chain: we assign all the collators to the first chain. This is the
153
    ///   orchestrator chain and we always want it to have collators.
154
    /// * lower than the sum of all the min: we cannot assign collators to all the chains. So remove chains until
155
    ///   we can. The order is important, the first chains will be assigned collators and the last ones will not.
156
    /// * lower than the sum of all the max: we can assign the min value to all the chains, and have some leftover.
157
    ///   We use the same order to decide where this extra collators will go, by filling the max of the first chain,
158
    ///   then the max of the second chain, and so on.
159
    /// * greater than the sum of all the max: all the chains will be assigned their max number of collators.
160
    ///
161
    /// # Params
162
    ///
163
    /// The first item of `chains` should be the orchestrator chain, because it will be the first one to be assigned
164
    /// collators.
165
    ///
166
    /// # Returns
167
    ///
168
    /// A list of `(para_id, num_collators)`.
169
3094
    pub fn select_chains_with_collators(
170
3094
        num_collators: u32,
171
3094
        chains: &[ChainNumCollators],
172
3094
    ) -> Vec<(ParaId, u32)> {
173
3094
        if chains.is_empty() {
174
            // Avoid panic if chains is empty
175
            return vec![];
176
3094
        }
177
3094
        // Let's count how many container chains we can support with the current number of collators
178
3094
        let mut available_collators = num_collators;
179
3094
        // Handle orchestrator chain in a special way, we always want to assign collators to it, even if we don't
180
3094
        // reach the min.
181
3094
        let min_orchestrator_collators = chains[0].min_collators;
182
3094
        available_collators = available_collators.saturating_sub(min_orchestrator_collators);
183
3094

            
184
3094
        let mut container_chains_with_collators = vec![chains[0]];
185
        // Skipping orchestrator chain because it was handled above
186
4831
        for cc in chains.iter().skip(1) {
187
4653
            if available_collators >= cc.min_collators {
188
2523
                available_collators -= cc.min_collators;
189
2523
                container_chains_with_collators.push(*cc);
190
2553
            } else if available_collators == 0 {
191
                // Do not break if there are still some available collators. Even if they were not enough to reach the
192
                // `min` of this chain, it is possible that one of the chains with less priority has a lower `min`, so
193
                // that chain should be assigned collators.
194
383
                break;
195
1747
            }
196
        }
197

            
198
3094
        let mut required_collators_min = 0;
199
8711
        for cc in &container_chains_with_collators {
200
5617
            required_collators_min += cc.min_collators;
201
5617
        }
202

            
203
3094
        if num_collators < min_orchestrator_collators {
204
            // Edge case: num collators less than min orchestrator collators: fill as much as we can
205
21
            vec![(chains[0].para_id, num_collators)]
206
        } else {
207
            // After assigning the min to all the chains we have this remainder. The remainder will be assigned until
208
            // all the chains reach the max value.
209
3073
            let mut required_collators_remainder = num_collators - required_collators_min;
210
3073
            let mut container_chains_variable = vec![];
211
8669
            for cc in &container_chains_with_collators {
212
5596
                // Each chain will have `min + extra` collators, where extra is capped so `min + extra <= max`.
213
5596
                let extra = cmp::min(
214
5596
                    required_collators_remainder,
215
5596
                    cc.max_collators.saturating_sub(cc.min_collators),
216
5596
                );
217
5596
                let num = cc.min_collators + extra;
218
5596
                required_collators_remainder -= extra;
219
5596
                container_chains_variable.push((cc.para_id, num));
220
5596
            }
221

            
222
3073
            container_chains_variable
223
        }
224
3094
    }
225

            
226
    /// Same as `prioritize_invulnerables` but return the invulnerables instead of inserting them into `old_assigned`.
227
    ///
228
    /// Mutates `old_assigned` by removing invulnerables from their old chain, even if they will later be assigned to
229
    /// the same chain.
230
3090
    pub fn remove_invulnerables(
231
3090
        collators: &[T::AccountId],
232
3090
        orchestrator_chain: ChainNumCollators,
233
3090
        old_assigned: &mut BTreeMap<ParaId, Vec<T::AccountId>>,
234
3090
    ) -> Vec<T::AccountId> {
235
3090
        // TODO: clean this up, maybe change remove_invulnerables trait into something more ergonomic
236
3090
        let min_orchestrator_collators = orchestrator_chain.min_collators as usize;
237
3090
        let invulnerables_already_assigned = T::RemoveInvulnerables::remove_invulnerables(
238
3090
            &mut old_assigned
239
3090
                .get(&orchestrator_chain.para_id)
240
3090
                .cloned()
241
3090
                .unwrap_or_default(),
242
3090
            min_orchestrator_collators,
243
3090
        );
244
3090
        let mut new_invulnerables = invulnerables_already_assigned;
245
3090
        if new_invulnerables.len() >= min_orchestrator_collators {
246
            // We already had invulnerables, we will just move them to the front of the list if they weren't already
247
2290
            return new_invulnerables;
248
800
        }
249
800

            
250
800
        // Not enough invulnerables currently assigned, get rest from new_collators
251
800
        let mut new_collators = collators.to_vec();
252
938
        for (_id, cs) in old_assigned.iter() {
253
4004
            new_collators.retain(|c| !cs.contains(c));
254
938
        }
255
800
        let num_missing_invulnerables = min_orchestrator_collators - new_invulnerables.len();
256
800
        let invulnerables_not_assigned = T::RemoveInvulnerables::remove_invulnerables(
257
800
            &mut new_collators,
258
800
            num_missing_invulnerables,
259
800
        );
260
800
        new_invulnerables.extend(invulnerables_not_assigned);
261
800

            
262
800
        if new_invulnerables.len() >= min_orchestrator_collators {
263
            // Got invulnerables from new_collators, and maybe some were already assigned
264
613
            return new_invulnerables;
265
187
        }
266
187

            
267
187
        // Still not enough invulnerables, try to get an invulnerable that is currently assigned somewhere else
268
187
        let num_missing_invulnerables = min_orchestrator_collators - new_invulnerables.len();
269
187
        let mut collators = collators.to_vec();
270
187
        let new_invulnerables_set = BTreeSet::from_iter(new_invulnerables.iter().cloned());
271
1240
        collators.retain(|c| {
272
1240
            // Remove collators already selected
273
1240
            !new_invulnerables_set.contains(c)
274
1240
        });
275
187
        let invulnerables_assigned_elsewhere =
276
187
            T::RemoveInvulnerables::remove_invulnerables(&mut collators, num_missing_invulnerables);
277
187

            
278
187
        if invulnerables_assigned_elsewhere.is_empty() {
279
            // If at this point we still do not have enough invulnerables, it means that there are no
280
            // enough invulnerables, so no problem, but return the invulnerables
281
184
            return new_invulnerables;
282
3
        }
283
3

            
284
3
        new_invulnerables.extend(invulnerables_assigned_elsewhere.iter().cloned());
285
3

            
286
3
        // In this case we must delete the old assignment of the invulnerables
287
3
        let reassigned_invulnerables_set = BTreeSet::from_iter(invulnerables_assigned_elsewhere);
288
        // old_assigned.remove_collators_in_set
289
8
        for (_id, cs) in old_assigned.iter_mut() {
290
17
            cs.retain(|c| !reassigned_invulnerables_set.contains(c));
291
8
        }
292

            
293
3
        new_invulnerables
294
3090
    }
295

            
296
    /// Ensure orchestrator chain has `min_orchestrator` invulnerables. If that's not possible, it tries to add as
297
    /// many invulnerables as possible.
298
    ///
299
    /// Get invulnerables from:
300
    /// * old_assigned in orchestrator
301
    /// * new collators
302
    /// * old_assigned elsewhere
303
    ///
304
    /// In that order.
305
    ///
306
    /// Mutates `old_assigned` because invulnerables will be inserted there, and if invulnerables were already
307
    /// assigned to some other chain, they will be removed from that other chain as well.
308
    ///
309
    /// # Params
310
    ///
311
    /// * `old_assigned` must be a subset of `collators`
312
    /// * `old_assigned` must not have duplicate collators.
313
    ///
314
    /// # Returns
315
    ///
316
    /// The number of invulnerables assigned to the orchestrator chain, capped to `min_collators`.
317
3090
    pub fn prioritize_invulnerables(
318
3090
        collators: &[T::AccountId],
319
3090
        orchestrator_chain: ChainNumCollators,
320
3090
        old_assigned: &mut BTreeMap<ParaId, Vec<T::AccountId>>,
321
3090
    ) -> usize {
322
3090
        let new_invulnerables =
323
3090
            Self::remove_invulnerables(collators, orchestrator_chain, old_assigned);
324
3090

            
325
3090
        if !new_invulnerables.is_empty() {
326
2552
            Self::insert_invulnerables(
327
2552
                old_assigned.entry(orchestrator_chain.para_id).or_default(),
328
2552
                &new_invulnerables,
329
2552
            );
330
2999
        }
331

            
332
3090
        new_invulnerables.len()
333
3090
    }
334

            
335
    /// Assign collators assuming that the number of collators is greater than or equal to the required.
336
    /// The order of both container chains and collators is important to ensure randomness when `old_assigned` is
337
    /// empty.
338
    ///
339
    /// # Params
340
    ///
341
    /// * `old_assigned` does not need to be a subset of `collators`: collators are checked and removed.
342
    /// * `old_assigned` does not need to be a subset of `chains`, unused para ids are removed. Collators
343
    ///   assigned to a para_id not present in `chains` may be reassigned to another para_id.
344
    /// * `chains` `num_collators` can be 0. In that case an empty vec is returned for that para id.
345
    /// * `old_assigned` must not have duplicate collators.
346
    /// * `shuffle` is used to shuffle the list collators. The list will be truncated to only have
347
    ///   the number of required collators, to ensure that shuffling doesn't cause a collator with low
348
    ///   priority to be assigned instead of a collator with higher priority.
349
    ///
350
    /// # Returns
351
    ///
352
    /// The collator assigment, a map from `ParaId` to `Vec<T>`.
353
    ///
354
    /// Or an error if the number of collators is not enough to fill all the chains, or if the required number
355
    /// of collators overflows a `u32`.
356
3092
    pub fn assign_full<TShuffle>(
357
3092
        collators: Vec<T::AccountId>,
358
3092
        chains: Vec<(ParaId, u32)>,
359
3092
        mut old_assigned: BTreeMap<ParaId, Vec<T::AccountId>>,
360
3092
        shuffle: Option<TShuffle>,
361
3092
    ) -> Result<BTreeMap<ParaId, Vec<T::AccountId>>, AssignmentError>
362
3092
    where
363
3092
        TShuffle: FnOnce(&mut Vec<T::AccountId>),
364
3092
    {
365
3092
        let mut required_collators = 0usize;
366
5607
        for (_para_id, num_collators) in chains.iter() {
367
5607
            let num_collators =
368
5607
                usize::try_from(*num_collators).map_err(|_| AssignmentError::NotEnoughCollators)?;
369
5607
            required_collators = required_collators
370
5607
                .checked_add(num_collators)
371
5607
                .ok_or(AssignmentError::NotEnoughCollators)?;
372
        }
373

            
374
        // This check is necessary to ensure priority: if the number of collators is less than required, it is
375
        // possible that the chain with the least priority could be assigned collators (since they are in
376
        // old_assigned), while some chains with higher priority might have no collators.
377
3092
        if collators.len() < required_collators {
378
1
            return Err(AssignmentError::NotEnoughCollators);
379
3091
        }
380
3091
        // We checked that the sum of all `num_collators` fits in `usize`, so we can safely use `as usize`.
381
3091

            
382
3091
        // Remove invalid collators and para ids from `old_assigned`
383
3091
        let para_ids_set =
384
5605
            BTreeSet::from_iter(chains.iter().map(|(para_id, _num_collators)| *para_id));
385
3091
        let collators_set = BTreeSet::from_iter(collators.iter().cloned());
386
3091
        Self::retain_valid_old_assigned(&mut old_assigned, &para_ids_set, &collators_set);
387

            
388
        // Truncate num collators to required
389
5605
        for (para_id, num_collators) in chains.iter() {
390
5605
            let entry = old_assigned.entry(*para_id).or_default();
391
5605
            entry.truncate(*num_collators as usize);
392
5605
        }
393

            
394
        // Count number of needed new collators. This is equivalent to:
395
        // `required_collators - old_assigned.iter().map(|cs| cs.len()).sum()`.
396
3091
        let mut needed_new_collators = 0;
397
5605
        for (para_id, num_collators) in chains.iter() {
398
5605
            let cs = old_assigned.entry(*para_id).or_default();
399
5605
            needed_new_collators += (*num_collators as usize).saturating_sub(cs.len());
400
5605
        }
401

            
402
3091
        let assigned_collators: BTreeSet<T::AccountId> = old_assigned
403
3091
            .iter()
404
5605
            .flat_map(|(_para_id, para_collators)| para_collators.iter().cloned())
405
3091
            .collect();
406
3091

            
407
3091
        // Truncate list of new_collators to `needed_new_collators` and shuffle it.
408
3091
        // This has the effect of keeping collator priority (the first collator of that list is more
409
3091
        // likely to be assigned to a chain than the last collator of that list), while also
410
3091
        // ensuring randomness (the original order does not directly affect which chain the
411
3091
        // collators are assigned to).
412
3091
        let mut new_collators: Vec<_> = collators
413
3091
            .into_iter()
414
3474
            .filter(|x| {
415
2280
                // Keep collators not already assigned
416
2280
                !assigned_collators.contains(x)
417
3474
            })
418
3091
            .take(needed_new_collators)
419
3091
            .collect();
420
3091
        if let Some(shuffle) = shuffle {
421
51
            shuffle(&mut new_collators);
422
3040
        }
423
3091
        let mut new_collators = new_collators.into_iter();
424

            
425
        // Fill missing collators
426
5605
        for (para_id, num_collators) in chains.iter() {
427
5605
            let cs = old_assigned.entry(*para_id).or_default();
428

            
429
7141
            while cs.len() < *num_collators as usize {
430
                // This error should never happen because we calculated `needed_new_collators`
431
                // using the same algorithm
432
1536
                let nc = new_collators
433
1536
                    .next()
434
1536
                    .ok_or(AssignmentError::NotEnoughCollators)?;
435
1536
                cs.push(nc);
436
            }
437
        }
438

            
439
3091
        Ok(old_assigned)
440
3092
    }
441

            
442
    /// Insert invulnerables ensuring that they are always the first in the list.
443
    /// The order of both lists is preserved.
444
    /// `assigned` may already contain the invulnerables, in that case they are only moved to the front.
445
    ///
446
    /// Invulnerables need to be the first of the list because we may truncate the list of collators if the number of
447
    /// collators changes, and in that case we want invulnerables to stay assigned there.
448
2552
    pub fn insert_invulnerables(assigned: &mut Vec<T::AccountId>, invulnerables: &[T::AccountId]) {
449
2572
        assigned.retain(|item| !invulnerables.contains(item));
450
2552

            
451
2552
        let mut new_assigned = invulnerables.to_vec();
452
2552
        new_assigned.extend(mem::take(assigned));
453
2552

            
454
2552
        *assigned = new_assigned;
455
2552
    }
456

            
457
    /// Removes invalid entries from `old_assigned`:
458
    ///
459
    /// * para ids not in `chains_with_collators`
460
    /// * collators not in `collators`
461
6175
    pub fn retain_valid_old_assigned(
462
6175
        old_assigned: &mut BTreeMap<ParaId, Vec<T::AccountId>>,
463
6175
        chains_with_collators: &BTreeSet<ParaId>,
464
6175
        collators: &BTreeSet<T::AccountId>,
465
6175
    ) {
466
6175
        // old_assigned.remove_container_chains_not_in_set
467
11933
        old_assigned.retain(|id, _cs| chains_with_collators.contains(id));
468
        // old_assigned.remove_collators_not_in_set
469
9987
        for (_id, cs) in old_assigned.iter_mut() {
470
13743
            cs.retain(|c| collators.contains(c));
471
9987
        }
472
6175
    }
473
}
474

            
475
/// Errors than can happen during collator assignment
476
#[derive(Debug, Clone, PartialEq, Eq)]
477
pub enum AssignmentError {
478
    /// An empty list of collators was passed to `assign_collators_always_keep_old`
479
    ZeroCollators,
480
    /// The required number of collators for `assign_full` is greater than the provided number of collators.
481
    /// Also includes possible overflows in number of collators.
482
    NotEnoughCollators,
483
    /// No collators were assigned to orchestrator chain
484
    EmptyOrchestrator,
485
}
486

            
487
/// A `ParaId` and a range of collators that need to be assigned to it.
488
/// This can be a container chain, a parathread, or the orchestrator chain.
489
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
490
pub struct ChainNumCollators {
491
    pub para_id: ParaId,
492
    pub min_collators: u32,
493
    // This will only be filled if all the other min have been reached
494
    pub max_collators: u32,
495
}