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
//! # Collator Assignment Pallet
18
//!
19
//! This pallet assigns a list of collators to:
20
//!    - the orchestrator chain
21
//!    - a set of container chains
22
//!
23
//! The set of container chains is retrieved thanks to the GetContainerChains trait
24
//! The number of collators to assign to the orchestrator chain and the number
25
//! of collators to assign to each container chain is retrieved through the GetHostConfiguration
26
//! trait.
27
//!  
28
//! The pallet uses the following approach:
29
//!
30
//! - First, it aims at filling the necessary collators to serve the orchestrator chain
31
//! - Second, it aims at filling in-order (FIFO) the existing containerChains
32
//!
33
//! Upon new session, this pallet takes whatever assignation was in the PendingCollatorContainerChain
34
//! storage, and assigns it as the current CollatorContainerChain. In addition, it takes the next
35
//! queued set of parachains and collators and calculates the assignment for the next session, storing
36
//! it in the PendingCollatorContainerChain storage item.
37
//!
38
//! The reason for the collator-assignment pallet to work with a one-session delay assignment is because
39
//! we want collators to know at least one session in advance the container chain/orchestrator that they
40
//! are assigned to.
41

            
42
#![cfg_attr(not(feature = "std"), no_std)]
43

            
44
use {
45
    crate::assignment::{Assignment, ChainNumCollators},
46
    core::ops::Mul,
47
    frame_support::{pallet_prelude::*, traits::Currency},
48
    frame_system::pallet_prelude::BlockNumberFor,
49
    rand::{seq::SliceRandom, SeedableRng},
50
    rand_chacha::ChaCha20Rng,
51
    sp_runtime::{
52
        traits::{AtLeast32BitUnsigned, One, Zero},
53
        Perbill, Saturating,
54
    },
55
    sp_std::{collections::btree_set::BTreeSet, fmt::Debug, prelude::*, vec},
56
    tp_traits::{
57
        CollatorAssignmentTip, ForSession, FullRotationModes, GetContainerChainAuthor,
58
        GetContainerChainsWithCollators, GetHostConfiguration, GetSessionContainerChains, ParaId,
59
        ParaIdAssignmentHooks, RemoveInvulnerables, ShouldRotateAllCollators, Slot,
60
    },
61
};
62
pub use {dp_collator_assignment::AssignedCollators, pallet::*};
63

            
64
mod assignment;
65
#[cfg(feature = "runtime-benchmarks")]
66
mod benchmarking;
67
pub mod weights;
68
pub use weights::WeightInfo;
69

            
70
#[cfg(test)]
71
mod mock;
72

            
73
#[cfg(test)]
74
mod tests;
75

            
76
#[derive(Encode, Decode, Debug, TypeInfo)]
77
pub struct CoreAllocationConfiguration {
78
    pub core_count: u32,
79
    pub max_parachain_percentage: Perbill,
80
}
81

            
82
6240
#[frame_support::pallet]
83
pub mod pallet {
84
    use super::*;
85

            
86
2816
    #[pallet::pallet]
87
    pub struct Pallet<T>(_);
88

            
89
    /// Configure the pallet by specifying the parameters and types on which it depends.
90
    #[pallet::config]
91
    pub trait Config: frame_system::Config {
92
        /// The overarching event type.
93
        type RuntimeEvent: From<Event<Self>> + IsType<<Self as frame_system::Config>::RuntimeEvent>;
94
        type SessionIndex: parity_scale_codec::FullCodec
95
            + TypeInfo
96
            + Copy
97
            + AtLeast32BitUnsigned
98
            + Debug;
99
        // `SESSION_DELAY` is used to delay any changes to Paras registration or configurations.
100
        // Wait until the session index is 2 larger then the current index to apply any changes,
101
        // which guarantees that at least one full session has passed before any changes are applied.
102
        type HostConfiguration: GetHostConfiguration<Self::SessionIndex>;
103
        type ContainerChains: GetSessionContainerChains<Self::SessionIndex>;
104
        type SelfParaId: Get<ParaId>;
105
        type ShouldRotateAllCollators: ShouldRotateAllCollators<Self::SessionIndex>;
106
        type Randomness: CollatorAssignmentRandomness<BlockNumberFor<Self>>;
107
        type RemoveInvulnerables: RemoveInvulnerables<Self::AccountId>;
108
        type ParaIdAssignmentHooks: ParaIdAssignmentHooks<BalanceOf<Self>, Self::AccountId>;
109
        type Currency: Currency<Self::AccountId>;
110
        type CollatorAssignmentTip: CollatorAssignmentTip<BalanceOf<Self>>;
111
        type ForceEmptyOrchestrator: Get<bool>;
112
        type CoreAllocationConfiguration: Get<Option<CoreAllocationConfiguration>>;
113
        /// The weight information of this pallet.
114
        type WeightInfo: WeightInfo;
115
    }
116

            
117
624
    #[pallet::event]
118
4878
    #[pallet::generate_deposit(pub(super) fn deposit_event)]
119
    pub enum Event<T: Config> {
120
        NewPendingAssignment {
121
            random_seed: [u8; 32],
122
            full_rotation: bool,
123
            target_session: T::SessionIndex,
124
            full_rotation_mode: FullRotationModes,
125
        },
126
    }
127

            
128
308312
    #[pallet::storage]
129
    #[pallet::unbounded]
130
    pub(crate) type CollatorContainerChain<T: Config> =
131
        StorageValue<_, AssignedCollators<T::AccountId>, ValueQuery>;
132

            
133
    /// Pending configuration changes.
134
    ///
135
    /// This is a list of configuration changes, each with a session index at which it should
136
    /// be applied.
137
    ///
138
    /// The list is sorted ascending by session index. Also, this list can only contain at most
139
    /// 2 items: for the next session and for the `scheduled_session`.
140
24826
    #[pallet::storage]
141
    #[pallet::unbounded]
142
    pub(crate) type PendingCollatorContainerChain<T: Config> =
143
        StorageValue<_, Option<AssignedCollators<T::AccountId>>, ValueQuery>;
144

            
145
    /// Randomness from previous block. Used to shuffle collators on session change.
146
    /// Should only be set on the last block of each session and should be killed on the on_initialize of the next block.
147
    /// The default value of [0; 32] disables randomness in the pallet.
148
12930
    #[pallet::storage]
149
    pub(crate) type Randomness<T: Config> = StorageValue<_, [u8; 32], ValueQuery>;
150

            
151
    /// Ratio of assigned collators to max collators.
152
10380
    #[pallet::storage]
153
    pub type CollatorFullnessRatio<T: Config> = StorageValue<_, Perbill, OptionQuery>;
154

            
155
624
    #[pallet::call]
156
    impl<T: Config> Pallet<T> {}
157

            
158
    /// A struct that holds the assignment that is active after the session change and optionally
159
    /// the assignment that becomes active after the next session change.
160
    pub struct SessionChangeOutcome<T: Config> {
161
        /// New active assignment.
162
        pub active_assignment: AssignedCollators<T::AccountId>,
163
        /// Next session active assignment.
164
        pub next_assignment: AssignedCollators<T::AccountId>,
165
        /// Total number of registered parachains before filtering them out, used as a weight hint
166
        pub num_total_registered_paras: u32,
167
    }
168

            
169
    impl<T: Config> Pallet<T> {
170
4886
        pub(crate) fn enough_collators_for_all_chains(
171
4886
            bulk_paras: &[ChainNumCollators],
172
4886
            pool_paras: &[ChainNumCollators],
173
4886
            target_session_index: T::SessionIndex,
174
4886
            number_of_collators: u32,
175
4886
            collators_per_container: u32,
176
4886
            collators_per_parathread: u32,
177
4886
        ) -> bool {
178
4886
            number_of_collators
179
4886
                >= T::HostConfiguration::min_collators_for_orchestrator(target_session_index)
180
4886
                    .saturating_add(collators_per_container.saturating_mul(bulk_paras.len() as u32))
181
4886
                    .saturating_add(
182
4886
                        collators_per_parathread.saturating_mul(pool_paras.len() as u32),
183
4886
                    )
184
4886
        }
185

            
186
        /// Takes the bulk paras (parachains) and pool paras (parathreads)
187
        /// and checks if we if a) Do we have enough collators? b) Do we have enough cores?
188
        /// If either of the answer is yes. We  separately sort bulk_paras and pool_paras and
189
        /// then append the two vectors.
190
1200
        pub(crate) fn order_paras_with_core_config(
191
1200
            mut bulk_paras: Vec<ChainNumCollators>,
192
1200
            mut pool_paras: Vec<ChainNumCollators>,
193
1200
            core_allocation_configuration: &CoreAllocationConfiguration,
194
1200
            target_session_index: T::SessionIndex,
195
1200
            number_of_collators: u32,
196
1200
            collators_per_container: u32,
197
1200
            collators_per_parathread: u32,
198
1200
        ) -> (Vec<ChainNumCollators>, bool) {
199
1200
            let core_count = core_allocation_configuration.core_count;
200
1200
            let max_number_of_bulk_paras = core_allocation_configuration
201
1200
                .max_parachain_percentage
202
1200
                .mul(core_count);
203
1200

            
204
1200
            let enough_cores_for_bulk_paras = bulk_paras.len() <= max_number_of_bulk_paras as usize;
205
1200

            
206
1200
            let enough_collators = Self::enough_collators_for_all_chains(
207
1200
                &bulk_paras,
208
1200
                &pool_paras,
209
1200
                target_session_index,
210
1200
                number_of_collators,
211
1200
                collators_per_container,
212
1200
                collators_per_parathread,
213
1200
            );
214

            
215
            // We should charge tip if parachain demand exceeds the `max_number_of_bulk_paras` OR
216
            // if `num_collators` is not enough to satisfy  collation need of all paras.
217
1200
            let should_charge_tip = !enough_cores_for_bulk_paras || !enough_collators;
218

            
219
            // Currently, we are sorting both bulk and pool paras by tip, even when for example
220
            // only number of bulk paras are restricted due to core availability since we deduct tip from
221
            // all paras.
222
            // We need to sort both separately as we have fixed space for parachains at the moment
223
            // which means even when we have some parathread cores empty we cannot schedule parachain there.
224
1200
            if should_charge_tip {
225
916
                bulk_paras.sort_by(|a, b| {
226
447
                    T::CollatorAssignmentTip::get_para_tip(b.para_id)
227
447
                        .cmp(&T::CollatorAssignmentTip::get_para_tip(a.para_id))
228
916
                });
229
859

            
230
888
                pool_paras.sort_by(|a, b| {
231
106
                    T::CollatorAssignmentTip::get_para_tip(b.para_id)
232
106
                        .cmp(&T::CollatorAssignmentTip::get_para_tip(a.para_id))
233
888
                });
234
859
            }
235

            
236
1200
            bulk_paras.truncate(max_number_of_bulk_paras as usize);
237
1200
            // We are not truncating pool paras, since their workload is not continuous one core
238
1200
            // can be shared by many paras during the session.
239
1200

            
240
1200
            let chains: Vec<_> = bulk_paras.into_iter().chain(pool_paras).collect();
241
1200

            
242
1200
            (chains, should_charge_tip)
243
1200
        }
244

            
245
3686
        pub(crate) fn order_paras(
246
3686
            bulk_paras: Vec<ChainNumCollators>,
247
3686
            pool_paras: Vec<ChainNumCollators>,
248
3686
            target_session_index: T::SessionIndex,
249
3686
            number_of_collators: u32,
250
3686
            collators_per_container: u32,
251
3686
            collators_per_parathread: u32,
252
3686
        ) -> (Vec<ChainNumCollators>, bool) {
253
3686
            // Are there enough collators to satisfy the minimum demand?
254
3686
            let enough_collators_for_all_chain = Self::enough_collators_for_all_chains(
255
3686
                &bulk_paras,
256
3686
                &pool_paras,
257
3686
                target_session_index,
258
3686
                number_of_collators,
259
3686
                collators_per_container,
260
3686
                collators_per_parathread,
261
3686
            );
262
3686

            
263
3686
            let mut chains: Vec<_> = bulk_paras.into_iter().chain(pool_paras).collect();
264
3686

            
265
3686
            // Prioritize paras by tip on congestion
266
3686
            // As of now this doesn't distinguish between bulk paras and pool paras
267
3686
            if !enough_collators_for_all_chain {
268
2193
                chains.sort_by(|a, b| {
269
2193
                    T::CollatorAssignmentTip::get_para_tip(b.para_id)
270
2193
                        .cmp(&T::CollatorAssignmentTip::get_para_tip(a.para_id))
271
2193
                });
272
2978
            }
273

            
274
3686
            (chains, !enough_collators_for_all_chain)
275
3686
        }
276

            
277
        /// Assign new collators
278
        /// collators should be queued collators
279
4878
        pub fn assign_collators(
280
4878
            current_session_index: &T::SessionIndex,
281
4878
            random_seed: [u8; 32],
282
4878
            collators: Vec<T::AccountId>,
283
4878
        ) -> SessionChangeOutcome<T> {
284
4878
            let maybe_core_allocation_configuration = T::CoreAllocationConfiguration::get();
285
4878
            // We work with one session delay to calculate assignments
286
4878
            let session_delay = T::SessionIndex::one();
287
4878
            let target_session_index = current_session_index.saturating_add(session_delay);
288
4878

            
289
4878
            let collators_per_container =
290
4878
                T::HostConfiguration::collators_per_container(target_session_index);
291
4878
            let collators_per_parathread =
292
4878
                T::HostConfiguration::collators_per_parathread(target_session_index);
293
4878

            
294
4878
            // We get the containerChains that we will have at the target session
295
4878
            let container_chains =
296
4878
                T::ContainerChains::session_container_chains(target_session_index);
297
4878
            let num_total_registered_paras = container_chains
298
4878
                .parachains
299
4878
                .len()
300
4878
                .saturating_add(container_chains.parathreads.len())
301
4878
                as u32;
302
4878
            let mut container_chain_ids = container_chains.parachains;
303
4878
            let mut parathreads: Vec<_> = container_chains
304
4878
                .parathreads
305
4878
                .into_iter()
306
4878
                .map(|(para_id, _)| para_id)
307
4878
                .collect();
308
4878

            
309
4878
            // We read current assigned collators
310
4878
            let old_assigned = Self::read_assigned_collators();
311
4878
            let old_assigned_para_ids: BTreeSet<ParaId> =
312
4878
                old_assigned.container_chains.keys().cloned().collect();
313
4878

            
314
4878
            // Remove the containerChains that do not have enough credits for block production
315
4878
            T::ParaIdAssignmentHooks::pre_assignment(
316
4878
                &mut container_chain_ids,
317
4878
                &old_assigned_para_ids,
318
4878
            );
319
4878
            // TODO: parathreads should be treated a bit differently, they don't need to have the same amount of credits
320
4878
            // as parathreads because they will not be producing blocks on every slot.
321
4878
            T::ParaIdAssignmentHooks::pre_assignment(&mut parathreads, &old_assigned_para_ids);
322
4878

            
323
4878
            let mut shuffle_collators = None;
324
4878
            // If the random_seed is all zeros, we don't shuffle the list of collators nor the list
325
4878
            // of container chains.
326
4878
            // This should only happen in tests_without_core_config, and in the genesis block.
327
4878
            if random_seed != [0; 32] {
328
121
                let mut rng: ChaCha20Rng = SeedableRng::from_seed(random_seed);
329
121
                container_chain_ids.shuffle(&mut rng);
330
121
                parathreads.shuffle(&mut rng);
331
217
                shuffle_collators = Some(move |collators: &mut Vec<T::AccountId>| {
332
217
                    collators.shuffle(&mut rng);
333
217
                })
334
4757
            }
335

            
336
4878
            let orchestrator_chain: ChainNumCollators = if T::ForceEmptyOrchestrator::get() {
337
1192
                ChainNumCollators {
338
1192
                    para_id: T::SelfParaId::get(),
339
1192
                    min_collators: 0u32,
340
1192
                    max_collators: 0u32,
341
1192
                    parathread: false,
342
1192
                }
343
            } else {
344
3686
                ChainNumCollators {
345
3686
                    para_id: T::SelfParaId::get(),
346
3686
                    min_collators: T::HostConfiguration::min_collators_for_orchestrator(
347
3686
                        target_session_index,
348
3686
                    ),
349
3686
                    max_collators: T::HostConfiguration::max_collators_for_orchestrator(
350
3686
                        target_session_index,
351
3686
                    ),
352
3686
                    parathread: false,
353
3686
                }
354
            };
355

            
356
            // Initialize list of chains as `[container1, container2, parathread1, parathread2]`.
357
            // The order means priority: the first chain in the list will be the first one to get assigned collators.
358
            // Chains will not be assigned less than `min_collators`, except the orchestrator chain.
359
            // First all chains will be assigned `min_collators`, and then the first one will be assigned up to `max`,
360
            // then the second one, and so on.
361
4878
            let mut bulk_paras = vec![];
362
4878
            let mut pool_paras = vec![];
363

            
364
10571
            for para_id in &container_chain_ids {
365
5693
                bulk_paras.push(ChainNumCollators {
366
5693
                    para_id: *para_id,
367
5693
                    min_collators: collators_per_container,
368
5693
                    max_collators: collators_per_container,
369
5693
                    parathread: false,
370
5693
                });
371
5693
            }
372
5409
            for para_id in &parathreads {
373
531
                pool_paras.push(ChainNumCollators {
374
531
                    para_id: *para_id,
375
531
                    min_collators: collators_per_parathread,
376
531
                    max_collators: collators_per_parathread,
377
531
                    parathread: true,
378
531
                });
379
531
            }
380

            
381
4878
            let (chains, need_to_charge_tip) =
382
4878
                if let Some(core_allocation_configuration) = maybe_core_allocation_configuration {
383
1192
                    Self::order_paras_with_core_config(
384
1192
                        bulk_paras,
385
1192
                        pool_paras,
386
1192
                        &core_allocation_configuration,
387
1192
                        target_session_index,
388
1192
                        collators.len() as u32,
389
1192
                        collators_per_container,
390
1192
                        collators_per_parathread,
391
1192
                    )
392
                } else {
393
3686
                    Self::order_paras(
394
3686
                        bulk_paras,
395
3686
                        pool_paras,
396
3686
                        target_session_index,
397
3686
                        collators.len() as u32,
398
3686
                        collators_per_container,
399
3686
                        collators_per_parathread,
400
3686
                    )
401
                };
402

            
403
            // We assign new collators
404
            // we use the config scheduled at the target_session_index
405
4878
            let full_rotation =
406
4878
                T::ShouldRotateAllCollators::should_rotate_all_collators(target_session_index);
407
4878
            if full_rotation {
408
406
                log::info!(
409
372
                    "Collator assignment: rotating collators. Session {:?}, Seed: {:?}",
410
372
                    current_session_index.encode(),
411
                    random_seed
412
                );
413
            } else {
414
4472
                log::info!(
415
1884
                    "Collator assignment: keep old assigned. Session {:?}, Seed: {:?}",
416
1884
                    current_session_index.encode(),
417
                    random_seed
418
                );
419
            }
420

            
421
4878
            let full_rotation_mode = if full_rotation {
422
406
                T::HostConfiguration::full_rotation_mode(target_session_index)
423
            } else {
424
                // On sessions where there is no rotation, we try to keep all collators assigned to the same chains
425
4472
                FullRotationModes::keep_all()
426
            };
427

            
428
4878
            Self::deposit_event(Event::NewPendingAssignment {
429
4878
                random_seed,
430
4878
                full_rotation,
431
4878
                target_session: target_session_index,
432
4878
                full_rotation_mode: full_rotation_mode.clone(),
433
4878
            });
434
4878

            
435
4878
            let new_assigned = Assignment::<T>::assign_collators_always_keep_old(
436
4878
                collators,
437
4878
                orchestrator_chain,
438
4878
                chains,
439
4878
                old_assigned.clone(),
440
4878
                shuffle_collators,
441
4878
                full_rotation_mode,
442
4878
            );
443

            
444
4878
            let mut new_assigned = match new_assigned {
445
4874
                Ok(x) => x,
446
4
                Err(e) => {
447
4
                    log::error!(
448
4
                        "Error in collator assignment, will keep previous assignment. {:?}",
449
                        e
450
                    );
451

            
452
4
                    old_assigned.clone()
453
                }
454
            };
455

            
456
4878
            let mut assigned_containers = new_assigned.container_chains.clone();
457
7126
            assigned_containers.retain(|_, v| !v.is_empty());
458

            
459
            // On congestion, prioritized chains need to pay the minimum tip of the prioritized chains
460
4878
            let maybe_tip: Option<BalanceOf<T>> = if !need_to_charge_tip {
461
1935
                None
462
            } else {
463
2943
                assigned_containers
464
2943
                    .into_keys()
465
2943
                    .filter_map(T::CollatorAssignmentTip::get_para_tip)
466
2943
                    .min()
467
            };
468

            
469
            // TODO: this probably is asking for a refactor
470
            // only apply the onCollatorAssignedHook if sufficient collators
471
4878
            T::ParaIdAssignmentHooks::post_assignment(
472
4878
                &old_assigned_para_ids,
473
4878
                &mut new_assigned.container_chains,
474
4878
                &maybe_tip,
475
4878
            );
476
4878

            
477
4878
            Self::store_collator_fullness(
478
4878
                &new_assigned,
479
4878
                T::HostConfiguration::max_collators(target_session_index),
480
4878
            );
481
4878

            
482
4878
            let mut pending = PendingCollatorContainerChain::<T>::get();
483
4878

            
484
4878
            let old_assigned_changed = old_assigned != new_assigned;
485
4878
            let mut pending_changed = false;
486
            // Update CollatorContainerChain using last entry of pending, if needed
487
4878
            if let Some(current) = pending.take() {
488
1165
                pending_changed = true;
489
1165
                CollatorContainerChain::<T>::put(current);
490
3739
            }
491
4878
            if old_assigned_changed {
492
1568
                pending = Some(new_assigned.clone());
493
1568
                pending_changed = true;
494
4024
            }
495
            // Update PendingCollatorContainerChain, if it changed
496
4878
            if pending_changed {
497
2307
                PendingCollatorContainerChain::<T>::put(pending);
498
3743
            }
499

            
500
            // Only applies to session index 0
501
4878
            if current_session_index == &T::SessionIndex::zero() {
502
1340
                CollatorContainerChain::<T>::put(new_assigned.clone());
503
1340
                return SessionChangeOutcome {
504
1340
                    active_assignment: new_assigned.clone(),
505
1340
                    next_assignment: new_assigned,
506
1340
                    num_total_registered_paras,
507
1340
                };
508
3538
            }
509
3538

            
510
3538
            SessionChangeOutcome {
511
3538
                active_assignment: old_assigned,
512
3538
                next_assignment: new_assigned,
513
3538
                num_total_registered_paras,
514
3538
            }
515
4878
        }
516

            
517
        /// Count number of collators assigned to any chain, divide that by `max_collators` and store
518
        /// in pallet storage.
519
4878
        fn store_collator_fullness(
520
4878
            new_assigned: &AssignedCollators<T::AccountId>,
521
4878
            max_collators: u32,
522
4878
        ) {
523
4878
            // Count number of assigned collators
524
4878
            let mut num_collators = 0;
525
4878
            num_collators.saturating_accrue(new_assigned.orchestrator_chain.len());
526
7126
            for collators in new_assigned.container_chains.values() {
527
6192
                num_collators.saturating_accrue(collators.len());
528
6192
            }
529

            
530
4878
            let mut num_collators = num_collators as u32;
531
4878
            if num_collators > max_collators {
532
148
                // Shouldn't happen but just in case
533
148
                num_collators = max_collators;
534
4878
            }
535

            
536
4878
            let ratio = Perbill::from_rational(num_collators, max_collators);
537
4878

            
538
4878
            CollatorFullnessRatio::<T>::put(ratio);
539
4878
        }
540

            
541
        // Returns the assigned collators as read from storage.
542
        // If there is any item in PendingCollatorContainerChain, returns that element.
543
        // Otherwise, reads and returns the current CollatorContainerChain
544
4878
        fn read_assigned_collators() -> AssignedCollators<T::AccountId> {
545
4878
            let mut pending_collator_list = PendingCollatorContainerChain::<T>::get();
546

            
547
4878
            if let Some(assigned_collators) = pending_collator_list.take() {
548
1165
                assigned_collators
549
            } else {
550
                // Read current
551
3713
                CollatorContainerChain::<T>::get()
552
            }
553
4878
        }
554

            
555
4878
        pub fn initializer_on_new_session(
556
4878
            session_index: &T::SessionIndex,
557
4878
            collators: Vec<T::AccountId>,
558
4878
        ) -> SessionChangeOutcome<T> {
559
4878
            let random_seed = T::Randomness::take_randomness();
560
4878
            let num_collators = collators.len();
561
4878
            let assigned_collators = Self::assign_collators(session_index, random_seed, collators);
562
4878
            let num_total_registered_paras = assigned_collators.num_total_registered_paras;
563
4878

            
564
4878
            frame_system::Pallet::<T>::register_extra_weight_unchecked(
565
4878
                T::WeightInfo::new_session(num_collators as u32, num_total_registered_paras),
566
4878
                DispatchClass::Mandatory,
567
4878
            );
568
4878

            
569
4878
            assigned_collators
570
4878
        }
571

            
572
117826
        pub fn collator_container_chain() -> AssignedCollators<T::AccountId> {
573
117826
            CollatorContainerChain::<T>::get()
574
117826
        }
575

            
576
36
        pub fn pending_collator_container_chain() -> Option<AssignedCollators<T::AccountId>> {
577
36
            PendingCollatorContainerChain::<T>::get()
578
36
        }
579
    }
580

            
581
    impl<T: Config> GetContainerChainAuthor<T::AccountId> for Pallet<T> {
582
109107
        fn author_for_slot(slot: Slot, para_id: ParaId) -> Option<T::AccountId> {
583
109107
            let assigned_collators = Pallet::<T>::collator_container_chain();
584
109107
            let collators = if para_id == T::SelfParaId::get() {
585
85809
                Some(&assigned_collators.orchestrator_chain)
586
            } else {
587
23298
                assigned_collators.container_chains.get(&para_id)
588
            }?;
589

            
590
109107
            if collators.is_empty() {
591
                // Avoid division by zero below
592
                return None;
593
109107
            }
594
109107
            let author_index = u64::from(slot) % collators.len() as u64;
595
109107
            collators.get(author_index as usize).cloned()
596
109107
        }
597

            
598
        #[cfg(feature = "runtime-benchmarks")]
599
        fn set_authors_for_para_id(para_id: ParaId, authors: Vec<T::AccountId>) {
600
            let mut assigned_collators = Pallet::<T>::collator_container_chain();
601
            assigned_collators.container_chains.insert(para_id, authors);
602
            CollatorContainerChain::<T>::put(assigned_collators);
603
        }
604
    }
605

            
606
66623
    #[pallet::hooks]
607
    impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
608
35358
        fn on_initialize(n: BlockNumberFor<T>) -> Weight {
609
35358
            let mut weight = Weight::zero();
610
35358

            
611
35358
            // Account reads and writes for on_finalize
612
35358
            weight.saturating_accrue(T::Randomness::prepare_randomness_weight(n));
613
35358

            
614
35358
            weight
615
35358
        }
616

            
617
34816
        fn on_finalize(n: BlockNumberFor<T>) {
618
34816
            // If the next block is a session change, read randomness and store in pallet storage
619
34816
            T::Randomness::prepare_randomness(n);
620
34816
        }
621
    }
