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
extern crate alloc;
44

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            
205
1212
            let enough_cores_for_bulk_paras = bulk_paras.len() <= max_number_of_bulk_paras as usize;
206
1212

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

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

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

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

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

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

            
243
1212
            (chains, should_charge_tip)
244
1212
        }
245

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

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

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

            
275
6108
            (chains, !enough_collators_for_all_chain)
276
6108
        }
277

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

            
290
7312
            let collators_per_container =
291
7312
                T::HostConfiguration::collators_per_container(target_session_index);
292
7312
            let collators_per_parathread =
293
7312
                T::HostConfiguration::collators_per_parathread(target_session_index);
294
7312

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            
457
7312
            let mut assigned_containers = new_assigned.container_chains.clone();
458
9683
            assigned_containers.retain(|_, v| !v.is_empty());
459

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

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

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

            
483
7312
            let mut pending = PendingCollatorContainerChain::<T>::get();
484
7312

            
485
7312
            let old_assigned_changed = old_assigned != new_assigned;
486
7312
            let mut pending_changed = false;
487
            // Update CollatorContainerChain using last entry of pending, if needed
488
7312
            if let Some(current) = pending.take() {
489
1979
                pending_changed = true;
490
1979
                CollatorContainerChain::<T>::put(current);
491
5359
            }
492
7312
            if old_assigned_changed {
493
2836
                pending = Some(new_assigned.clone());
494
2836
                pending_changed = true;
495
5293
            }
496
            // Update PendingCollatorContainerChain, if it changed
497
7312
            if pending_changed {
498
4060
                PendingCollatorContainerChain::<T>::put(pending);
499
5489
            }
500

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

            
511
5079
            SessionChangeOutcome {
512
5079
                active_assignment: old_assigned,
513
5079
                next_assignment: new_assigned,
514
5079
                num_total_registered_paras,
515
5079
            }
516
7312
        }
517

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

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

            
537
7312
            let ratio = Perbill::from_rational(num_collators, max_collators);
538
7312

            
539
7312
            CollatorFullnessRatio::<T>::put(ratio);
540
7312
        }
541

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

            
548
7312
            if let Some(assigned_collators) = pending_collator_list.take() {
549
1979
                assigned_collators
550
            } else {
551
                // Read current
552
5333
                CollatorContainerChain::<T>::get()
553
            }
554
7312
        }
555

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

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

            
570
7312
            assigned_collators
571
7312
        }
572

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

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

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

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

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

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

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

            
615
38009
            weight
616
38009
        }
617

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

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

            
634
33514
            chains.container_chains.into_iter().collect()
635
33514
        }
636

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

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

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

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

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

            
688
7145
        if period == 0 {
689
            // A period of 0 disables rotation
690
1840
            false
691
        } else {
692
5305
            session_index % Period::get() == 0
693
        }
694
7145
    }
695
}
696

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

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

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

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

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

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

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

            
753
29481
        if T::should_end_session(n.saturating_add(One::one())) {
754
2879
            weight.saturating_accrue(Runtime::DbWeight::get().reads_writes(1, 1));
755
26602
        }
756

            
757
29481
        weight
758
29481
    }
759

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

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

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

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

            
786
7117
    fn prepare_randomness(_n: BlockNumber) {}
787

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

            
796
1204
        T::get()
797
1204
    }
798
}