622

            
623
    impl<T: Config> GetContainerChainsWithCollators<T::AccountId> for Pallet<T> {
624
29682
        fn container_chains_with_collators(
625
29682
            for_session: ForSession,
626
29682
        ) -> Vec<(ParaId, Vec<T::AccountId>)> {
627
29682
            // If next session has None then current session data will stay.
628
29682
            let chains = (for_session == ForSession::Next)
629
29682
                .then(PendingCollatorContainerChain::<T>::get)
630
29682
                .flatten()
631
29682
                .unwrap_or_else(CollatorContainerChain::<T>::get);
632
29682

            
633
29682
            chains.container_chains.into_iter().collect()
634
29682
        }
635

            
636
2987
        fn get_all_collators_assigned_to_chains(for_session: ForSession) -> BTreeSet<T::AccountId> {
637
2987
            let mut all_chains: Vec<T::AccountId> =
638
2987
                Self::container_chains_with_collators(for_session)
639
2987
                    .iter()
640
5054
                    .flat_map(|(_para_id, collators)| collators.iter())
641
2987
                    .cloned()
642
2987
                    .collect();
643
2987
            all_chains.extend(
644
2987
                Self::collator_container_chain()
645
2987
                    .orchestrator_chain
646
2987
                    .iter()
647
2987
                    .cloned(),
648
2987
            );
649
2987
            all_chains.into_iter().collect()
650
2987
        }
651

            
652
        #[cfg(feature = "runtime-benchmarks")]
653
        fn set_container_chains_with_collators(
654
            for_session: ForSession,
655
            container_chains: &[(ParaId, Vec<T::AccountId>)],
656
        ) {
657
            match for_session {
658
                ForSession::Current => {
659
                    let mut collators = CollatorContainerChain::<T>::get();
660
                    collators.container_chains = container_chains.iter().cloned().collect();
661
                    CollatorContainerChain::<T>::put(collators);
662
                }
663
                ForSession::Next => {
664
                    let mut collators =
665
                        PendingCollatorContainerChain::<T>::get().unwrap_or_default();
666
                    collators.container_chains = container_chains.iter().cloned().collect();
667
                    PendingCollatorContainerChain::<T>::put(Some(collators));
668
                }
669
            }
670
        }
671
    }
672
}
673

            
674
/// Balance used by this pallet
675
pub type BalanceOf<T> =
676
    <<T as Config>::Currency as Currency<<T as frame_system::Config>::AccountId>>::Balance;
677

            
678
pub struct RotateCollatorsEveryNSessions<Period>(PhantomData<Period>);
679

            
680
impl<Period> ShouldRotateAllCollators<u32> for RotateCollatorsEveryNSessions<Period>
681
where
682
    Period: Get<u32>,
683
{
684
4711
    fn should_rotate_all_collators(session_index: u32) -> bool {
685
4711
        let period = Period::get();
686
4711

            
687
4711
        if period == 0 {
688
            // A period of 0 disables rotation
689
1716
            false
690
        } else {
691
2995
            session_index % Period::get() == 0
692
        }
693
4711
    }
694
}
695

            
696
/// Only works on parachains because in relaychains it is not possible to know for sure if the next
697
/// block will be in the same session as the current one, as it depends on slots and validators can
698
/// skip slots.
699
pub trait GetRandomnessForNextBlock<BlockNumber> {
700
    fn should_end_session(block_number: BlockNumber) -> bool;
701
    fn get_randomness() -> [u8; 32];
702
}
703

            
704
impl<BlockNumber> GetRandomnessForNextBlock<BlockNumber> for () {
705
    fn should_end_session(_block_number: BlockNumber) -> bool {
706
        false
707
    }
708

            
709
    fn get_randomness() -> [u8; 32] {
710
        [0; 32]
711
    }
712
}
713

            
714
pub trait CollatorAssignmentRandomness<BlockNumber> {
715
    /// Called in on_initialize, returns weight needed by prepare_randomness call.
716
    fn prepare_randomness_weight(n: BlockNumber) -> Weight;
717
    /// Called in on_finalize.
718
    /// Prepares randomness for the next block if the next block is a new session start.
719
    fn prepare_randomness(n: BlockNumber);
720
    /// Called once at the start of each session in on_initialize of pallet_initializer
721
    fn take_randomness() -> [u8; 32];
722
}
723

            
724
impl<BlockNumber> CollatorAssignmentRandomness<BlockNumber> for () {
725
1057
    fn prepare_randomness_weight(_n: BlockNumber) -> Weight {
726
1057
        Weight::zero()
727
1057
    }
728
987
    fn prepare_randomness(_n: BlockNumber) {}
729
167
    fn take_randomness() -> [u8; 32] {
730
167
        [0; 32]
731
167
    }
732
}
733

            
734
/// Parachain randomness impl.
735
///
736
/// Reads relay chain randomness in the last block of the session and stores it in pallet storage.
737
/// When new session starts, takes that value from storage removing it.
738
/// Relay randomness cannot be accessed in `on_initialize`, so `prepare_randomness` is executed in
739
/// `on_finalize`, with `prepare_randomness_weight` reserving the weight needed.
740
pub struct ParachainRandomness<T, Runtime>(PhantomData<(T, Runtime)>);
741

            
742
impl<BlockNumber, T, Runtime> CollatorAssignmentRandomness<BlockNumber>
743
    for ParachainRandomness<T, Runtime>
744
where
745
    BlockNumber: Saturating + One,
746
    T: GetRandomnessForNextBlock<BlockNumber>,
747
    Runtime: frame_system::Config + crate::Config,
748
{
749
26960
    fn prepare_randomness_weight(n: BlockNumber) -> Weight {
750
26960
        let mut weight = Weight::zero();
751
26960

            
752
26960
        if T::should_end_session(n.saturating_add(One::one())) {
753
2637
            weight.saturating_accrue(Runtime::DbWeight::get().reads_writes(1, 1));
754
24323
        }
755

            
756
26960
        weight
757
26960
    }
758

            
759
26834
    fn prepare_randomness(n: BlockNumber) {
760
26834
        if T::should_end_session(n.saturating_add(One::one())) {
761
2634
            let random_seed = T::get_randomness();
762
2634
            Randomness::<Runtime>::put(random_seed);
763
24200
        }
764
26834
    }
765

            
766
3519
    fn take_randomness() -> [u8; 32] {
767
3519
        Randomness::<Runtime>::take()
768
3519
    }
769
}
770

            
771
/// Solochain randomness.
772
///
773
/// Uses current block randomness. This randomness exists in `on_initialize` so we don't need to
774
/// `prepare_randomness` in the previous block.
775
pub struct SolochainRandomness<T>(PhantomData<T>);
776

            
777
impl<BlockNumber, T> CollatorAssignmentRandomness<BlockNumber> for SolochainRandomness<T>
778
where
779
    T: Get<[u8; 32]>,
780
{
781
7341
    fn prepare_randomness_weight(_n: BlockNumber) -> Weight {
782
7341
        Weight::zero()
783
7341
    }
784

            
785
6995
    fn prepare_randomness(_n: BlockNumber) {}
786

            
787
1192
    fn take_randomness() -> [u8; 32] {
788
1192
        #[cfg(feature = "runtime-benchmarks")]
789
1192
        if let Some(x) =
790
1192
            frame_support::storage::unhashed::take(b"__bench_collator_assignment_randomness")
791
1192
        {
792
1192
            return x;
793
1192
        }
794
1192

            
795
1192
        T::get()
796
1192
    }
797
}